huangfu/InteractionImitation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Use ray tune in gail

Browse Source
This commit is contained in:
Johannes Fischer
2021-10-08 18:19:08 +02:00
parent 6d2ab54b6e
commit c59da3611c
2 changed files with 286 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

115
scratch/etienne/intersimple/gail_flat_ray.py Normal file
View File

@@ -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)

171
scratch/etienne/intersimple/gail_image_random_ray.py Normal file
View File

@@ -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)
# %%