From 01752fac125d020f80f312e9fe8af28b6bd300f9 Mon Sep 17 00:00:00 2001 From: Arec Date: Mon, 11 Oct 2021 07:14:03 -0700 Subject: [PATCH] adding model to append actions to encoded image state for discriminator, files to test different environment experiments, and a script to render an environment from a saved model --- scratch/etienne/intersimple/data/generate.sh | 2 + .../etienne/intersimple/gail/discriminator.py | 41 +++++++++++ .../gail_image_multiagent_nocollision.py | 70 +++++++++++++++++++ .../gail_image_singleagent_nocollision.py | 70 +++++++++++++++++++ .../intersimple/render_env_from_model.py | 21 ++++++ 5 files changed, 204 insertions(+) mode change 100644 => 100755 scratch/etienne/intersimple/data/generate.sh create mode 100644 scratch/etienne/intersimple/gail_image_multiagent_nocollision.py create mode 100644 scratch/etienne/intersimple/gail_image_singleagent_nocollision.py create mode 100644 scratch/etienne/intersimple/render_env_from_model.py diff --git a/scratch/etienne/intersimple/data/generate.sh b/scratch/etienne/intersimple/data/generate.sh old mode 100644 new mode 100755 index 4b329fc..d7387b2 --- a/scratch/etienne/intersimple/data/generate.sh +++ b/scratch/etienne/intersimple/data/generate.sh @@ -3,3 +3,5 @@ #python -m expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.001.pkl' python -m expert --env=NRasterized --min_timesteps=200 --env_args='{agent:51,width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl' # python -m expert --env=NRasterizedRandomAgent --min_timesteps=10000 --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_NRasterizedRandomAgentw36h36mppx2.pkl' +#python -m expert --env=NRasterized --min_timesteps=200 --env_args='{agent:51,width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl' +#python -m expert --env=NRasterized --min_timesteps=3000 --video --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedRandomAgentw36h36mppx2.pkl' diff --git a/scratch/etienne/intersimple/gail/discriminator.py b/scratch/etienne/intersimple/gail/discriminator.py index f3c16e0..02cb83a 100644 --- a/scratch/etienne/intersimple/gail/discriminator.py +++ b/scratch/etienne/intersimple/gail/discriminator.py @@ -39,6 +39,47 @@ class CnnDiscriminator(torch.nn.Module): assert sa.ndim == 4 return self.cnn(sa).squeeze(1) +class CnnDiscriminatorFlatAction(torch.nn.Module): + """ConvNet similar to stable_baselines3.common.policies.ActorCriticCnnPolicy.""" + + def __init__(self, env): + super().__init__() + + obs_channels, _, _ = env.observation_space.shape + (action_size,) = env.action_space.shape + in_channels = obs_channels + + self.cnn = torch.nn.Sequential( + torch.nn.Conv2d(in_channels, 32, kernel_size=(8, 8), stride=(4, 4)), # in_channels -> 32 + torch.nn.ReLU(), + torch.nn.Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2)), # 32 -> 64 + torch.nn.ReLU(), + torch.nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1)), # 64 -> 64 + torch.nn.ReLU(), + torch.nn.Flatten(start_dim=1, end_dim=-1), + torch.nn.LazyLinear(128), # 28224 -> 128 + ) + self.decoder = torch.nn.Sequential( + torch.nn.LazyLinear(64), #128 + 2 -> 64 + torch.nn.ReLU(), + torch.nn.LazyLinear(64), #64 -> 64 + torch.nn.ReLU(), + torch.nn.LazyLinear(1) #64 -> 1 + ) + + @staticmethod + def _concatenate(state, action): + b, s= state.shape + b, a = action.shape + sa = torch.cat((state, action), -1) + return sa + + def forward(self, state, action): + s = self.cnn(state) + sa = self._concatenate(s, action) + assert sa.ndim == 2 + return self.decoder(sa).squeeze(1) + class MlpDiscriminator(torch.nn.Module): """MLP similar to stable_baselines3.common.policies.ActorCriticPolicy.""" diff --git a/scratch/etienne/intersimple/gail_image_multiagent_nocollision.py b/scratch/etienne/intersimple/gail_image_multiagent_nocollision.py new file mode 100644 index 0000000..4619b5a --- /dev/null +++ b/scratch/etienne/intersimple/gail_image_multiagent_nocollision.py @@ -0,0 +1,70 @@ +# %% +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 NRasterized + +from gail.discriminator import CnnDiscriminatorFlatAction + +model_name = 'gail_image_multiagent_nocollision' + +# %% +# Load pickled test demonstrations. +with open("data/NormalizedIntersimpleExpertMu.001_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) + +venv = make_vec_env(NRasterized, n_envs=2, env_kwargs={'stop_on_collision':False, 'width': 36, 'height': 36, 'm_per_px': 2}) + +tempdir = tempfile.TemporaryDirectory(prefix="quickstart") +tempdir_path = pathlib.Path(tempdir.name) +print(f"All Tensorboards and logging are being written inside {tempdir_path}/.") + +# 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. +logger.configure(tempdir_path / "GAIL/") +gail_trainer = adversarial.GAIL( + venv, + expert_data=transitions, + expert_batch_size=32, + #n_disc_updates_per_round=2048, + discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)}, + gen_algo=sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=1024), + allow_variable_horizon=True, +) +gail_trainer.train(total_timesteps=100000) +gail_trainer.gen_algo.save(model_name) + +#del gail_trainer + +# %% +model = sb3.PPO.load(model_name) + +env = NRasterized(stop_on_collision=False, width=36, height=36, m_per_px=2) + +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) diff --git a/scratch/etienne/intersimple/gail_image_singleagent_nocollision.py b/scratch/etienne/intersimple/gail_image_singleagent_nocollision.py new file mode 100644 index 0000000..1c8cec4 --- /dev/null +++ b/scratch/etienne/intersimple/gail_image_singleagent_nocollision.py @@ -0,0 +1,70 @@ +# %% +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 NRasterized + +from gail.discriminator import CnnDiscriminator + +model_name = 'gail_image_singleagent_nocollision' + +# %% +# Load pickled test demonstrations. +with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.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(NRasterized, n_envs=2, env_kwargs={'agent':51, 'stop_on_collision':False, 'width': 36, 'height': 36, 'm_per_px': 2}) + +tempdir = tempfile.TemporaryDirectory(prefix="quickstart") +tempdir_path = pathlib.Path(tempdir.name) +print(f"All Tensorboards and logging are being written inside {tempdir_path}/.") + +# 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. +logger.configure(tempdir_path / "GAIL/") +gail_trainer = adversarial.GAIL( + venv, + expert_data=transitions, + expert_batch_size=32, + #n_disc_updates_per_round=2048, + discrim_kwargs={'discrim_net': CnnDiscriminator(venv)}, + gen_algo=sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=1024), + allow_variable_horizon=True, +) +gail_trainer.train(total_timesteps=100000) +gail_trainer.gen_algo.save(model_name) + +#del gail_trainer + +# %% +model = sb3.PPO.load(model_name) + +env = NRasterized(agent=51, width=36, height=36, m_per_px=2, stop_on_collision=False) + +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) diff --git a/scratch/etienne/intersimple/render_env_from_model.py b/scratch/etienne/intersimple/render_env_from_model.py new file mode 100644 index 0000000..f44a670 --- /dev/null +++ b/scratch/etienne/intersimple/render_env_from_model.py @@ -0,0 +1,21 @@ + +import stable_baselines3 as sb3 +from intersim.envs.intersimple import NRasterized + + +model_name = 'gail_image_singleagent_nocollision' +model = sb3.PPO.load(model_name) + +env = NRasterized(stop_on_collision=False, width=36, height=36, m_per_px=2, agent=51) + +obs = env.reset() +i=0 +while True and i < 600: + i+=1 + 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