Use ray tune in gail

scratch/etienne/intersimple/gail_image_random_ray.py

@@ -0,0 +1,171 @@
+# %%
+import pathlib
+import pickle
+import tempfile
+import os
+import random
+import numpy as np
+import torch
+
+# set up ray tune
+import ray
+from ray import tune
+from ray.tune import Analysis, ExperimentAnalysis
+from ray.tune.schedulers import ASHAScheduler
+from ray.tune.suggest.hyperopt import HyperOptSearch
+from ray.tune.suggest import ConcurrencyLimiter
+
+import stable_baselines3 as sb3
+from stable_baselines3.common.env_util import make_vec_env
+
+
+from imitation.algorithms import adversarial, bc
+from imitation.data import rollout
+from imitation.util import logger
+
+from intersim.envs.intersimple import NRasterizedRandomAgent, IntersimpleReward, speed_reward, NRasterized, NRasterizedRandomAgentVerbose
+import functools
+from stable_baselines3.common.evaluation import evaluate_policy
+from gym.wrappers import TimeLimit
+
+from gail.discriminator import CnnDiscriminator
+
+model_name = 'gail_image_random_ray'
+env_kwargs={'width': 36, 'height': 36, 'm_per_px': 2}
+
+# %%
+
+import argparse
+parser = argparse.ArgumentParser()
+parser.add_argument("--outdir", help="result directory", default='ray')
+parser.add_argument("--test", help="test run", default=False, action="store_true")
+args = parser.parse_args()
+outdir = args.outdir
+
+# %%
+# Load pickled test demonstrations.
+with open("data/NormalizedIntersimpleExpertMu.001N10000_NRasterizedRandomAgentw36h36mppx2.pkl", "rb") as f:
+    # This is a list of `imitation.data.types.Trajectory`, where
+    # every instance contains observations and actions for a single expert
+    # demonstration.
+    trajectories = pickle.load(f)
+
+# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
+# This is a more general dataclass containing unordered
+# (observation, actions, next_observation) transitions.
+transitions = rollout.flatten_trajectories(trajectories)
+# Store transitions in shared ray memory
+ray_transitions = ray.put(transitions)
+
+# %%
+venv = make_vec_env(NRasterizedRandomAgent, n_envs=2, env_kwargs=env_kwargs)
+
+tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
+tempdir_path = pathlib.Path(tempdir.name)
+print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
+logger.configure(tempdir_path / "GAIL/")
+
+def get_ray_config(test=False):
+    if test:
+        return {
+            'expert_batch_size': 2,
+            'ppo_n_steps': 2,
+            'ppo_batch_size': 2,
+            'ppo_n_epochs': 1,
+            'total_timesteps': 10,
+        }
+    else:
+        return {
+            'expert_batch_size': tune.choice([2**x for x in range(6,10)]),
+            'ppo_n_steps': tune.choice([2048, 3072, 4096]),
+            'ppo_batch_size': tune.choice([2**x for x in range(9,13)]),
+            'ppo_n_epochs': tune.choice([6,10]),
+            'total_timesteps': 400_000,
+        }
+
+
+def ray_train(config, checkpoint_dir=None):
+    # Train GAIL on expert data.
+    # GAIL, and AIRL also accept as `expert_data` any Pytorch-style DataLoader that
+    # iterates over dictionaries containing observations, actions, and next_observations.
+
+    discriminator = CnnDiscriminator(venv)
+    if checkpoint_dir:
+        discriminator.load_state_dict(torch.load(os.path.join(checkpoint_dir, 'disc_checkpoint')))
+        generator = sb3.PPO.load(os.path.join(checkpoint_dir, 'gen_checkpoint'))
+    else:
+        generator = sb3.PPO(
+            "CnnPolicy", venv, verbose=0, 
+            n_steps=config["ppo_n_steps"], 
+            batch_size=config["ppo_batch_size"], 
+            n_epochs=config["ppo_n_epochs"]
+        )
+    gail_trainer = adversarial.GAIL(
+        venv,
+        expert_data=ray.get(ray_transitions),
+        expert_batch_size=config["expert_batch_size"],
+        #n_disc_updates_per_round=2048,
+        discrim_kwargs={'discrim_net': discriminator},
+        gen_algo=generator,
+        allow_variable_horizon=True,
+    )
+    def callback(round):
+        # eval_env = NRasterized(agent=51, reward=functools.partial(speed_reward, collision_penalty=0.), **env_kwargs)
+        eval_env = TimeLimit(NRasterizedRandomAgent(reward=functools.partial(speed_reward, collision_penalty=0.), **env_kwargs), max_episode_steps=1000)
+        episode_rewards, episode_lengths = evaluate_policy(generator, eval_env, return_episode_rewards=True)
+        tune.report(
+            reward=np.mean(episode_rewards),
+            length=np.mean(episode_lengths),
+            training_iteration=round,
+        )
+        with tune.checkpoint_dir(step=round) as checkpoint_dir:
+            gail_trainer.gen_algo.save(os.path.join(checkpoint_dir, 'gen_checkpoint'))
+            torch.save(discriminator.state_dict(), os.path.join(checkpoint_dir, 'disc_checkpoint'))
+
+    gail_trainer.train(total_timesteps=config['total_timesteps'], callback=callback)
+
+
+ray_config = get_ray_config(args.test)
+search = HyperOptSearch(ray_config, metric='length', mode="max",)
+search = ConcurrencyLimiter(search, max_concurrent=10)
+custom_scheduler = ASHAScheduler(time_attr='training_iteration', metric='length', mode="max", grace_period=15)
+
+analysis = tune.run(
+    ray_train, 
+    # config=ray_config,
+    search_alg=search,
+    scheduler=custom_scheduler,
+    local_dir=outdir,
+    resources_per_trial={"cpu":10, "gpu": 0.2},
+    num_samples=1 if args.test else 100,
+)
+
+del analysis
+
+# %%
+# outdir = "ray/ray_train_2021-09-20_13-33-50/ray_train_f06785b0_33_expert_batch_size=128,ppo_batch_size=1024,ppo_n_epochs=6,ppo_n_steps=2048,total_timesteps=400000_2021-09-20_15-52-05"
+
+# %%
+analysis = Analysis(outdir, default_metric="length", default_mode="max")
+filepath = analysis.get_best_logdir()
+print("Best ray experiment:", filepath)
+config = analysis.get_best_config()
+print("Best config:", config)
+
+# %%
+
+model = sb3.PPO.load(os.path.join(analysis.get_last_checkpoint(), 'gen_checkpoint'))
+
+# env = NRasterized(agent=51, **env_kwargs)
+env = TimeLimit(NRasterizedRandomAgent(reward=functools.partial(speed_reward, collision_penalty=0.), **env_kwargs), max_episode_steps=1000)
+obs = env.reset()
+while True:
+    action, _states = model.predict(obs)
+    obs, rewards, done, info = env.step(action)
+    env.render(mode='post')
+    if done:
+        break
+
+env.env.close(filestr='render/'+model_name)
+# %%
+