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

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

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."""
 

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)

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)

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)