From c59da3611c7c66ede5049b8af954e49c7ba308cf Mon Sep 17 00:00:00 2001 From: Johannes Fischer Date: Fri, 8 Oct 2021 18:19:08 +0200 Subject: [PATCH] Use ray tune in gail --- scratch/etienne/intersimple/gail_flat_ray.py | 115 ++++++++++++ .../intersimple/gail_image_random_ray.py | 171 ++++++++++++++++++ 2 files changed, 286 insertions(+) create mode 100644 scratch/etienne/intersimple/gail_flat_ray.py create mode 100644 scratch/etienne/intersimple/gail_image_random_ray.py diff --git a/scratch/etienne/intersimple/gail_flat_ray.py b/scratch/etienne/intersimple/gail_flat_ray.py new file mode 100644 index 0000000..2d37c41 --- /dev/null +++ b/scratch/etienne/intersimple/gail_flat_ray.py @@ -0,0 +1,115 @@ +# %% +import pathlib +import pickle +import tempfile + +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 IntersimpleReward, speed_reward + +from gail.discriminator import MlpDiscriminator +import numpy as np +import functools +from stable_baselines3.common.evaluation import evaluate_policy +from ray import tune +import os +import torch + +model_name = 'gail_flat' + +# %% +# Load pickled test demonstrations. +#with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51.pkl", "rb") as f: +with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.001.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) + +venv = make_vec_env(IntersimpleReward, n_envs=2, env_kwargs={'agent': 51}) + +tempdir = tempfile.TemporaryDirectory(prefix="quickstart") +tempdir_path = pathlib.Path(tempdir.name) +print(f"All Tensorboards and logging are being written inside {tempdir_path}/.") + +def training_function(config, checkpoint_dir=None): + logger.configure(tempdir_path / "GAIL/") + + discriminator = MlpDiscriminator() + 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("MlpPolicy", venv, verbose=1, n_steps=config['n_steps']) + + gail_trainer = adversarial.GAIL( + venv, + expert_data=transitions, + expert_batch_size=config['expert_batch_size'], + n_disc_updates_per_round=config['n_disc_updates_per_round'], + discrim_kwargs={'discrim_net': MlpDiscriminator()}, + gen_algo=generator, + allow_variable_horizon=True, + ) + + def callback(epoch): + print("callback") + eval_env = IntersimpleReward(agent=51, reward=functools.partial(speed_reward, collision_penalty=0.)) + #sync_envs_normalization(self.training_env, self.eval_env) + 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=epoch, + ) + + with tune.checkpoint_dir(step=epoch) 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=40000, callback=callback) + +analysis = tune.run( + training_function, + config = { + 'expert_batch_size': tune.randint(1, 22), #220, + 'n_disc_updates_per_round': tune.randint(2, 100), #16, + 'n_steps': tune.randint(1, 10000), #4096, + }, + resources_per_trial={ + 'cpu': 1, + # 'gpu': 1, + }, + local_dir='ray', + num_samples=10, +) + +print('Best config', analysis.get_best_config(metric='progress', mode='max')) + +#del gail_trainer + +# %% +model = sb3.PPO.load(model_name) + +env = IntersimpleReward(agent=51) + +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.close(filestr='render/'+model_name) \ No newline at end of file diff --git a/scratch/etienne/intersimple/gail_image_random_ray.py b/scratch/etienne/intersimple/gail_image_random_ray.py new file mode 100644 index 0000000..3c00c5c --- /dev/null +++ b/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) +# %% +