Smaller conv nets

- support predict()
- move discount to OptionsEnv
- fix RenderOptions

.gitignore

@@ -1,3 +1,8 @@
+*.pkl
+*.pt
+*.zip
+**/ray/*
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
@@ -145,3 +150,4 @@ output/
 
 # Dependencies
 InteractionSimulator/
+imitation/

generate_demos.sh

@@ -0,0 +1,10 @@
+#DEFAULT PARAMETERS:
+# locs:list=None, (default to all locations)
+# tracks:list=None, (default to all tracks)
+# env_class:str='NRasterizedIncrementingAgent',
+# env_args:dict={width:36,height:36,m_per_px:2},
+# expert_class:str='NRasterizedRouteIncrementingAgent',  
+# expert_args:dict={mu:0.001}):
+
+# python -m src.data.expert --locs='[DR_USA_Roundabout_FT]' --tracks='[0]'
+python -m src.data.expert --locs='[DR_USA_Roundabout_FT]' --tracks='[0]'

requirements.txt

@@ -4,6 +4,7 @@ sklearn
 pytest
 json5
 tqdm
-tensorboard
 ray[tune]
-hyperopt
+hyperopt
+psutil
+fire

scratch/arec/intersimple/commands.txt

@@ -0,0 +1,22 @@
+python -m render_options --model_name='gail_options_image_mid_wcollision' --env='NRasterizedRoute' --options=True --width=36 --height=36 --m_per_px=2 --agent=50 --stop_on_collision=False
+
+import torch, os
+from src.data import load_experts
+folder = 'expert_data/DR_USA_Roundabout_FT/track0000'
+single_agent = os.path.join(folder, 'expert.pkl')
+multi_agent = os.path.join(folder,'joint_expert_states.pt')
+multi_agent_actions = os.path.join(folder,'joint_expert_actions.pt')
+demonstrations = load_experts([single_agent], flatten=False)
+demonstrations[0].__dict__.keys()
+len(demonstrations[0].obs)
+single_agent_lengths = [len(demonstration.obs) for demonstration in demonstrations]
+states = torch.load(multi_agent)
+actions = torch.load(multi_agent_actions)
+multi_agent_lengths = [sum(~torch.isnan(states[:,i,0])).item() for i in range(states.shape[1])]
+
+single_agent_actions = [demonstration.acts for demonstration in demonstrations]
+multi_agent_actions = [actions[~torch.isnan(actions[:,i,0])] for i in range(actions.shape[1])]
+
+import pickle
+with open(single_agent, "rb") as f:
+    new_trajectories = pickle.load(f)

scratch/arec/intersimple/data/expert.py

@@ -0,0 +1,145 @@
+from intersim.envs.intersimple import Intersimple
+from stable_baselines3.common.policies import BasePolicy
+import gym
+import intersim.envs.intersimple
+import imitation.data.rollout as rollout
+from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
+from imitation.data.wrappers import RolloutInfoWrapper
+
+class IntersimExpert(BasePolicy):
+    
+    def __init__(self, intersim_env, mu=0, *args, **kwargs):
+        super().__init__(
+            observation_space=gym.spaces.Space(),
+            action_space=gym.spaces.Space(),
+            *args, **kwargs
+        )
+        self._intersim = intersim_env
+        self._mu = mu
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def _predict(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def _action(self):
+        target_t = min(self._intersim._ind + 1, len(self._intersim._svt.simstate) - 1)
+        target_state = self._intersim._svt.simstate[target_t]
+        return self._intersim.target_state(target_state, mu=self._mu)
+
+    def predict(self, *args, **kwargs):
+        return self._action(), None
+
+class IntersimpleExpert(BasePolicy):
+
+    def __init__(self, intersimple_env, mu=0, *args, **kwargs):
+        super().__init__(
+            observation_space=intersimple_env.observation_space,
+            action_space=intersimple_env.action_space,
+            *args, **kwargs
+        )
+        self._intersimple = intersimple_env
+        self._intersim_expert = IntersimExpert(intersimple_env._env, mu=mu)
+    
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def _predict(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def _action(self):
+        return self._intersim_expert._action()[self._intersimple._agent]
+
+    def predict(self, *args, **kwargs):
+        return self._action(), None
+
+class NormalizedIntersimpleExpert(IntersimpleExpert):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def predict(self, *args, **kwargs):
+        action, _ = super().predict(*args, **kwargs)
+        return self._intersimple._normalize(action), None
+
+class DummyVecEnvPolicy(BasePolicy):
+
+    def __init__(self, experts):
+        self._experts = [e() for e in experts]
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def _predict(self, *args, **kwargs):
+        raise NotImplementedError()
+    
+    def predict(self, *args, **kwargs):
+        predictions = [e.predict() for e in self._experts]
+        actions = [p[0] for p in predictions]
+        states = [p[1] for p in predictions]
+        return actions, states
+    
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def _predict(self, *args, **kwargs):
+        raise NotImplementedError()
+
+def save_video(env, expert):
+    env.reset()
+    env.render()
+    done = False
+    while not done:
+        actions, _ = expert.predict()
+        _, _, done, _ = env.step(actions)
+        env.render()
+    env.close()
+
+def demonstrations(expert='NormalizedIntersimpleExpert', env='NRasterizedIncrementingAgent', path=None, min_timesteps=None, min_episodes=None, video=False, env_args={}, policy_args={}):
+    """Rollout and save expert demos.
+    
+    Usage:
+        python -m intersimple.expert <flags>
+    Args: 
+        expert (class): class of expert
+        env (class): class of env intersim.envs.intersimple
+        path (str): path to store output
+        min_timesteps (int): min number of timesteps for call to rollout.rollout_and_save
+        min_episodes (int): min number of episodes for call to rollout.rollout_and_save
+        video (bool): whether to save a video of the expert until a single environment instantiation stops
+        env_args (dict): dictionary of kwargs when instantiating environment class
+        policy_args (dict): dictionary of kwargs when instantiating Expert policy
+    """
+    
+    Env = intersim.envs.intersimple.__dict__[env]
+    Expert = globals()[expert]
+
+    env = Env(**env_args)
+    info_env = RolloutInfoWrapper(env) # getting rollout info (dictionary) from environment
+    venv = DummyVecEnv([lambda: info_env]) # making a DummyVecEnv with a list of a function that when called returns the rollout info
+
+    policy = Expert(env, **policy_args) # instantiate an expert policy from specified class with instantiated environment and policy kwargs
+    venv_policy = DummyVecEnvPolicy([lambda: policy]) # make a DummyVecEnvPolicy with a list of a function that when called returns the Expert policy
+
+    if min_timesteps is None and min_episodes is None:
+        min_episodes = env.nv # one episode per vehicle being controlled in environment (hopefully an incrementing agent environment)
+
+    if video:
+        save_video(env, policy)
+    
+    path = path or (policy.__class__.__name__ + '_' + env.__class__.__name__ + '.pkl')
+    suntil = rollout.make_sample_until(
+            min_timesteps=min_timesteps,
+            min_episodes=min_episodes,
+        )
+    rollout.rollout_and_save(
+        path=path,
+        policy=venv_policy,
+        venv=venv,
+        sample_until=suntil
+    )
+
+if __name__ == '__main__':
+    import fire
+    fire.Fire(demonstrations)

scratch/arec/intersimple/data/generate.sh

@@ -0,0 +1,9 @@
+#python -m expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51.pkl'
+#python -m expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --policy_args='{mu:0.005}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.005.pkl'
+#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'
+#python -m expert --env=NRasterizedIncrementingAgent --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedIncrementingAgentw36h36mppx2.pkl'
+python -m process_all_experts --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}'

scratch/arec/intersimple/data/load_experts.py

@@ -0,0 +1,23 @@
+import pickle
+import imitation.data.rollout as rollout
+from tqdm import tqdm
+
+def load_experts(expert_files=[]):
+    """
+    Load expert trajectories from files and combine their transitions into a single RB
+
+    Args:
+        expert_files (list): list of expert file strings
+    Returns:
+        transitions (list): list of combined expert episode transitions 
+    """
+    transitions = []
+    for file in tqdm(expert_files):
+        with open(file, "rb") as f:
+            trajectories = pickle.load(f)
+        transitions = transitions + rollout.flatten_trajectories(trajectories)
+    return transitions
+
+if __name__=='__main__':
+    import fire
+    fire.Fire(load_experts)

scratch/arec/intersimple/data/process_all_experts.py

@@ -0,0 +1,48 @@
+import tqdm
+import expert
+import copy
+import os
+import intersim
+from tqdm import tqdm
+
+def process_all_experts(filename='expert.pkl',env_args={}, policy_args={}):
+    """
+    Process all experts in the Interaction Dataset
+    For now, using NormalizedIntersimpleExpert with NRasterizedIncrementingAgent environment
+
+    Args:
+        filename (str): name for track file
+        env_args (dict): default environment kwargs
+        policy_args (dict): default policy kwargs
+    """
+    I, J = len(intersim.LOCATIONS), intersim.MAX_TRACKS
+    pbar = tqdm(total=I*J)
+    for loc in range(I):
+        for track in range(J):
+            
+            it_env_args = copy.deepcopy(env_args)
+            it_env_args.update({
+                'loc':loc,
+                'track':track,
+            })
+            out_folder = os.path.join(intersim.LOCATIONS[loc], 'track%04i'%(track))
+            if not os.path.isdir(out_folder):
+                os.makedirs(out_folder) 
+            it_path = os.path.join(out_folder,filename)
+
+            expert.demonstrations(
+                expert='NormalizedIntersimpleExpert', 
+                env='NRasterizedIncrementingAgent', 
+                path=it_path, 
+                env_args=it_env_args, 
+                policy_args=policy_args,
+                )
+            pbar.update(1)
+    pbar.close()
+    
+
+if __name__=='__main__':
+    import fire
+    fire.Fire(process_all_experts)
+
+

scratch/arec/intersimple/gail/discriminator.py

@@ -0,0 +1,101 @@
+import torch
+
+# imitation.rewards.discrim_nets.DiscrimNetGAIL is composed of self.discriminator (nn.Module),
+# which gets called with inputs (state, action) when needed.
+
+class CnnDiscriminator(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 + action_size
+
+        self.cnn = torch.nn.Sequential(
+            torch.nn.Conv2d(in_channels, 32, kernel_size=(8, 8), stride=(4, 4)), # 5+1 -> 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(512), # 28224 -> 512
+            torch.nn.ReLU(),
+            torch.nn.LazyLinear(1), # 512 -> 1
+        )
+    
+    @staticmethod
+    def _concatenate(state, action):
+        b, _, h, w = state.shape
+        _, a = action.shape
+        act = action.unsqueeze(-1).unsqueeze(-1).expand((b, a, h, w))
+        sa = torch.cat((state, act), -3)
+        return sa
+
+    def forward(self, state, action):
+        sa = self._concatenate(state, action)
+        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.float())
+        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."""
+
+    def __init__(self, env=None):
+        super().__init__()
+        self.flatten = torch.nn.Flatten(start_dim=1, end_dim=-1)
+        self.mlp = torch.nn.Sequential(
+            torch.nn.LazyLinear(64), # 42 -> 64
+            torch.nn.Tanh(),
+            torch.nn.LazyLinear(64), # 64 -> 64
+            torch.nn.Tanh(),
+            torch.nn.LazyLinear(1),  # 64 -> 1
+        )
+    
+    def forward(self, state, action):
+        flat = self.flatten(state)
+        sa = torch.cat((action, flat), -1)
+        assert sa.ndim == 2
+        return self.mlp(sa).squeeze(1)

scratch/arec/intersimple/gail/test_discriminator.py

@@ -0,0 +1,45 @@
+from intersim.envs.intersimple import NRasterized
+from discriminator import CnnDiscriminator
+import torch
+
+def test_image_concatenation():
+    env = NRasterized()
+    disc = CnnDiscriminator(env)
+    s = torch.tensor(env.reset()).unsqueeze(0)
+    a = torch.tensor([[0.5]])
+    sa = disc._concatenate(s, a)
+
+    assert s.shape == (1, 5, 200, 200)
+    assert a.shape == (1, 1)
+    assert sa.shape == (1, 6, 200, 200)
+    assert torch.allclose(sa[:, :5], 1.0 * s)
+    assert (sa[:, 5] == a.unsqueeze(-1)).all()
+
+def test_image_concatenation3():
+    env = NRasterized()
+    disc = CnnDiscriminator(env)
+
+    s1 = env.reset()
+    a1 = 0.15
+    s2, _, _, _ = env.step(0.9)
+    a2 = 0.25
+    s3, _, _, _ = env.step(-0.9)
+    a3 = 0.35
+    
+    s = torch.stack([
+        torch.tensor(s1),
+        torch.tensor(s2),
+        torch.tensor(s3)
+    ], axis=0)
+    a = torch.tensor([
+        [a1],
+        [a2],
+        [a3],
+    ])
+    sa = disc._concatenate(s, a)
+
+    assert s.shape == (3, 5, 200, 200)
+    assert a.shape == (3, 1)
+    assert sa.shape == (3, 6, 200, 200)
+    assert torch.allclose(sa[:, :5], 1.0 * s)
+    assert (sa[:, 5] == a.unsqueeze(-1)).all()

scratch/arec/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/arec/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/arec/intersimple/gail_options_image.py

@@ -0,0 +1,172 @@
+# %%
+import sys
+sys.path.append('../../../')
+from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
+from src.policies import OptionsCnnPolicy
+from src.util import render_env
+from src.data import load_experts
+from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
+from src.gail.train import train_discriminator, train_generator
+
+from imitation.algorithms import adversarial
+from imitation.util import logger
+import imitation.data.rollout as rollout
+
+import stable_baselines3
+from stable_baselines3.common.env_util import make_vec_env
+
+import torch
+import torch.utils.data
+import numpy as np
+import itertools
+import gym
+import pickle
+import tempfile
+import pathlib
+from tqdm import tqdm
+
+from intersim.envs.intersimple import NRasterized, NRasterizedRoute, NRasterizedRandomAgent, NRasterizedIncrementingAgent, NRasterizedRouteRandomAgent
+
+ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10]] # option 0 is safe fallback
+
+def flatten_transitions(transitions):
+    return {
+        'obs': np.stack(list(t['obs'] for t in transitions), axis=0),
+        'next_obs': np.stack(list(t['next_obs'] for t in transitions), axis=0),
+        'acts': np.stack(list(t['acts'] for t in transitions), axis=0),
+        'dones': np.stack(list(t['dones'] for t in transitions), axis=0),
+    }
+
+def train(expert_data, env_class=NRasterizedRouteRandomAgent, env_settings={}, 
+        epochs=10, discrim_batch_size=32, generator_steps=2048, discount=0.99):
+    """
+    Args: 
+        expert_data: list of transitions
+        env_class: environment class
+        env_settings: environment settings
+        epochs: number of epochs to train for
+        discrim_batch_size: discriminator batch size
+        generator_steps: number of steps taken in generator
+        discount: discount factor
+    Returns:
+        generator (stable_baselines3.PPO): options policy
+    """
+    env = env_class(**env_settings)
+    env.discount = discount
+
+    tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
+    tempdir_path = pathlib.Path(tempdir.name)
+    logger.configure(tempdir_path / "GAIL/")
+    print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
+
+    venv = make_vec_env(env_class, n_envs=1, env_kwargs=env_settings)
+    discriminator = adversarial.GAIL(
+        expert_data=expert_data,
+        expert_batch_size=discrim_batch_size,
+        discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
+        #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
+        venv=venv, # unused
+        gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
+    )
+
+    generator = stable_baselines3.PPO(
+        OptionsCnnPolicy,
+        OptionsEnv(env, options=ALL_OPTIONS),
+        verbose=1,
+        n_steps=generator_steps,
+    )
+
+    # PPO.train requires logger as set up in
+    # PPO._setup_learn (called by PPO.learn)
+    generator._logger = stable_baselines3.common.utils.configure_logger(
+        generator.verbose,
+        generator.tensorboard_log,
+    )
+
+    for _ in tqdm(range(epochs)):
+        train_discriminator(LLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=discrim_batch_size)
+        train_generator(HLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=generator_steps)
+    
+    return generator
+
+# %%
+if __name__ == '__main__':
+    # %%
+    model_name = 'gail_options_image_mid_wcollision'
+    env_class = NRasterizedRouteRandomAgent
+    env_settings = {'width': 36, 'height': 36, 'm_per_px': 2, 'stop_on_collision': False}
+    
+    #env_class = NRasterized
+    #env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2}
+    files = ['../../../expert_data/DR_USA_Roundabout_FT/track%04i/expert.pkl'%(i) for i in range(5)]
+    transitions=load_experts(files)
+
+    generator = train(
+        transitions,
+        env_class=env_class,
+        env_settings=env_settings,
+        epochs=2, 
+        discrim_batch_size=256, 
+        generator_steps=10,#256, 
+        discount=0.99
+        )
+
+    generator.save(model_name)
+
+    # Render
+    render_settings = {'width': 36, 'height': 36, 'm_per_px': 2, 'agent':51, 'stop_on_collision': False}
+    render_env(model_name=model_name, env='NRasterizedRoute', options=True, options_list=ALL_OPTIONS,
+        **render_settings)
+
+
+# %% Tests
+
+def test_ll_expert_data():
+    with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
+        expert_trajectories = pickle.load(f)
+    expert_transitions = rollout.flatten_trajectories(expert_trajectories)
+
+    env = LLOptions(NRasterized(agent=51, width=36, height=36, m_per_px=2))
+
+    gen_transitions = list(itertools.islice(env.sample_ll(
+        policy=stable_baselines3.PPO(
+            OptionsCnnPolicy,
+            OptionsEnv(env),
+            verbose=1,
+        )
+    ), 10))
+    gen_transitions = flatten_transitions(gen_transitions)
+
+    assert expert_transitions[:10].obs.shape == gen_transitions['obs'].shape
+    assert expert_transitions[:10].next_obs.shape == gen_transitions['next_obs'].shape
+    assert expert_transitions[:10].acts.shape == gen_transitions['acts'].shape
+    assert expert_transitions[:10].dones.shape == gen_transitions['dones'].shape
+
+def test_ll_states():
+    env = NRasterized()
+    policy = stable_baselines3.PPO(
+        OptionsCnnPolicy,
+        OptionsEnv(env),
+        verbose=1,
+    )
+    llenv = LLOptions(env)
+    transitions = list(itertools.islice(llenv.sample_ll(policy=policy), 100))
+
+    env2 = NRasterized()
+    s2 = env2.reset()
+    for i, t in enumerate(transitions):
+        assert i == 0 or np.array_equal(t['obs'], transitions[i-1]['next_obs'])
+        assert np.array_equal(t['obs'], s2)
+        assert t['acts'].shape == (1,)
+
+        nexts2, _, done2, _ = env2.step(t['acts'])
+        assert np.array_equal(t['next_obs'], nexts2)
+        assert np.array_equal(t['dones'], done2)
+
+        if done2:
+            break
+
+        s2 = nexts2
+
+def test_hl_transitions():
+    pass

scratch/arec/intersimple/gail_options_scratch.py

@@ -0,0 +1,559 @@
+# %%
+from gail.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
+from imitation.algorithms import adversarial
+import stable_baselines3
+import torch.utils.data
+import numpy as np
+from intersim.envs.intersimple import NRasterized, NRasterizedRandomAgent
+import itertools
+from torch.distributions import Categorical
+import gym
+import torch
+import pickle
+import imitation.data.rollout as rollout
+import tempfile
+import pathlib
+from imitation.util import logger
+from stable_baselines3.common.env_util import make_vec_env
+from tqdm import tqdm
+
+import logging 
+logging.basicConfig(level=logging.DEBUG)
+
+ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback
+
+class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy):
+    """
+    Class for high-level options policy (generator)
+    """
+    def __init__(self, observation_space, *args, **kwargs):
+        super().__init__(observation_space['obs'], *args, **kwargs)
+
+    def _prior_distribution(self, s):
+        """
+        Return prior distribution over high-level options (before masking)
+        Args:
+            s (torch.tensor): observation
+        Returns:
+            values (torch.tensor): values from critic
+            dist (torch.distributions): prior distribution over actions
+        """
+        latent_pi, latent_vf, latent_sde = self._get_latent(s) 
+        distribution = self._get_action_dist_from_latent(latent_pi, latent_sde) 
+        values = self.value_net(latent_vf)
+        return values, distribution.distribution
+
+    def predict(self, obs):
+        """
+        Will mask invalid states before making action selections
+        Args:
+            obs: dict with keys:
+                obs (torch.tensor): (B,o)  true observations
+                mask (torch.tensor): (B,m) mask over valid actions
+        Returns:
+            ch (torch.tensor): (B,a) sampled actions
+            values (torch.tensor): (B,) predicted value at observation
+            log_probs (torch.tensor): (B,) log probabilities of selected actions
+        """
+        s, m = obs['obs'], obs['mask']
+        values, prior = self._prior_distribution(s)
+        posterior = Categorical(prior.probs * m)
+        ch = posterior.sample()
+        return ch, values, posterior.log_prob(ch)
+
+    def evaluate_actions(self, obs, ch):
+        """
+        Evaluate particular actions
+        Args:
+            obs: dict with keys:
+                obs (torch.tensor): (B,o) true observations
+                mask (torch.tensor): (B,m) masks over valid actions
+            ch (torch.tensor): (B,a) selected actions
+        Returns:
+            values (torch.tensor): (B,) predicted value at observation
+            log_probs (torch.tensor): (B,) log probabilities of selected actions
+            ent (torch.tensor): (B,) entropy of each distribution over actions
+        """
+        s, m = obs['obs'], obs['mask']
+        values, prior = self._prior_distribution(s)
+        posterior = Categorical(prior.probs * m)
+        return values, posterior.log_prob(ch), posterior.entropy() # additional values used by PPO.train
+
+class OptionsEnv(gym.Wrapper):
+    """
+    Wrap an intersimple environment with an options generator
+    """
+    def __init__(self, env, *args, **kwargs):
+        """
+        Initialize wrapped environment and set high-level action and observation spaces
+        """
+        super().__init__(env, *args, **kwargs)
+        num_hl_options = len(ALL_OPTIONS)
+        self.action_space = gym.spaces.Discrete(num_hl_options)
+        self.observation_space = gym.spaces.Dict({
+            'obs': env.observation_space,
+            'mask': gym.spaces.Box(low=0, high=1, shape=(num_hl_options,)),
+        })
+
+    def _after_choice(self):
+        pass
+    
+    def _after_step(self):
+        pass
+
+    def _transitions(self):
+        raise NotImplementedError('Use `LLOptions` or `HLOptions` for sampling.')
+    
+    def sample(self, generator):
+        """
+        yield transitions using a generator
+        Args:
+            generator (sb3.PPO)
+        Yields:
+
+        """
+        self.done = True
+        while True:
+            self.episode_start = False
+
+            if self.done:
+                # reset environment
+                self.s = self.env.reset()
+                self.m = available_actions(self.env)
+                self.done = False
+                self.episode_start = True
+
+            # set the action, the value of the start state, and the logprob of the action
+            # according to the current environment state and mask
+            self.ch, self.value, self.log_prob = generator.policy.predict({
+                'obs': torch.tensor(self.s).unsqueeze(0).to(generator.policy.device),
+                'mask': torch.tensor(self.m).unsqueeze(0).to(generator.policy.device),
+            })
+
+            # store a float list of actions to take given the option selected in the environment
+            self.plan = list(map(float, generate_plan(self.env, self.ch)))
+
+            # run whatever _after_choice might dictate in a child class
+            self._after_choice()
+
+            # some checks
+            assert not self.done
+            assert self.plan
+            assert feasible(self.env, self.plan, self.ch)
+
+            # execute the option so long as the episode isn't complete and the plan is still feasible
+            while not self.done and self.plan and feasible(self.env, self.plan, self.ch):
+                
+                # pop first action
+                self.a, self.plan = self.plan[0], self.plan[1:]
+                
+                # normalize action ??
+                self.a = self.env._normalize(self.a)
+                
+                # step through environment
+                self.nexts, _, self.done, _ = self.env.step(self.a)
+                self.nextm = available_actions(self.env)
+
+                # run whatever _after_step might dictate in child class
+                self._after_step()
+
+                # update state and mask to current
+                self.s = self.nexts
+                self.m = self.nextm
+
+            # transitions yielded from self._transitions() functions specied in child classes
+            yield from self._transitions()
+
+            ### NOTE: only yields after a full option has been executed / exited
+
+class LLOptions(OptionsEnv):
+    """Sample low-level (state, action) tuples for discriminator training."""
+    
+    def __init__(self, *args, **kwargs):
+        """
+        LLOption uses the true LL observations
+        """
+        super().__init__(*args, **kwargs)
+        # overwrite observation space to just output obs directly
+        self.observation_space = self.observation_space['obs'] 
+
+    def _after_choice(self):
+        """
+        After each option choice, initialize/reset the transition buffer
+        """
+        self._transition_buffer = []
+
+    def _after_step(self):
+        """
+        After each ll action, append s, s', a, done to transition buffer
+        """
+        self._transition_buffer.append({
+            'obs': self.s,
+            'next_obs': self.nexts,
+            'acts': np.array((self.a,)),
+            'dones': np.array(self.done),
+        })
+
+    def _transitions(self):
+        """
+        Yield from the transition buffer
+        """
+        yield from self._transition_buffer
+    
+    def sample_ll(self, policy):
+        """
+        Args:
+            policy
+        Returns:
+            gen: iterable which samples low-level transitions from the environment
+        """
+        return self.sample(policy)
+
+class HLOptions(OptionsEnv):
+    """Sample high-level (state, action, reward) tuples for generator training."""
+    
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _after_choice(self):
+        """
+        After an option selection, initialize total reward and number of steps
+        """
+        self.r = 0
+        self.steps = 0
+
+    def _after_step(self):
+        """
+        After each low-level action, add the discounted discriminated reward score (given a discriminator)
+        """
+        self.r += self.discount**self.steps * self.discriminator.discrim_net.reward_train(
+            state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()),
+            action=torch.tensor([[self.a]]).to(self.discriminator.discrim_net.device()),
+            next_state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused
+            done=torch.tensor(self.done).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused
+        )
+        self.steps += 1
+    
+    def _transitions(self):
+        """
+        Yield a single dictionary per high-level selected action
+        Fields:
+            obs: high-level state and mask at selection
+            action: chosen high-level action
+            reward: accumulated option reward
+            episode_start: whether the action was chosen at the episode start
+            value: the value estimate from the starting state
+            log_prob: the log_prob of the selected action from the starting state
+            done: whether the episode has ended
+
+        """
+        yield {
+            'obs': {'obs': self.s, 'mask': self.m},
+            'action': self.ch,
+            'reward': self.r.detach(),
+            'episode_start': self.episode_start,
+            'value': self.value.detach(),
+            'log_prob': self.log_prob.detach(),
+            'done': self.done,
+        }
+    
+    def sample_hl(self, policy, discriminator):
+        """
+        Args:
+            policy
+            discriminator: function with which to score rewards
+        Returns:
+            gen: iterable which samples high-level transitions from the environment
+        """
+        self.discriminator = discriminator
+        return self.sample(policy)
+
+class RenderOptions(LLOptions):
+
+    def _after_step(self):
+        """
+        Render the environment after each low-level step
+        """
+        super()._after_step()
+        self.env.render()
+    
+    def close(self, *args, **kwargs):
+        """
+        On 'close', close the environment
+        """
+        self.env.close(*args, **kwargs)
+
+def available_actions(env):
+    """Return mask of available actions given current `env` state."""
+    valid = np.array([feasible(env, generate_plan(env, i), i) for i in range(len(ALL_OPTIONS))])
+    return valid
+
+def target_velocity_plan(current_v: float, target_v: float, t: int, dt: float):
+    """Smoothly target a velocity in a given number of steps"""
+    # for now, constant acceleration
+    a = (target_v - current_v)  / (t * dt)
+    return a*np.ones((t,))
+
+def generate_plan(env, i):
+    """Generate input profile for high-level action `i`."""
+    assert i < len(ALL_OPTIONS), "Invalid option index {i}"
+    target_v, t = ALL_OPTIONS[i]
+    current_v = env._env.state[env._agent, 1].item() # extract from env
+    plan = target_velocity_plan(current_v, target_v, t, env._env._dt)
+    assert len(plan) == t, "incorrect plan length"
+    return plan
+
+def check_future_collisions_fast(env, actions):
+    """Checks whether `env._agent` would collide with other agents assuming `actions` as input.
+    
+    Vehicles are (over-)approximated by single circles.
+
+    Args:
+        env (gym.Env): current environment state
+        actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles
+    Returns:
+        feasible (torch.Tensor): tensor of shape (B,) indicating whether the respective action profiles are collision-free
+    """
+    B, (T, nv, _) = len(actions), actions[0].shape
+
+    states = torch.stack(env._env.propagate_action_profile(actions), axis=0)
+    assert states.shape == (B, T, nv, 5)
+
+    distance = ((states[:, :, :, :2] - states[:, :, env._agent:env._agent+1, :2])**2).sum(-1).sqrt()
+    distance = torch.where(distance.isnan(), np.inf*torch.ones_like(distance), distance) # only collide with spawned agents
+    distance[:, :, env._agent] = np.inf # cannot collide with itself
+    assert distance.shape == (B, T, nv)
+
+    radius = (env._env._lengths**2 + env._env._widths**2).sqrt() / 2
+    min_distance = radius[env._agent] + radius
+    min_distance = min_distance.unsqueeze(0).unsqueeze(0)
+    assert min_distance.shape == (1, 1, nv)
+
+    return (distance > min_distance).all(-1).all(-1)
+
+def check_future_collisions_circles(env, actions, n_circles:int=2):
+    """Checks whether `env._agent` would collide with other agents assuming `actions` as input.
+    
+    Vehicles are (over-)approximated by multiple circles.
+
+    Args:
+        env (gym.Env): current environment state
+        actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles
+    Returns:
+        feasible (torch.Tensor): tensor of shape (B,) indicating whether the respective action profiles are collision-free
+    """
+    assert n_circles >= 2
+    B, (T, nv, _) = len(actions), actions[0].shape
+
+    states = torch.stack(env._env.propagate_action_profile(actions), axis=0)
+    assert states.shape == (B, T, nv, 5)
+    centers = states[:, :, :, :2]
+    psi = states[:, :, :, 3]
+    lon = torch.stack([psi.cos(), psi.sin()],dim=-1) # (B, T, nv, 2)
+
+    # offset between [-env._env.lengths+env._env.widths/2, env._env.lengths/2-env._env.widths/2]
+    back = (-env._env._lengths/2+env._env._widths/2).unsqueeze(-1) # (nv, 1)
+    length = (env._env._lengths-env._env._widths).unsqueeze(-1) # (nv, 1)
+    diff_d =  back + length*(torch.arange(n_circles)/(n_circles-1)).unsqueeze(0)  # (nv, n_circles) 
+    assert diff_d.shape == (nv, n_circles)
+    
+    offsets = diff_d[None, None, :, :, None] *  lon[:, :, :, None, :]
+    assert offsets.shape == (B, T, nv, n_circles, 2)
+
+    expanded_centers=centers.unsqueeze(-2) + offsets #(B, T, nv, n_circles, 2)
+    assert expanded_centers.shape == (B, T, nv, n_circles, 2)
+    agent_centers = expanded_centers[:,:,env._agent:env._agent+1,:,:] #(B, T, 1, n_circles, 2)
+    ds = expanded_centers.reshape((B, T, nv*n_circles, 1, 2)) - agent_centers #(B, T, nv*nc,1, 2) - (B, T, 1, nc, 2) = (B, T, nv*nc, nc, 2) 
+
+    distance = (ds**2).sum(-1).sqrt().reshape((B, T, nv, n_circles, n_circles)) # (B, T, nv, nc, nc)
+    distance = torch.where(distance.isnan(), np.inf*torch.ones_like(distance), distance) # only collide with spawned agents
+    distance[:, :, env._agent] = np.inf # cannot collide with itself
+    assert distance.shape == (B, T, nv, n_circles, n_circles)
+
+    radius = env._env._widths*np.sqrt(2) / 2
+    min_distance = radius[env._agent] + radius
+    min_distance = min_distance[None, None, :, None, None]
+    assert min_distance.shape == (1, 1, nv, 1, 1)
+
+    return (distance > min_distance).all(-1).all(-1).all(-1).all(-1)
+
+def feasible(env, plan, ch):
+    """Check if input profile is feasible given current `env` state. Action `ch=0` is safe fallback."""
+    
+    # zero pad plan - Take (T,) np plan and convert it to (T, nv, 1) torch.Tensor
+    full_plan = torch.zeros(len(plan), env._env._nv, 1)
+    full_plan[:, env._agent, 0] = torch.tensor(plan)
+    # valid = check_future_collisions_fast(env, [full_plan]) # check_future_collisions_fast takes in B-list and outputs (B,) bool tensor
+    valid = check_future_collisions_circles(env, [full_plan])
+    return ch == 0 or valid.item()
+
+def flatten_transitions(transitions):
+    return {
+        'obs': np.stack(list(t['obs'] for t in transitions), axis=0),
+        'next_obs': np.stack(list(t['next_obs'] for t in transitions), axis=0),
+        'acts': np.stack(list(t['acts'] for t in transitions), axis=0),
+        'dones': np.stack(list(t['dones'] for t in transitions), axis=0),
+    }
+
+def train_discriminator(env, generator, discriminator, num_samples):
+    transitions = list(itertools.islice(env.sample_ll(generator), num_samples))
+    generator_samples = flatten_transitions(transitions)
+    discriminator.train_disc(gen_samples=generator_samples)
+
+def train_generator(env, generator, discriminator, num_samples):
+    generator_samples = list(itertools.islice(env.sample_hl(generator, discriminator), num_samples+1))
+    
+    generator.rollout_buffer.reset()
+    for s in generator_samples[:-1]:
+        generator.rollout_buffer.add(
+            obs=s['obs'],
+            action=s['action'].cpu(),
+            reward=s['reward'].cpu(),
+            episode_start=s['episode_start'],
+            value=s['value'],
+            log_prob=s['log_prob'],
+        )
+    
+    generator.rollout_buffer.compute_returns_and_advantage(
+        last_values=generator_samples[-1]['value'],
+        dones=generator_samples[-1]['done'],
+    )
+    
+    generator.train()
+
+def train(expert_data, env_class=NRasterizedRandomAgent, env_settings={}, epochs=10, discrim_batch_size=32, generator_steps=2048, discount=0.99):
+    """
+    Args: 
+        expert_data: list of transitions
+        env_class: environment class
+        env_settings: environment settings
+        epochs: number of epochs to train for
+        discrim_batch_size: discriminator batch size
+        generator_steps: number of steps taken in generator
+        discount: discount factor
+    Returns:
+        generator (stable_baselines3.PPO): options policy
+    """
+    env = env_class(**env_settings)
+    env.discount = discount
+
+    tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
+    tempdir_path = pathlib.Path(tempdir.name)
+    logger.configure(tempdir_path / "GAIL/")
+    print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
+
+    venv = make_vec_env(env_class, n_envs=1, env_kwargs=env_settings)
+    discriminator = adversarial.GAIL(
+        expert_data=expert_data,
+        expert_batch_size=discrim_batch_size,
+        discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
+        #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
+        venv=venv, # unused
+        gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
+    )
+
+    generator = stable_baselines3.PPO(
+        OptionsCnnPolicy,
+        OptionsEnv(env),
+        verbose=1,
+        n_steps=generator_steps,
+    )
+
+    # PPO.train requires logger as set up in
+    # PPO._setup_learn (called by PPO.learn)
+    generator._logger = stable_baselines3.common.utils.configure_logger(
+        generator.verbose,
+        generator.tensorboard_log,
+    )
+
+    for _ in tqdm(range(epochs)):
+        train_discriminator(LLOptions(env), generator, discriminator, num_samples=discrim_batch_size)
+        train_generator(HLOptions(env), generator, discriminator, num_samples=generator_steps)
+    
+    return generator
+
+# %%
+if __name__ == '__main__':
+    # %%
+    model_name = 'gail_options_image'
+    env_class = NRasterizedRandomAgent
+    env_settings = {'width': 36, 'height': 36, 'm_per_px': 2}
+
+    with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedIncrementingAgentw36h36mppx2.pkl", "rb") as f:
+        trajectories = pickle.load(f)
+    #import pdb 
+    #pdb.set_trace()
+    transitions = rollout.flatten_trajectories(trajectories)
+    generator = train(
+        transitions,
+        env_class=env_class,
+        env_settings=env_settings,
+        epochs=2, 
+        discrim_batch_size=32, 
+        generator_steps=2048, 
+        discount=0.99
+        )
+
+    generator.save(model_name) # save ppo sb3 generator class
+
+    # %%
+    model = stable_baselines3.PPO.load(model_name) # not actually used
+
+    env = RenderOptions(NRasterizedRandomAgent(**env_settings))
+    for s in env.sample_ll(generator):
+        if s['dones']:
+            break
+
+    env.close(filestr='render/'+model_name)
+
+# %% Tests
+
+def test_ll_expert_data():
+    with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
+        expert_trajectories = pickle.load(f)
+    expert_transitions = rollout.flatten_trajectories(expert_trajectories)
+
+    env = LLOptions(NRasterized(agent=51, width=36, height=36, m_per_px=2))
+
+    gen_transitions = list(itertools.islice(env.sample_ll(
+        policy=stable_baselines3.PPO(
+            OptionsCnnPolicy,
+            OptionsEnv(env),
+            verbose=1,
+        )
+    ), 10))
+    gen_transitions = flatten_transitions(gen_transitions)
+
+    assert expert_transitions[:10].obs.shape == gen_transitions['obs'].shape
+    assert expert_transitions[:10].next_obs.shape == gen_transitions['next_obs'].shape
+    assert expert_transitions[:10].acts.shape == gen_transitions['acts'].shape
+    assert expert_transitions[:10].dones.shape == gen_transitions['dones'].shape
+
+def test_ll_states():
+    env = NRasterized()
+    policy = stable_baselines3.PPO(
+        OptionsCnnPolicy,
+        OptionsEnv(env),
+        verbose=1,
+    )
+    llenv = LLOptions(env)
+    transitions = list(itertools.islice(llenv.sample_ll(policy=policy), 100))
+
+    env2 = NRasterized()
+    s2 = env2.reset()
+    for i, t in enumerate(transitions):
+        assert i == 0 or np.array_equal(t['obs'], transitions[i-1]['next_obs'])
+        assert np.array_equal(t['obs'], s2)
+        assert t['acts'].shape == (1,)
+
+        nexts2, _, done2, _ = env2.step(t['acts'])
+        assert np.array_equal(t['next_obs'], nexts2)
+        assert np.array_equal(t['dones'], done2)
+
+        if done2:
+            break
+
+        s2 = nexts2
+
+def test_hl_transitions():
+    pass

scratch/arec/intersimple/options_gail.py

@@ -0,0 +1,510 @@
+# %%
+from gail.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
+from imitation.algorithms import adversarial
+import stable_baselines3
+import torch.utils.data
+import numpy as np
+from intersim.envs.intersimple import NRasterized, NRasterizedRandomAgent
+import itertools
+from torch.distributions import Categorical
+import gym
+import torch
+import pickle
+import imitation.data.rollout as rollout
+import tempfile
+import pathlib
+from imitation.util import logger
+from stable_baselines3.common.env_util import make_vec_env
+from tqdm import tqdm
+
+import logging 
+logging.basicConfig(level=logging.DEBUG)
+
+ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback
+
+class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy):
+    """
+    Class for high-level options policy (generator)
+    """
+    def __init__(self, observation_space, *args, **kwargs):
+        super().__init__(observation_space['obs'], *args, **kwargs)
+
+    def _prior_distribution(self, s):
+        """
+        Return prior distribution over high-level options (before masking)
+        Args:
+            s (torch.tensor): observation
+        Returns:
+            values (torch.tensor): values from critic
+            dist (torch.distributions): prior distribution over actions
+        """
+        latent_pi, latent_vf, latent_sde = self._get_latent(s) 
+        distribution = self._get_action_dist_from_latent(latent_pi, latent_sde) 
+        values = self.value_net(latent_vf)
+        return values, distribution.distribution
+
+    def predict(self, obs):
+        """
+        Will mask invalid states before making action selections
+        Args:
+            obs: dict with keys:
+                obs (torch.tensor): (B,o)  true observations
+                mask (torch.tensor): (B,m) mask over valid actions
+        Returns:
+            ch (torch.tensor): (B,a) sampled actions
+            values (torch.tensor): (B,) predicted value at observation
+            log_probs (torch.tensor): (B,) log probabilities of selected actions
+        """
+        s, m = obs['obs'], obs['mask']
+        values, prior = self._prior_distribution(s)
+        posterior = Categorical(prior.probs * m)
+        ch = posterior.sample()
+        return ch, values, posterior.log_prob(ch)
+
+    def evaluate_actions(self, obs, ch):
+        """
+        Evaluate particular actions
+        Args:
+            obs: dict with keys:
+                obs (torch.tensor): (B,o) true observations
+                mask (torch.tensor): (B,m) masks over valid actions
+            ch (torch.tensor): (B,a) selected actions
+        Returns:
+            values (torch.tensor): (B,) predicted value at observation
+            log_probs (torch.tensor): (B,) log probabilities of selected actions
+            ent (torch.tensor): (B,) entropy of each distribution over actions
+        """
+        s, m = obs['obs'], obs['mask']
+        values, prior = self._prior_distribution(s)
+        posterior = Categorical(prior.probs * m)
+        return values, posterior.log_prob(ch), posterior.entropy() # additional values used by PPO.train
+
+class OptionsEnv(gym.Wrapper):
+    """
+    Wrap an intersimple environment with an options generator
+    """
+    def __init__(self, env, render=False, *args, **kwargs):
+        """
+        Initialize wrapped environment and set high-level action and observation spaces
+        """
+        super().__init__(env, *args, **kwargs)
+        num_hl_options = len(ALL_OPTIONS)
+        self.action_space = gym.spaces.Discrete(num_hl_options)
+        self.observation_space = gym.spaces.Dict({
+            'obs': env.observation_space,
+            'mask': gym.spaces.Box(low=0, high=1, shape=(num_hl_options,)),
+        })
+        self._hl_transition_buffer = []
+        self._ll_transition_buffer = []
+        self.render=render
+
+    def _after_option_choice(self):
+        """
+        After initial option choice, 
+        """
+        self._hl_r = 0
+        self._hl_steps = 0
+    
+    def _after_step(self):
+        """
+        After each step, add the ll transition to the appropriate buffer, add to reward, add to steps, and possibly render
+        """
+
+        self._ll_transition_buffer.append({
+            'obs': self.s,
+            'next_obs': self.nexts,
+            'acts': np.array((self.a,)),
+            'dones': np.array(self.done),
+        })  
+        self.r += self.discount**self.steps * self.discriminator.discrim_net.reward_train(
+            state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()),
+            action=torch.tensor([[self.a]]).to(self.discriminator.discrim_net.device()),
+            next_state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused
+            done=torch.tensor(self.done).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused
+        )
+        self.steps += 1
+        if self.render:
+            self.env.render()
+    
+    def _after_option(self):
+        """
+        After each low-level action, add the discounted discriminated reward score (given a discriminator)
+        """
+        self._hl_transition_buffer.append({
+            'obs': {'obs': self.os, 'mask': self.m},
+            'action': self.ch,
+            'reward': self.r.detach(),
+            'episode_start': self.episode_start,
+            'value': self.value.detach(),
+            'log_prob': self.log_prob.detach(),
+            'done': self.done,
+            })
+
+    def close(self, *args, **kwargs):
+        """
+        On 'close', close the environment
+        """
+        self.env.close(*args, **kwargs)
+    
+    def sample(self, generator, controller):
+        """
+        yield transitions using a generator
+        Args:
+            generator (sb3.PPO)
+            controller (str): 'high' or 'low' to yield from proper buffer
+        Yields:
+
+        """
+        self.done = True
+        # DO I WANT TO EMPTY THE BUFFERS??? Probs naw
+        while True:
+            
+            # yield from buffers to empty what was stored previously
+            if controller = 'high':
+                yield from self._hl_transition_buffer
+            elif controller == 'low':
+                yield from self._ll_transition_buffer
+            else:
+                raise('Improper buffer')
+            
+            self.episode_start = False
+            if self.done:
+                # reset environment
+                self.s = self.env.reset()
+                self.done = False
+                self.episode_start = True
+
+            self.os = self.s.copy() # option start state
+            self.m = available_actions(self.env)
+
+            # set the action, the value of the start state, and the logprob of the action
+            # according to the current environment state and mask
+            self.ch, self.value, self.log_prob = generator.policy.predict({
+                'obs': torch.tensor(self.os).unsqueeze(0).to(generator.policy.device),
+                'mask': torch.tensor(self.m).unsqueeze(0).to(generator.policy.device),
+            })
+
+            # store a float list of actions to take given the option selected in the environment
+            self.plan = list(map(float, generate_plan(self.env, self.ch)))
+
+            # run whatever _after_choice might dictate in a child class
+            self._after_option_choice()
+
+            # some checks
+            assert not self.done
+            assert self.plan
+            assert feasible(self.env, self.plan, self.ch)
+
+            # execute the option so long as the episode isn't complete and the plan is still feasible
+            while not self.done and self.plan and feasible(self.env, self.plan, self.ch):
+                
+                # pop first action
+                self.a, self.plan = self.plan[0], self.plan[1:]
+                
+                # normalize action ??
+                self.a = self.env._normalize(self.a)
+                
+                # step through environment
+                self.nexts, _, self.done, _ = self.env.step(self.a)
+
+                # run whatever _after_step might dictate in child class
+                self._after_step()
+
+                # update state and mask to current
+                self.s = self.nexts
+            
+            # run whatever to do after option
+            self._after_option()
+
+    def sample_ll(self, policy):
+        """
+        Not quite sure how this works????
+        Why would you do this over LLOptions.sample(policy)
+        """
+        return self.sample(policy, 'low')
+
+    def sample_hl(self, policy, discriminator):
+        """
+        Args:
+            policy
+            discriminator: function with which to score rewards
+        Returns:
+            gen: an which samples high-level transitions from the environment
+        """
+        self.discriminator = discriminator
+        return self.sample(policy)
+
+def available_actions(env):
+    """Return mask of available actions given current `env` state."""
+    valid = np.array([feasible(env, generate_plan(env, i), i) for i in range(len(ALL_OPTIONS))])
+    return valid
+
+def target_velocity_plan(current_v: float, target_v: float, t: int, dt: float):
+    """Smoothly target a velocity in a given number of steps"""
+    # for now, constant acceleration
+    a = (target_v - current_v)  / (t * dt)
+    return a*np.ones((t,))
+
+def generate_plan(env, i):
+    """Generate input profile for high-level action `i`."""
+    assert i < len(ALL_OPTIONS), "Invalid option index {i}"
+    target_v, t = ALL_OPTIONS[i]
+    current_v = env._env.state[env._agent, 1].item() # extract from env
+    plan = target_velocity_plan(current_v, target_v, t, env._env._dt)
+    assert len(plan) == t, "incorrect plan length"
+    return plan
+
+def check_future_collisions_fast(env, actions):
+    """Checks whether `env._agent` would collide with other agents assuming `actions` as input.
+    
+    Vehicles are (over-)approximated by single circles.
+
+    Args:
+        env (gym.Env): current environment state
+        actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles
+    Returns:
+        feasible (torch.Tensor): tensor of shape (B,) indicating whether the respective action profiles are collision-free
+    """
+    B, (T, nv, _) = len(actions), actions[0].shape
+
+    states = torch.stack(env._env.propagate_action_profile(actions), axis=0)
+    assert states.shape == (B, T, nv, 5)
+
+    distance = ((states[:, :, :, :2] - states[:, :, env._agent:env._agent+1, :2])**2).sum(-1).sqrt()
+    distance = torch.where(distance.isnan(), np.inf*torch.ones_like(distance), distance) # only collide with spawned agents
+    distance[:, :, env._agent] = np.inf # cannot collide with itself
+    assert distance.shape == (B, T, nv)
+
+    radius = (env._env._lengths**2 + env._env._widths**2).sqrt() / 2
+    min_distance = radius[env._agent] + radius
+    min_distance = min_distance.unsqueeze(0).unsqueeze(0)
+    assert min_distance.shape == (1, 1, nv)
+
+    return (distance > min_distance).all(-1).all(-1)
+
+def check_future_collisions_circles(env, actions, n_circles:int=2):
+    """Checks whether `env._agent` would collide with other agents assuming `actions` as input.
+    
+    Vehicles are (over-)approximated by multiple circles.
+
+    Args:
+        env (gym.Env): current environment state
+        actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles
+    Returns:
+        feasible (torch.Tensor): tensor of shape (B,) indicating whether the respective action profiles are collision-free
+    """
+    assert n_circles >= 2
+    B, (T, nv, _) = len(actions), actions[0].shape
+
+    states = torch.stack(env._env.propagate_action_profile(actions), axis=0)
+    assert states.shape == (B, T, nv, 5)
+    centers = states[:, :, :, :2]
+    psi = states[:, :, :, 3]
+    lon = torch.stack([psi.cos(), psi.sin()],dim=-1) # (B, T, nv, 2)
+
+    # offset between [-env._env.lengths+env._env.widths/2, env._env.lengths/2-env._env.widths/2]
+    back = (-env._env._lengths/2+env._env._widths/2).unsqueeze(-1) # (nv, 1)
+    length = (env._env._lengths-env._env._widths).unsqueeze(-1) # (nv, 1)
+    diff_d =  back + length*(torch.arange(n_circles)/(n_circles-1)).unsqueeze(0)  # (nv, n_circles) 
+    assert diff_d.shape == (nv, n_circles)
+    
+    offsets = diff_d[None, None, :, :, None] *  lon[:, :, :, None, :]
+    assert offsets.shape == (B, T, nv, n_circles, 2)
+
+    expanded_centers=centers.unsqueeze(-2) + offsets #(B, T, nv, n_circles, 2)
+    assert expanded_centers.shape == (B, T, nv, n_circles, 2)
+    agent_centers = expanded_centers[:,:,env._agent:env._agent+1,:,:] #(B, T, 1, n_circles, 2)
+    ds = expanded_centers.reshape((B, T, nv*n_circles, 1, 2)) - agent_centers #(B, T, nv*nc,1, 2) - (B, T, 1, nc, 2) = (B, T, nv*nc, nc, 2) 
+
+    distance = (ds**2).sum(-1).sqrt().reshape((B, T, nv, n_circles, n_circles)) # (B, T, nv, nc, nc)
+    distance = torch.where(distance.isnan(), np.inf*torch.ones_like(distance), distance) # only collide with spawned agents
+    distance[:, :, env._agent] = np.inf # cannot collide with itself
+    assert distance.shape == (B, T, nv, n_circles, n_circles)
+
+    radius = env._env._widths*np.sqrt(2) / 2
+    min_distance = radius[env._agent] + radius
+    min_distance = min_distance[None, None, :, None, None]
+    assert min_distance.shape == (1, 1, nv, 1, 1)
+
+    return (distance > min_distance).all(-1).all(-1).all(-1).all(-1)
+
+def feasible(env, plan, ch):
+    """Check if input profile is feasible given current `env` state. Action `ch=0` is safe fallback."""
+    
+    # zero pad plan - Take (T,) np plan and convert it to (T, nv, 1) torch.Tensor
+    full_plan = torch.zeros(len(plan), env._env._nv, 1)
+    full_plan[:, env._agent, 0] = torch.tensor(plan)
+    # valid = check_future_collisions_fast(env, [full_plan]) # check_future_collisions_fast takes in B-list and outputs (B,) bool tensor
+    valid = check_future_collisions_circles(env, [full_plan])
+    return ch == 0 or valid.item()
+
+def flatten_transitions(transitions):
+    return {
+        'obs': np.stack(list(t['obs'] for t in transitions), axis=0),
+        'next_obs': np.stack(list(t['next_obs'] for t in transitions), axis=0),
+        'acts': np.stack(list(t['acts'] for t in transitions), axis=0),
+        'dones': np.stack(list(t['dones'] for t in transitions), axis=0),
+    }
+
+def train_discriminator(env, generator, discriminator, num_samples):
+    transitions = list(itertools.islice(env.sample_ll(generator), num_samples))
+    generator_samples = flatten_transitions(transitions)
+    discriminator.train_disc(gen_samples=generator_samples)
+
+def train_generator(env, generator, discriminator, num_samples):
+    generator_samples = list(itertools.islice(env.sample_hl(generator, discriminator), num_samples+1))
+    
+    generator.rollout_buffer.reset()
+    for s in generator_samples[:-1]:
+        generator.rollout_buffer.add(
+            obs=s['obs'],
+            action=s['action'].cpu(),
+            reward=s['reward'].cpu(),
+            episode_start=s['episode_start'],
+            value=s['value'],
+            log_prob=s['log_prob'],
+        )
+    
+    generator.rollout_buffer.compute_returns_and_advantage(
+        last_values=generator_samples[-1]['value'],
+        dones=generator_samples[-1]['done'],
+    )
+    
+    generator.train()
+
+def train(expert_data, env_class=NRasterizedRandomAgent, env_settings={}, epochs=10, discrim_batch_size=32, generator_steps=2048, discount=0.99):
+    """
+    Args: 
+        expert_data: list of transitions
+        env_class: environment class
+        env_settings: environment settings
+        epochs: number of epochs to train for
+        discrim_batch_size: discriminator batch size
+        generator_steps: number of steps taken in generator
+        discount: discount factor
+    Returns:
+        generator (stable_baselines3.PPO): options policy
+    """
+    env = env_class(**env_settings)
+    env.discount = discount
+
+    tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
+    tempdir_path = pathlib.Path(tempdir.name)
+    logger.configure(tempdir_path / "GAIL/")
+    print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
+
+    venv = make_vec_env(env_class, n_envs=1, env_kwargs=env_settings)
+    discriminator = adversarial.GAIL(
+        expert_data=expert_data,
+        expert_batch_size=discrim_batch_size,
+        discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
+        #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
+        venv=venv, # unused
+        gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
+    )
+
+    generator = stable_baselines3.PPO(
+        OptionsCnnPolicy,
+        OptionsEnv(env),
+        verbose=1,
+        n_steps=generator_steps,
+    )
+
+    # PPO.train requires logger as set up in
+    # PPO._setup_learn (called by PPO.learn)
+    generator._logger = stable_baselines3.common.utils.configure_logger(
+        generator.verbose,
+        generator.tensorboard_log,
+    )
+
+    for _ in tqdm(range(epochs)):
+        train_discriminator(LLOptions(env), generator, discriminator, num_samples=discrim_batch_size)
+        train_generator(HLOptions(env), generator, discriminator, num_samples=generator_steps)
+    
+    return generator
+
+# %%
+if __name__ == '__main__':
+    # %%
+    model_name = 'gail_options_image'
+    env_class = NRasterizedRandomAgent
+    env_settings = {'width': 36, 'height': 36, 'm_per_px': 2}
+
+    with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedIncrementingAgentw36h36mppx2.pkl", "rb") as f:
+        trajectories = pickle.load(f)
+    #import pdb 
+    #pdb.set_trace()
+    transitions = rollout.flatten_trajectories(trajectories)
+    generator = train(
+        transitions,
+        env_class=env_class,
+        env_settings=env_settings,
+        epochs=2, 
+        discrim_batch_size=32, 
+        generator_steps=2048, 
+        discount=0.99
+        )
+
+    generator.save(model_name) # save ppo sb3 generator class
+
+    # %%
+    model = stable_baselines3.PPO.load(model_name) # not actually used
+
+    env = OptionsGail(NRasterizedRandomAgent(**env_settings), render=True)
+    for s in env.sample_ll(generator):
+        if s['dones']:
+            break
+
+    env.close(filestr='render/'+model_name)
+
+# %% Tests
+
+def test_ll_expert_data():
+    with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
+        expert_trajectories = pickle.load(f)
+    expert_transitions = rollout.flatten_trajectories(expert_trajectories)
+
+    env = LLOptions(NRasterized(agent=51, width=36, height=36, m_per_px=2))
+
+    gen_transitions = list(itertools.islice(env.sample_ll(
+        policy=stable_baselines3.PPO(
+            OptionsCnnPolicy,
+            OptionsEnv(env),
+            verbose=1,
+        )
+    ), 10))
+    gen_transitions = flatten_transitions(gen_transitions)
+
+    assert expert_transitions[:10].obs.shape == gen_transitions['obs'].shape
+    assert expert_transitions[:10].next_obs.shape == gen_transitions['next_obs'].shape
+    assert expert_transitions[:10].acts.shape == gen_transitions['acts'].shape
+    assert expert_transitions[:10].dones.shape == gen_transitions['dones'].shape
+
+def test_ll_states():
+    env = NRasterized()
+    policy = stable_baselines3.PPO(
+        OptionsCnnPolicy,
+        OptionsEnv(env),
+        verbose=1,
+    )
+    llenv = LLOptions(env)
+    transitions = list(itertools.islice(llenv.sample_ll(policy=policy), 100))
+
+    env2 = NRasterized()
+    s2 = env2.reset()
+    for i, t in enumerate(transitions):
+        assert i == 0 or np.array_equal(t['obs'], transitions[i-1]['next_obs'])
+        assert np.array_equal(t['obs'], s2)
+        assert t['acts'].shape == (1,)
+
+        nexts2, _, done2, _ = env2.step(t['acts'])
+        assert np.array_equal(t['next_obs'], nexts2)
+        assert np.array_equal(t['dones'], done2)
+
+        if done2:
+            break
+
+        s2 = nexts2
+
+def test_hl_transitions():
+    pass

scratch/arec/intersimple/plan.txt

@@ -0,0 +1,55 @@
+Environment 
+-- each 'environment' follows a single roundabout and track id (recording of that roundabout)
+-- on reset, the environment we will use changes the vehicle to control while having the other agents follow their true data (expert controller)
+---- Note this can be problematic as it can lead to vehicles behind you crashing into you
+
+TRAINING
+---------
+1. Load pre-trained massive set of transitions
+-- For all roundabouts
+    -- For all tracks
+        -- For all vehicles
+            -- For all valid timesteps
+                -- Rasterized state (incl. path), action
+
+2. HGAIL
+-- For each epoch
+    -- INSTANTIATE A NEW ENVIRONMENT (Roundabout + Track) w/ randomized agent, from set of all expert environments
+        -- Train discriminator off training data + yielded low-level transitions in replay buffer
+        -- Train generator off yielded high-level transitions + summed low-level discriminator rewards
+
+TESTING
+----------
+1. Save average vehicle velocities for all expert vehicles (loop roundabout + track + vehicle, average over time)
+
+2. Run test suite for: expert, BC, GAIL, RAIL, HGAIL, (and hopefully HRAIL)
+-- For all roundabouts, tracks
+    -- Get expert velocities for track
+    -- Simulate incrementing agent environment (e.g. on reset, agent +=1)
+        -- Store low-level true joint states, actions, and controlled vehicle index
+    -- Per-vehicle statistics (v_all, v_mean, v_shortfall, a_all, jerk_all, n_collisions, T) 
+-- Aggregate statistics + joint 
+
+Problems
+-----------
+Should train without stopping for collisions, however when doing so, end up with policy that always takes decelerate option
+-- It seems safe at the start of each vehicles sim, but actually it isn't since a car will spawn and hit it
+Solutions:
+    -- Hold cars from spawning if their spawn location is full
+    -- Start simulations a few seconds later (after cars clear their spawn places) <- Preferred
+    
+Test could run indefinitely if stop_on_collision is off
+Solution:
+    -- Set maximum episode length in intersimple
+
+
+
+
+
+
+
+
+Save massive set of transition raw states beforehand (1 from training, but with raw states)
+# -- For all roundabouts, tracks
+#     -- For all vehicles, steps
+#         -- Raw vehicle state, action

scratch/arec/intersimple/render_env_from_model.py

@@ -0,0 +1,58 @@
+
+import stable_baselines3 as sb3
+from intersim.envs.intersimple import NRasterized
+
+
+def render_env(model_name='gail_image_multiagent_nocollision', agent=51, environment=NRasterized):
+    """
+    Render a video from an model, agent, and environment
+    Args:
+        model_name (str): name of the model
+        agent (int): agent to start the video from
+        environment (gym.Env): gym environment class to render environment on
+    """
+
+    model = sb3.PPO.load(model_name)
+
+    env = environment(stop_on_collision=False, width=36, height=36, m_per_px=2, agent=agent)
+
+    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+'_agent%i'%(agent)) 
+
+def render_options_env(model_name='gail_image_multiagent_nocollision', agent=51, environment=NRasterized):
+    """
+    Render a video from an model, agent, and environment
+    Args:
+        model_name (str): name of the model
+        agent (int): agent to start the video from
+        environment (gym.Env): gym environment class to render environment on
+    """
+
+    model = sb3.PPO.load(model_name)
+
+    env = environment(stop_on_collision=False, width=36, height=36, m_per_px=2, agent=agent)
+
+    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+'_agent%i'%(agent)) 
+
+if __name__ == '__main__':
+    import fire
+    fire.Fire(render_env)

scratch/arec/intersimple/render_options.py

@@ -0,0 +1,11 @@
+import sys
+sys.path.append('../../../')
+from src.util import render_env
+ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10]]
+
+def render_wrapper(**kwargs):
+    render_env(**kwargs, options_list=ALL_OPTIONS)
+
+if __name__=='__main__':
+    import fire
+    fire.Fire(render_wrapper)

scratch/arec/intersimple/test_model.py

@@ -0,0 +1,148 @@
+from tqdm import tqdm
+from copy import deepcopy
+import stable_baselines3 as sb3
+import intersim
+
+ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10]] # option 0 is safe fallback
+
+def load_model(model_path:str, method:str):
+    """
+    Load a model given a path and the method
+    
+    Args:
+        model_path (str): the path to the model
+        method (str): the method for the model
+    Returns:
+        model: the action model
+        is_heir (bool): whether the method is heirarchial
+    """
+    model = None
+    is_heir = False
+    if method == 'expert':
+        raise NotImplementedError
+    elif method == 'bc':
+        raise NotImplementedError
+    elif method == 'gail':
+        raise NotImplementedError
+    elif method == 'rail':
+        raise NotImplementedError
+    elif method == 'hgail':
+        is_heir = True
+        model = sb3.PPO.load(model_path)
+    elif method == 'hrail':
+        is_heir = True
+        raise NotImplementedError
+    else:
+        raise NotImplementedError
+    return model, is_heir
+
+def load_expert_states(roundabout, track):
+    """
+    Load expert states from roundabout/track info
+    Args:
+        roundabout (str): roundabout name
+        track (str): track id
+    Returns:
+        states (torch.tensor): (T+1, nv, 5) expert states for track file
+        actions (torch.tensor): (T, nv, 1) expert actions for track file
+    """
+    state_path = '../../../expert_data/%s/track%04i/joint_expert_states.pt'%(roundabout, track)] #FIXME when moving
+    action_path = '../../../expert_data/%s/track%04i/joint_expert_actions.pt'%(roundabout, track)] #FIXME when moving
+    states = torch.load(path)
+    actions = torch.load(path)
+    # nanify actions where vehicle's don't exist
+    import pdb
+    pdb.set_trace()
+    return states, actions
+
+def test_model(
+    locations=[(0,0)], 
+    model_name='gail_image_multiagent_nocollision', 
+    env='NRasterizedRouteIncrementingAgent', 
+    method='expert',
+    options_list=ALL_OPTIONS, 
+    **env_kwargs):
+    """
+    Test a particular model at different locations/tracks
+
+    Args:
+        locations (list of tuples): list of (roundabout, track) integer pairs
+        model_name (str): name of model to test
+        env (str): environment class
+        method (str): method (expert, bc, gail, rail, hgail, hrail)
+        options_list (list): list of options
+    """
+
+    # load policy
+    policy, is_heir = load_model(model_name, method)
+
+    # iterate through vehicles
+    all_vehicle_infos = []
+    for i, location in tqdm(enumerate(locations)):
+        
+        # add roundabout and track to environent 
+        roundabout, track = location
+        iround = intersim.LOCATIONS.index(roundabout)
+        it_env_kwargs = deepcopy(env_kwargs)
+        loc_kwargs = {
+            'loc':iround,
+            'track':track
+        }
+        it_env_kwargs.update(loc_kwargs)
+        
+        # load expert states and get average velocities
+        expert_states, expert_actions = load_expert_states(roundabout, track)
+        expert_vavg = torch.nanmean(expert_states[:,:,3], dim=-1)
+
+        # initialize environment
+        if not is_heir:
+            Env = src.options.envs.__dict__[env]
+        else: 
+            Env = intersim.envs.intersimple.__dict__[env]
+        env = Env(**env_kwargs)
+        s = env.reset()
+
+        # Iterate through every vehicle and time
+        vehicle_infos, done = [], False
+        for iv in range(env.nv):
+            v_number = env.agent
+            i_vehicle_infos = {'s':[], 'a':[], 'it':[]}
+            while not done:
+                a = policy(s)
+                sp, r, done, info   = env.step(a)
+                i_vehicle_infos['s'].append(env._env.state) # FIX
+                i_vehicle_infos['a'].append(a) 
+                i_vehicle_infos['it'].append(env._env.it) # FIX
+            i_vehicle_info.update({
+                'vehicle_id': env.agent,
+                'n_steps': len(i_vehicle_infos['a']),
+                'T': len(i_vehicle_infos['a'])*env._env.dt, # FIX
+                'n_collisions': collision.check(i_vehicle_infos['s'], env._env.lengths. env._env.widths), # FIX
+                'expert_vavg': expert_vavg[env.agent]
+            })
+            vehicle_infos.append(i_vehicle_info)
+            env.reset()
+
+        all_vehicle_infos.append({
+            'loc': location,
+            'track': track,
+            'stats': vehicle_infos
+            })
+        env.close()
+
+    # print and save model-specific metrics
+    outfolder = 'test_metrics'
+    print_and_save(all_vehicle_infos, method, model, outfolder)
+
+def print_and_save(stats, method, model, outfolder):
+    """
+    Print and save stats
+    """
+    pass
+
+def load_compare():
+    pass
+
+if __name__=='__main__':
+    import fire
+    fire.Fire()

scratch/etienne/intersimple/data/expert.py

@@ -1,7 +1,8 @@
-from intersim.envs.intersimple import Intersimple
+from intersim.envs.intersimple import Intersimple, InfoFilter
 from stable_baselines3.common.policies import BasePolicy
 import gym
-import intersim.envs.intersimple
+from intersim.envs.intersimple import *
+from gail.envs import *
 import imitation.data.rollout as rollout
 from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
 from imitation.data.wrappers import RolloutInfoWrapper
@@ -49,6 +50,8 @@ class IntersimpleExpert(BasePolicy):
         raise NotImplementedError()
 
     def _action(self):
+        # RandomLocation mixin re-initializes the intersim sub-env
+        self._intersim_expert._intersim = self._intersimple._env
         return self._intersim_expert._action()[self._intersimple._agent]
 
     def predict(self, *args, **kwargs):
@@ -63,16 +66,28 @@ class NormalizedIntersimpleExpert(IntersimpleExpert):
         action, _ = super().predict(*args, **kwargs)
         return self._intersimple._normalize(action), None
 
-class DummyVecEnvPolicy():
+class DummyVecEnvPolicy(BasePolicy):
 
     def __init__(self, experts):
         self._experts = [e() for e in experts]
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def _predict(self, *args, **kwargs):
+        raise NotImplementedError()
     
     def predict(self, *args, **kwargs):
         predictions = [e.predict() for e in self._experts]
         actions = [p[0] for p in predictions]
         states = [p[1] for p in predictions]
         return actions, states
+    
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def _predict(self, *args, **kwargs):
+        raise NotImplementedError()
 
 def save_video(env, expert):
     env.reset()
@@ -91,7 +106,7 @@ def demonstrations(expert='NormalizedIntersimpleExpert', env='NRasterizedRandomA
         python -m intersimple.expert <flags>
 
     """
-    Env = intersim.envs.intersimple.__dict__[env]
+    Env = globals()[env]
     Expert = globals()[expert]
 
     env = Env(**env_args)
@@ -105,6 +120,7 @@ def demonstrations(expert='NormalizedIntersimpleExpert', env='NRasterizedRandomA
         save_video(env, policy)
     
     path = path or (policy.__class__.__name__ + '_' + env.__class__.__name__ + '.pkl')
+    include_infos = isinstance(env, InfoFilter)
 
     rollout.rollout_and_save(
         path=path,
@@ -113,7 +129,8 @@ def demonstrations(expert='NormalizedIntersimpleExpert', env='NRasterizedRandomA
         sample_until=rollout.make_sample_until(
             min_timesteps=min_timesteps,
             min_episodes=min_episodes,
-        )
+        ),
+        exclude_infos=not include_infos,
     )
 
 if __name__ == '__main__':

scratch/etienne/intersimple/data/generate.sh

@@ -1,3 +1,15 @@
-#python -m intersimple.expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51.pkl'
+#python -m expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51.pkl'
-#python -m intersimple.expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --policy_args='{mu:0.005}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.005.pkl'
+#python -m expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --policy_args='{mu:0.005}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.005.pkl'
-python -m intersimple.expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.001.pkl' --video
+#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=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'
+#python -m expert --env=NRasterizedRandomAgent --min_timesteps=200 --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedRandomAgentw36h36mppx2.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=NRasterizedRouteRandomAgent --min_timesteps=10000 --env_args='{width:70,height:70,m_per_px:1}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteRandomAgentw70h70mppx1.pkl'
+#python -m expert --env=NRasterizedRouteRandomAgentLocation --min_timesteps=100000 --env_args='{width:70,height:70,m_per_px:1}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N100000_NRasterizedRouteRandomAgentLocationw70h70mppx1.pkl'
+#python -m expert --env=NRasterizedRouteRandomAgentLocation --min_timesteps=100000 --env_args='{width:70,height:70,m_per_px:1,map_color:128}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N100000_NRasterizedRouteRandomAgentLocationw70h70mppx1mapc128.pkl'
+#python -m expert --env=NRasterizedRouteSpeedRandomAgentLocation --min_timesteps=10000 --env_args='{width:70,height:70,m_per_px:1,map_color:128,mu:0.001}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteSpeedRandomAgentLocationw70h70mppx1mapc128mu.001.pkl'
+#python -m data.expert --env=NRasterizedRouteSpeedRandomAgentLocation --min_timesteps=10000 --env_args='{width:70,height:70,m_per_px:1,map_color:128,mu:0.001,skip_frames:5}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteSpeedRandomAgentLocationw70h70mppx1mapc128mu.001skip5.pkl'
+#python -m data.expert --env=TLNRasterizedRouteRandomAgentLocation --min_timesteps=100000 --env_args='{width:70,height:70,m_per_px:1,mu:0.001,random_skip:True,max_episode_steps:50}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N100000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl'
+python -m data.expert --env=TLNRasterizedRouteRandomAgentLocation --min_timesteps=50000 --env_args='{width:70,height:70,m_per_px:1,mu:0.001,random_skip:True,max_episode_steps:50}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N50000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl'

scratch/etienne/intersimple/gail/discriminator.py

@@ -1,50 +0,0 @@
-import torch
-
-class CnnDiscriminator(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 + action_size
-
-        self.cnn = torch.nn.Sequential(
-            torch.nn.Conv2d(in_channels, 32, kernel_size=(8, 8), stride=(4, 4)), # 5+1 -> 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(512), # 28224 -> 512
-            torch.nn.ReLU(),
-            torch.nn.LazyLinear(1), # 512 -> 1
-        )
-    
-    def forward(self, state, action):
-        b, _, h, w = state.shape
-        _, a = action.shape
-        act_layer = action.unsqueeze(-1).unsqueeze(-1).expand((b, a, h, w))
-        sa = torch.cat((act_layer, state), -3)
-        return self.cnn(sa).squeeze()
-
-class MlpDiscriminator(torch.nn.Module):
-    """MLP similar to stable_baselines3.common.policies.ActorCriticPolicy."""
-
-    def __init__(self, env=None):
-        super().__init__()
-        self.flatten = torch.nn.Flatten(start_dim=1, end_dim=-1)
-        self.mlp = torch.nn.Sequential(
-            torch.nn.LazyLinear(64), # 42 -> 64
-            torch.nn.Tanh(),
-            torch.nn.LazyLinear(64), # 64 -> 64
-            torch.nn.Tanh(),
-            torch.nn.LazyLinear(1),  # 64 -> 1
-        )
-    
-    def forward(self, state, action):
-        flat = self.flatten(state)
-        sa = torch.cat((action, flat), -1)
-        return self.mlp(sa).squeeze()

scratch/etienne/intersimple/gail/envs.py

@@ -0,0 +1,46 @@
+import gym
+from gym.wrappers.time_limit import TimeLimit
+import numpy as np
+from intersim.envs.intersimple import NRasterizedRouteRandomAgentLocation, RandomLocation, RandomAgent, RewardVisualization, Reward, \
+    ImageObservationAnimation, RasterizedRoute, NObservations, RasterizedObservation, \
+    NormalizedActionSpace, ActionVisualization, InteractionSimulatorMarkerViz, ImitationCompat, Intersimple
+
+class RasterizedSpeed:
+
+    def __init__(self, max_speed=12, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        channels, height, width = self.observation_space.shape
+        self.observation_space = gym.spaces.Box(
+            low=0,
+            high=255,
+            shape=(channels+1, height, width),
+            dtype=np.uint8
+        )
+        self._max_speed = max_speed
+
+    def _simple_obs(self, intersim_obs, intersim_info):
+        img = super()._simple_obs(intersim_obs, intersim_info)
+
+        ego_speed = intersim_obs['state'][self._agent, 2]
+        scaled_speed = (255 * ego_speed) // self._max_speed
+        speed_layer = scaled_speed * np.ones_like(img[:1], dtype=np.uint8)
+        speed_layer = speed_layer.clamp(0, 255)
+        
+        obs = np.concatenate((img, speed_layer), axis=0)
+        return obs
+
+class NRasterizedRouteSpeedRandomAgentLocation(RandomLocation, RandomAgent, RewardVisualization,
+        Reward, ImageObservationAnimation, RasterizedRoute, NObservations, RasterizedSpeed, RasterizedObservation,
+        NormalizedActionSpace, ActionVisualization, InteractionSimulatorMarkerViz, ImitationCompat, Intersimple):
+    pass
+
+class TransparentTimeLimit(TimeLimit):
+
+    def __getattr__(self, name):
+        return getattr(self.env, name)
+    
+    def close(self, *args, **kwargs):
+        return self.env.close(*args, **kwargs)
+
+def TLNRasterizedRouteRandomAgentLocation(max_episode_steps, *args, **kwargs):
+    return TransparentTimeLimit(NRasterizedRouteRandomAgentLocation(*args, **kwargs), max_episode_steps=max_episode_steps)

scratch/etienne/intersimple/gail/options2.py

@@ -0,0 +1,127 @@
+import gym
+import torch
+from src.util.collisions import feasible
+import numpy as np
+from collections import deque
+
+def imitation_discriminator(discriminator):
+    return lambda obs, action, next_obs, done: discriminator.discrim_net.predict_reward_train(
+        state=torch.tensor(obs).unsqueeze(0).to(discriminator.discrim_net.device()),
+        action=torch.tensor([[action]]).to(discriminator.discrim_net.device()),
+        next_state=torch.tensor(next_obs).unsqueeze(0).to(discriminator.discrim_net.device()), # unused
+        done=torch.tensor(done).unsqueeze(0).to(discriminator.discrim_net.device()), # unused
+    ).item()
+
+class OptionsEnv(gym.Wrapper):
+
+    def __init__(self, env, options, discriminator, discount, ll_buffer, *args, **kwargs):
+        super().__init__(env, *args, **kwargs)
+
+        self.options = options
+        num_hl_options = len(self.options)
+        self.action_space = gym.spaces.Discrete(num_hl_options)
+        self.observation_space = gym.spaces.Dict({
+            'obs': env.observation_space,
+            'mask': gym.spaces.Box(low=0, high=1, shape=(num_hl_options,)),
+        })
+
+        self.discriminator = discriminator
+        self.discount = discount
+        self.ll_buffer = ll_buffer
+    
+    @staticmethod
+    def _hl_observation(obs, mask):
+        return {
+            'obs': obs,
+            'mask': mask,
+        }
+    
+    def reset(self):
+        self.done = False
+        self.obs = self.env.reset()
+        self.m = available_actions(self.env, self.options)
+        return self._hl_observation(self.obs, self.m)
+    
+    def _ll_step(self, action):
+        return self.env.step(action)
+    
+    def step(self, action):
+        assert self.m[action]
+        assert not self.done
+
+        plan = list(map(float, generate_plan(self.env, action, self.options)))
+        reward = 0
+        steps = 0
+
+        while not self.done and plan and \
+            (feasible(self.env, safety_plan(self.env, plan)) or self.m.sum() == 1):
+            
+            a, plan = plan[0], plan[1:]
+            a = self.env._normalize(a)
+
+            next_obs, _, self.done, info = self._ll_step(a)
+
+            reward += self.discount**steps * self.discriminator(self.obs, a, next_obs, self.done)
+
+            self.ll_buffer.append({
+                'obs': self.obs,
+                'next_obs': next_obs,
+                'acts': np.array((a,)),
+                'dones': np.array(self.done),
+            })
+
+            steps += 1
+            self.obs = next_obs
+
+        self.m = available_actions(self.env, self.options)
+
+        return self._hl_observation(self.obs, self.m), reward, self.done, info
+
+class RenderOptions(OptionsEnv):
+
+    def __init__(self, env, options, *args, **kwargs):
+        super().__init__(env, options, discriminator=lambda s, a, n, d: 0, discount=1, ll_buffer=deque(maxlen=0), *args, **kwargs)
+
+    def _ll_step(self, action):
+        out = super()._ll_step(action)
+        self.env.render(mode='post')
+        return out
+    
+    def close(self, *args, **kwargs):
+        self.env.close(*args, **kwargs)
+
+def safety_plan(env, plan):
+    return np.concatenate((plan, np.array(5 * [env._env._min_acc])), axis=0)
+
+def available_actions(env, options):
+    """Return mask of available actions given current `env` state.
+    Action 0 is considered safe fallback.
+    """
+    plans = [generate_plan(env, i, options) for i, _ in enumerate(options)]
+    # is emergency braking still possible?
+    plans = list(map(lambda p: safety_plan(env, p), plans))
+
+    T = max(len(p) for p in plans)
+    plans = [np.pad(p, ((0, T-len(p)),), constant_values=np.nan) for p in plans]
+    plans = np.stack(plans, axis=0)
+    
+    valid = feasible(env, plans)
+    if not valid.any():
+        valid[0] = True
+
+    return valid
+
+def target_velocity_plan(current_v: float, target_v: float, t: int, dt: float):
+    """Smoothly target a velocity in a given number of steps"""
+    # for now, constant acceleration
+    a = (target_v - current_v)  / (t * dt)
+    return a*np.ones((t,))
+
+def generate_plan(env, i, options):
+    """Generate input profile for high-level action `i`."""
+    assert i < len(options), "Invalid option index {i}"
+    target_v, t = options[i]
+    current_v = env._env.state[env._agent, 1].item() # extract from env
+    plan = target_velocity_plan(current_v, target_v, t, env._env._dt)
+    assert len(plan) == t, "incorrect plan length"
+    return plan

scratch/etienne/intersimple/gail_flat.py

@@ -43,12 +43,13 @@ logger.configure(tempdir_path / "GAIL/")
 gail_trainer = adversarial.GAIL(
     venv,
     expert_data=transitions,
-    expert_batch_size=220,
+    expert_batch_size=150,
-    #n_disc_updates_per_round=32,
+    n_disc_updates_per_round=32,
     discrim_kwargs={'discrim_net': MlpDiscriminator()},
-    gen_algo=sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=4096),
+    gen_algo=sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=4530),
+    allow_variable_horizon=True,
 )
-gail_trainer.train(total_timesteps=80000)
+gail_trainer.train(total_timesteps=400000)
 gail_trainer.gen_algo.save(model_name)
 
 #del gail_trainer
@@ -66,4 +67,4 @@ while True:
     if done:
         break
 
-env.close(filestr='render/'+model_name)
+env.close(filestr='render/'+model_name)

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)

scratch/etienne/intersimple/gail_image.py

@@ -18,7 +18,7 @@ model_name = 'gail_image'
 
 # %%
 # Load pickled test demonstrations.
-with open("data/NormalizedIntersimpleExpert_NRasterizedAgent51.pkl", "rb") as f:
+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.
@@ -30,7 +30,7 @@ with open("data/NormalizedIntersimpleExpert_NRasterizedAgent51.pkl", "rb") as f:
 # (observation, actions, next_observation) transitions.
 transitions = rollout.flatten_trajectories(trajectories)
 
-venv = make_vec_env(NRasterized, n_envs=2, env_kwargs={'agent': 51})
+venv = make_vec_env(NRasterized, n_envs=2, env_kwargs={'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2})
 
 tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
 tempdir_path = pathlib.Path(tempdir.name)
@@ -43,10 +43,11 @@ logger.configure(tempdir_path / "GAIL/")
 gail_trainer = adversarial.GAIL(
     venv,
     expert_data=transitions,
-    expert_batch_size=200,
+    expert_batch_size=32,
-    n_disc_updates_per_round=2048,
+    #n_disc_updates_per_round=2048,
     discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
-    gen_algo=sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=128),
+    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)
@@ -56,7 +57,7 @@ gail_trainer.gen_algo.save(model_name)
 # %%
 model = sb3.PPO.load(model_name)
 
-env = NRasterized(agent=51)
+env = NRasterized(agent=51, width=36, height=36, m_per_px=2)
 
 obs = env.reset()
 while True:

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_random.py

@@ -0,0 +1,79 @@
+# %%
+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 NRasterizedRandomAgent, IntersimpleReward, speed_reward
+import functools
+from stable_baselines3.common.evaluation import evaluate_policy
+
+from gail.discriminator import CnnDiscriminator
+
+model_name = 'gail_image_random'
+
+# %%
+# 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)
+
+env_kwargs = {'width': 36, 'height': 36, 'm_per_px': 2}
+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}/.")
+
+# 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/")
+generator = sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=1024)
+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=generator,
+    allow_variable_horizon=True,
+)
+def callback(round):
+    eval_env = NRasterizedRandomAgent(reward=functools.partial(speed_reward, collision_penalty=0.), **env_kwargs)
+    #sync_envs_normalization(self.training_env, self.eval_env)
+    episode_rewards, episode_lengths = evaluate_policy(generator, eval_env, return_episode_rewards=True)
+
+gail_trainer.train(total_timesteps=100000, callback=callback)
+gail_trainer.gen_algo.save(model_name)
+
+#del gail_trainer
+
+# %%
+model = sb3.PPO.load(model_name)
+
+env = NRasterizedRandomAgent(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_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)
+# %%
+

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/gail_options.py

@@ -1,116 +0,0 @@
-from gail.discriminator import MlpDiscriminator
-from imitation.algorithms import adversarial
-import stable_baselines3
-import torch.utils.data
-import numpy as np
-from intersim.envs.intersimple import Intersimple
-import itertools
-from torch.distributions import Categorical
-import gym
-
-class OptionsMlpPolicy:
-
-    def __init__(self, *args, **kwargs):
-        self._policy = stable_baselines3.common.policies.ActorCriticPolicy(
-            *args, **kwargs
-        )
-
-    def _prior_distribution(self, s):
-        latent_pi, _, latent_sde = self._policy._get_latent(s)
-        distribution = self._policy._get_action_dist_from_latent(latent_pi, latent_sde)
-        return distribution.distribution
-
-    def predict(self, obs):
-        s, m = obs
-        prior = self._prior_distribution(s)
-        posterior = Categorical(prior.probs * m)
-        ch = posterior.sample()
-        return ch
-
-    def evaluate_actions(self, obs, ch):
-        s, m = obs
-        values = self._policy.value_net(s)
-        prior = self._prior_distribution(s)
-        posterior = Categorical(prior.probs * m)
-        return values, posterior.logprob(ch), posterior.entropy() # additional values used by PPO.train
-
-def available_actions(env):
-    """Return mask of available actions given current `env` state."""
-    return np.ones((env.num_hl_actions,))
-
-def generate_plan(env, i):
-    """Generate input profile for high-level action `i`."""
-    return np.zeros((env.num_hl_steps,))
-
-def feasible(env, plan):
-    """Check if input profile is feasible given current `env` state."""
-    return True
-
-def sample_ll(env, generator):
-    """Sample low-level (state, action) pairs for discriminator training."""
-    done = True
-    while True:
-        if done:
-            s = env.reset()
-
-        m = available_actions(env)
-        ch = generator.policy.predict((s, m))
-        plan = list(generate_plan(env, ch))
-
-        while not done and plan and feasible(env, plan):
-            a = plan.pop()
-            yield (s, a)
-            s, _, done, _ = env.step(a)
-
-def train_discriminator(env, expert_data, generator, discriminator, generator_batch_size):
-    expert_samples = next(expert_data)
-    generator_samples = itertools.islice(sample_ll(env, generator), generator_batch_size)
-    discriminator.train_disc(expert_samples, generator_samples)
-
-def sample_hl(env, generator, discriminator):
-    """Sample high-level (state, action, reward) tuples for generator training."""
-    done = True
-    while True:
-        if done:
-            s = env.reset()
-
-        m = available_actions(env)
-        obs = (s, m)
-        ch = generator.policy.predict((s, m))
-        plan = list(generate_plan(env, ch))
-        r = 0
-        discount = 1
-
-        while not done and plan and feasible(env, plan):
-            a = plan.pop()
-            r += discount * discriminator.discrim_net(s, a)
-            discount *= env.discount
-            s, _, done, _ = env.step(a)
-        
-        yield (obs, ch, r)
-
-def train_generator(env, generator, discriminator, generator_batch_size):
-    generator_samples = itertools.islice(sample_hl(env, generator, discriminator), generator_batch_size)
-    generator.rollout_buffer.reset()
-    generator.rollout_buffer.add(generator_samples)
-    generator.train()
-
-class OptionsEnv(gym.Wrapper):
-    
-    def __init__(self, env):
-        super().__init__(env)
-        self.action_space = gym.spaces.Discrete(env.num_hl_options)
-
-def train(expert_data, epochs=10, generator_batch_size=1024, expert_batch_size=1024, num_hl_options=10, num_hl_steps=10, discount=0.99):
-    env = Intersimple()
-    env.num_hl_options = num_hl_options
-    env.num_hl_steps = num_hl_steps
-    env.discount = discount
-
-    discriminator = adversarial.GAIL(discrim_kwargs={'discrim_net': MlpDiscriminator()})
-    generator = stable_baselines3.PPO(OptionsMlpPolicy, OptionsEnv(env))
-    expert_data = torch.utils.data.DataLoader(expert_data, expert_batch_size)
-
-    for _ in range(epochs):
-        train_discriminator(env, expert_data, generator, discriminator, generator_batch_size)
-        train_generator(env, generator, discriminator, generator_batch_size)

scratch/etienne/intersimple/gail_options_image.py

@@ -0,0 +1,90 @@
+# %%
+import sys
+sys.path.append('../../../')
+
+from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
+from imitation.algorithms import adversarial
+import stable_baselines3
+import torch.utils.data
+import numpy as np
+from intersim.envs.intersimple import NRasterized
+import itertools
+from torch.distributions import Categorical
+import gym
+import torch
+import pickle
+import imitation.data.rollout as rollout
+import tempfile
+import pathlib
+from imitation.util import logger
+from stable_baselines3.common.env_util import make_vec_env
+from tqdm import tqdm
+from src.policies.options import OptionsCnnPolicy
+from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
+from src.gail.train import train_discriminator, train_generator
+
+model_name = 'gail_options_image'
+env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2}
+
+ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10]] # option 0 is safe fallback
+
+def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, discount=0.99):
+    env = NRasterized(**env_settings)
+    env.discount = discount
+
+    tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
+    tempdir_path = pathlib.Path(tempdir.name)
+    logger.configure(tempdir_path / "GAIL/")
+    print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
+
+    venv = make_vec_env(NRasterized, n_envs=1, env_kwargs=env_settings)
+    discriminator = adversarial.GAIL(
+        expert_data=expert_data,
+        expert_batch_size=expert_batch_size,
+        discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
+        #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
+        venv=venv, # unused
+        gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
+    )
+
+    generator = stable_baselines3.PPO(
+        OptionsCnnPolicy,
+        OptionsEnv(env, options=ALL_OPTIONS),
+        verbose=1,
+        n_steps=generator_steps,
+    )
+
+    # PPO.train requires logger as set up in
+    # PPO._setup_learn (called by PPO.learn)
+    generator._logger = stable_baselines3.common.utils.configure_logger(
+        generator.verbose,
+        generator.tensorboard_log,
+    )
+
+    for _ in tqdm(range(epochs)):
+        train_discriminator(LLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=expert_batch_size)
+        train_generator(HLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=generator_steps)
+    
+    return generator
+
+# %%
+if __name__ == '__main__':
+    # %%
+ 
+    with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
+        trajectories = pickle.load(f)
+    transitions = rollout.flatten_trajectories(trajectories)
+    generator = train(transitions)
+
+    generator.save(model_name)
+
+    # %%
+    model = stable_baselines3.PPO.load(model_name)
+
+    env = RenderOptions(NRasterized(**env_settings), options=ALL_OPTIONS)
+
+    for s in env.sample_ll(model):
+        if s['dones']:
+            break
+
+    env.close(filestr='render/'+model_name)

scratch/etienne/intersimple/gail_options_image_alltracks.py

@@ -0,0 +1,410 @@
+# %%
+import sys
+sys.path.append('../../../')
+from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
+from src.policies import OptionsCnnPolicy
+from src.util import feasible
+from src.data import load_experts
+
+from imitation.algorithms import adversarial
+from imitation.util import logger
+import imitation.data.rollout as rollout
+
+import stable_baselines3
+from stable_baselines3.common.env_util import make_vec_env
+
+import torch
+import torch.utils.data
+import numpy as np
+import itertools
+import gym
+import pickle
+import tempfile
+import pathlib
+from tqdm import tqdm
+
+from intersim.envs.intersimple import NRasterized, NRasterizedRandomAgent, NRasterizedIncrementingAgent
+
+ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback
+
+class OptionsEnv(gym.Wrapper):
+    """
+    Wrap an intersimple environment with an options generator
+    """
+    def __init__(self, env, *args, **kwargs):
+        """
+        Initialize wrapped environment and set high-level action and observation spaces
+        """
+        super().__init__(env, *args, **kwargs)
+        num_hl_options = len(ALL_OPTIONS)
+        self.action_space = gym.spaces.Discrete(num_hl_options)
+        self.observation_space = gym.spaces.Dict({
+            'obs': env.observation_space,
+            'mask': gym.spaces.Box(low=0, high=1, shape=(num_hl_options,)),
+        })
+
+    def _after_choice(self):
+        pass
+    
+    def _after_step(self):
+        pass
+
+    def _transitions(self):
+        raise NotImplementedError('Use `LLOptions` or `HLOptions` for sampling.')
+    
+    def sample(self, generator):
+        """
+        yield transitions using a generator
+        Args:
+            generator (sb3.PPO)
+        Yields:
+
+        """
+        self.done = True
+        while True:
+            self.episode_start = False
+            if self.done:
+                self.s = self.env.reset()
+                self.done = False
+                self.episode_start = True
+
+            self.m = available_actions(self.env)
+            self.ch, self.value, self.log_prob = generator.policy.predict({
+                'obs': torch.tensor(self.s).unsqueeze(0).to(generator.policy.device),
+                'mask': torch.tensor(self.m).unsqueeze(0).to(generator.policy.device),
+            })
+            self.plan = list(map(float, generate_plan(self.env, self.ch)))
+
+            self._after_choice()
+
+            assert not self.done
+            assert self.plan
+            #assert feasible(self.env, self.plan, self.ch)
+
+            while not self.done and self.plan and feasible(self.env, self.plan, self.ch):
+                self.a, self.plan = self.plan[0], self.plan[1:]
+                self.a = self.env._normalize(self.a)
+                self.nexts, _, self.done, _ = self.env.step(self.a)
+
+                self._after_step()
+
+                self.s = self.nexts
+            
+            yield from self._transitions()
+
+class LLOptions(OptionsEnv):
+    """Sample low-level (state, action) tuples for discriminator training."""
+    
+    def __init__(self, *args, **kwargs):
+        """
+        LLOption uses the true LL observations
+        """
+        super().__init__(*args, **kwargs)
+        # overwrite observation space to just output obs directly
+        self.observation_space = self.observation_space['obs']
+
+    def _after_choice(self):
+        """
+        After each option choice, initialize/reset the transition buffer
+        """
+        self._transition_buffer = []
+
+    def _after_step(self):
+        """
+        After each ll action, append s, s', a, done to transition buffer
+        """
+        self._transition_buffer.append({
+            'obs': self.s,
+            'next_obs': self.nexts,
+            'acts': np.array((self.a,)),
+            'dones': np.array(self.done),
+        })
+
+    def _transitions(self):
+        """
+        Yield from the transition buffer
+        """
+        yield from self._transition_buffer
+    
+    def sample_ll(self, policy):
+        """
+        Args:
+            policy
+        Returns:
+            gen: iterable which samples low-level transitions from the environment
+        """
+        return self.sample(policy)
+
+class HLOptions(OptionsEnv):
+    """Sample high-level (state, action, reward) tuples for generator training."""
+    
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _after_choice(self):
+        """
+        After an option selection, initialize total reward and number of steps
+        """
+        self.obs = {'obs': np.copy(self.s), 'mask': np.copy(self.m)}
+        self.r = 0
+        self.steps = 0
+
+    def _after_step(self):
+        """
+        After each low-level action, add the discounted discriminated reward score (given a discriminator)
+        """
+        self.r += self.discount**self.steps * self.discriminator.discrim_net.reward_train(
+            state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()),
+            action=torch.tensor([[self.a]]).to(self.discriminator.discrim_net.device()),
+            next_state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused
+            done=torch.tensor(self.done).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused
+        )
+        self.steps += 1
+    
+    def _transitions(self):
+        """
+        Yield a single dictionary per high-level selected action
+        Fields:
+            obs: high-level state and mask at selection
+            action: chosen high-level action
+            reward: accumulated option reward
+            episode_start: whether the action was chosen at the episode start
+            value: the value estimate from the starting state
+            log_prob: the log_prob of the selected action from the starting state
+            done: whether the episode has ended
+
+        """
+        yield {
+            'obs': self.obs,
+            'action': self.ch,
+            'reward': self.r.detach(),
+            'episode_start': self.episode_start,
+            'value': self.value.detach(),
+            'log_prob': self.log_prob.detach(),
+            'done': self.done,
+        }
+    
+    def sample_hl(self, policy, discriminator):
+        """
+        Args:
+            policy
+            discriminator: function with which to score rewards
+        Returns:
+            gen: iterable which samples high-level transitions from the environment
+        """
+        self.discriminator = discriminator
+        return self.sample(policy)
+
+class RenderOptions(LLOptions):
+
+    def _after_step(self):
+        """
+        Render the environment after each low-level step
+        """
+        super()._after_step()
+        self.env.render()
+    
+    def close(self, *args, **kwargs):
+        """
+        On 'close', close the environment
+        """
+        self.env.close(*args, **kwargs)
+
+def available_actions(env):
+    """Return mask of available actions given current `env` state."""
+    valid = np.array([feasible(env, generate_plan(env, i), i) for i in range(len(ALL_OPTIONS))])
+    return valid
+
+def target_velocity_plan(current_v: float, target_v: float, t: int, dt: float):
+    """Smoothly target a velocity in a given number of steps"""
+    # for now, constant acceleration
+    a = (target_v - current_v)  / (t * dt)
+    return a*np.ones((t,))
+
+def generate_plan(env, i):
+    """Generate input profile for high-level action `i`.
+
+    Args:
+        env (gym.Env): current environment state
+        i (int): high-level action `i`
+    Returns:
+        plan (np.array): length T array of acceleration values
+    """
+    assert i < len(ALL_OPTIONS), "Invalid option index {i}"
+    target_v, t = ALL_OPTIONS[i]
+    current_v = env._env.state[env._agent, 1].item() # extract from env
+    plan = target_velocity_plan(current_v, target_v, t, env._env._dt)
+    assert len(plan) == t, "incorrect plan length"
+    return plan
+
+def flatten_transitions(transitions):
+    return {
+        'obs': np.stack(list(t['obs'] for t in transitions), axis=0),
+        'next_obs': np.stack(list(t['next_obs'] for t in transitions), axis=0),
+        'acts': np.stack(list(t['acts'] for t in transitions), axis=0),
+        'dones': np.stack(list(t['dones'] for t in transitions), axis=0),
+    }
+
+def train_discriminator(env, generator, discriminator, num_samples):
+    transitions = list(itertools.islice(env.sample_ll(generator), num_samples))
+    generator_samples = flatten_transitions(transitions)
+    discriminator.train_disc(gen_samples=generator_samples)
+
+def train_generator(env, generator, discriminator, num_samples):
+    generator_samples = list(itertools.islice(env.sample_hl(generator, discriminator), num_samples+1))
+    
+    generator.rollout_buffer.reset()
+    for s in generator_samples[:-1]:
+        generator.rollout_buffer.add(
+            obs=s['obs'],
+            action=s['action'].cpu(),
+            reward=s['reward'].cpu(),
+            episode_start=s['episode_start'],
+            value=s['value'],
+            log_prob=s['log_prob'],
+        )
+    
+    generator.rollout_buffer.compute_returns_and_advantage(
+        last_values=generator_samples[-1]['value'],
+        dones=generator_samples[-1]['done'],
+    )
+    
+    generator.train()
+
+def train(expert_data, env_class=NRasterizedRandomAgent, env_settings={}, epochs=10, discrim_batch_size=32, generator_steps=2048, discount=0.99):
+    """
+    Args: 
+        expert_data: list of transitions
+        env_class: environment class
+        env_settings: environment settings
+        epochs: number of epochs to train for
+        discrim_batch_size: discriminator batch size
+        generator_steps: number of steps taken in generator
+        discount: discount factor
+    Returns:
+        generator (stable_baselines3.PPO): options policy
+    """
+    env = env_class(**env_settings)
+    env.discount = discount
+
+    tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
+    tempdir_path = pathlib.Path(tempdir.name)
+    logger.configure(tempdir_path / "GAIL/")
+    print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
+
+    venv = make_vec_env(env_class, n_envs=1, env_kwargs=env_settings)
+    discriminator = adversarial.GAIL(
+        expert_data=expert_data,
+        expert_batch_size=discrim_batch_size,
+        discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
+        #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
+        venv=venv, # unused
+        gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
+    )
+
+    generator = stable_baselines3.PPO(
+        OptionsCnnPolicy,
+        OptionsEnv(env),
+        verbose=1,
+        n_steps=generator_steps,
+    )
+
+    # PPO.train requires logger as set up in
+    # PPO._setup_learn (called by PPO.learn)
+    generator._logger = stable_baselines3.common.utils.configure_logger(
+        generator.verbose,
+        generator.tensorboard_log,
+    )
+
+    for _ in tqdm(range(epochs)):
+        train_discriminator(LLOptions(env), generator, discriminator, num_samples=discrim_batch_size)
+        train_generator(HLOptions(env), generator, discriminator, num_samples=generator_steps)
+    
+    return generator
+
+# %%
+if __name__ == '__main__':
+    # %%
+    model_name = 'gail_options_image'
+    env_class = NRasterizedRandomAgent
+    env_settings = {'width': 36, 'height': 36, 'm_per_px': 2}
+    
+    #env_class = NRasterized
+    #env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2}
+    files = ['../../../expert_data/DR_USA_Roundabout_FT/track%04i/expert.pkl'%(i) for i in range(5)]
+    transitions=load_experts(files)
+
+    generator = train(
+        transitions,
+        env_class=env_class,
+        env_settings=env_settings,
+        epochs=10, 
+        discrim_batch_size=32, 
+        generator_steps=2048, 
+        discount=0.99
+        )
+
+    generator.save(model_name)
+
+    # %%
+    model = stable_baselines3.PPO.load(model_name)
+
+    env = RenderOptions(NRasterizedRandomAgent(**env_args))
+
+    for s in env.sample_ll(model):
+        if s['dones']:
+            break
+
+    env.close(filestr='render/'+model_name)
+
+# %% Tests
+
+def test_ll_expert_data():
+    with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
+        expert_trajectories = pickle.load(f)
+    expert_transitions = rollout.flatten_trajectories(expert_trajectories)
+
+    env = LLOptions(NRasterized(agent=51, width=36, height=36, m_per_px=2))
+
+    gen_transitions = list(itertools.islice(env.sample_ll(
+        policy=stable_baselines3.PPO(
+            OptionsCnnPolicy,
+            OptionsEnv(env),
+            verbose=1,
+        )
+    ), 10))
+    gen_transitions = flatten_transitions(gen_transitions)
+
+    assert expert_transitions[:10].obs.shape == gen_transitions['obs'].shape
+    assert expert_transitions[:10].next_obs.shape == gen_transitions['next_obs'].shape
+    assert expert_transitions[:10].acts.shape == gen_transitions['acts'].shape
+    assert expert_transitions[:10].dones.shape == gen_transitions['dones'].shape
+
+def test_ll_states():
+    env = NRasterized()
+    policy = stable_baselines3.PPO(
+        OptionsCnnPolicy,
+        OptionsEnv(env),
+        verbose=1,
+    )
+    llenv = LLOptions(env)
+    transitions = list(itertools.islice(llenv.sample_ll(policy=policy), 100))
+
+    env2 = NRasterized()
+    s2 = env2.reset()
+    for i, t in enumerate(transitions):
+        assert i == 0 or np.array_equal(t['obs'], transitions[i-1]['next_obs'])
+        assert np.array_equal(t['obs'], s2)
+        assert t['acts'].shape == (1,)
+
+        nexts2, _, done2, _ = env2.step(t['acts'])
+        assert np.array_equal(t['next_obs'], nexts2)
+        assert np.array_equal(t['dones'], done2)
+
+        if done2:
+            break
+
+        s2 = nexts2
+
+def test_hl_transitions():
+    pass

scratch/etienne/intersimple/gail_options_image_random.py

@@ -0,0 +1,91 @@
+# %%
+import sys
+sys.path.append('../../../')
+
+from src.discriminator import CnnDiscriminatorFlatAction
+from imitation.algorithms import adversarial
+import stable_baselines3
+import torch.utils.data
+import numpy as np
+from intersim.envs.intersimple import NRasterizedRouteRandomAgent
+import itertools
+from torch.distributions import Categorical
+import gym
+import torch
+import pickle
+import imitation.data.rollout as rollout
+import tempfile
+import pathlib
+from imitation.util import logger
+from stable_baselines3.common.env_util import make_vec_env
+from tqdm import tqdm
+from src.policies.options import OptionsCnnPolicy
+from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
+from src.gail.train import train_discriminator, train_generator
+
+model_name = 'gail_options_image_random'
+env_settings = {'width': 70, 'height': 70, 'm_per_px': 1}
+
+ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback
+
+def train(expert_data, epochs=100, expert_batch_size=64, generator_steps=1024, discount=0.99):
+    env = NRasterizedRouteRandomAgent(**env_settings)
+    env.discount = discount
+
+    tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
+    tempdir_path = pathlib.Path(tempdir.name)
+    logger.configure(tempdir_path / "GAIL/")
+    print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
+
+    venv = make_vec_env(NRasterizedRouteRandomAgent, n_envs=1, env_kwargs=env_settings)
+    discriminator = adversarial.GAIL(
+        expert_data=expert_data,
+        expert_batch_size=expert_batch_size,
+        discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
+        #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
+        venv=venv, # unused
+        gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
+    )
+
+    generator = stable_baselines3.PPO(
+        OptionsCnnPolicy,
+        OptionsEnv(env, options=ALL_OPTIONS),
+        verbose=1,
+        n_steps=generator_steps,
+    )
+
+    # PPO.train requires logger as set up in
+    # PPO._setup_learn (called by PPO.learn)
+    generator._logger = stable_baselines3.common.utils.configure_logger(
+        generator.verbose,
+        generator.tensorboard_log,
+    )
+
+    for _ in tqdm(range(epochs)):
+        train_discriminator(LLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=expert_batch_size)
+        train_generator(HLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=generator_steps)
+        generator.save(model_name)
+    
+    return generator
+
+def video(model_name, env):
+    model = stable_baselines3.PPO.load(model_name)
+    env = RenderOptions(env, options=ALL_OPTIONS)
+    for s in env.sample_ll(model):
+        if s['dones']:
+            break
+    env.close(filestr='render/'+model_name)
+
+def evaluate():
+    video(
+        model_name=model_name,
+        env=NRasterizedRouteRandomAgent(**env_settings)
+    )
+
+# %%
+if __name__ == '__main__':
+
+    with open("data/NormalizedIntersimpleExpertMu.001N10000_NRasterizedRouteRandomAgentw70h70mppx1.pkl", "rb") as f:
+        trajectories = pickle.load(f)
+    transitions = rollout.flatten_trajectories(trajectories)
+    train(transitions)

scratch/etienne/intersimple/gail_options_image_random_location.py

@@ -0,0 +1,144 @@
+# %%
+from collections import deque
+import sys
+sys.path.append('../../../')
+
+from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
+from imitation.algorithms import adversarial
+import stable_baselines3
+import pickle
+import imitation.data.rollout as rollout
+import tempfile
+import pathlib
+from imitation.util import logger
+from tqdm import tqdm
+from src.policies.options import OptionsCnnPolicy
+from src.gail.train import flatten_transitions
+from gail.options2 import OptionsEnv, RenderOptions, imitation_discriminator
+from gail.envs import TLNRasterizedRouteRandomAgentLocation
+from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
+from stable_baselines3.common.env_util import make_vec_env
+import torch
+import numpy as np
+
+model_name = 'gail_options_image_random_location'
+env_settings = {'width': 70, 'height': 70, 'm_per_px': 1, 'mu': 0.001, 'random_skip': True, 'max_episode_steps': 200}
+
+ALL_OPTIONS = [(v,t) for v in [0,2,4,8,10] for t in [5, 10, 20]] # option 0 is safe fallback
+
+class NoisyDiscriminator(CnnDiscriminatorFlatAction):
+
+    def __init__(self, *args, std=0.0, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.std = std
+
+    def forward(self, state, action):
+        noise = self.std * torch.randn(*action.shape, device=action.device)
+        return super().forward(state, action + noise)
+
+class LLBuffer(deque):
+
+    def sample(self, n):
+        assert n <= self.maxlen, f'Sample size of {n} exceeds buffer capacity of {self.maxlen}'
+        assert n <= len(self), f'Sample size of {n} exceeds buffer size of {len(self)}'
+        ind = np.random.randint(len(self), size=n)
+        return list(self[i] for i in ind)
+
+def train(
+        expert_data,
+        expert_batch_size=4096,
+        discriminator_updates_per_round=20,
+        generator_steps=1024,
+        generator_batch_size=1024,
+        generator_total_steps=8192,
+        generator_updates_per_round=10,
+        discount=1.0,
+        epochs=200,
+    ):
+    env = TLNRasterizedRouteRandomAgentLocation(**env_settings)
+
+    tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
+    tempdir_path = pathlib.Path(tempdir.name)
+    logger.configure(tempdir_path / "GAIL/")
+    print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
+
+    venv = DummyVecEnv([lambda: env])
+    discriminator = adversarial.GAIL(
+        expert_data=expert_data,
+        expert_batch_size=expert_batch_size,
+        #discrim_kwargs={'discrim_net': NoisyDiscriminator(venv, std=0.25)},
+        disc_opt_cls=torch.optim.RMSprop,
+        disc_opt_kwargs={'lr': 0.0001, 'weight_decay': 0.003},
+        discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
+        venv=venv, # unused
+        gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
+    )
+
+    ll_buffer = LLBuffer(maxlen=expert_batch_size*10)
+
+    options_env = make_vec_env(
+        OptionsEnv,
+        n_envs=1,
+        #vec_env_cls=SubprocVecEnv,
+        env_kwargs={
+            'env': env,
+            'options': ALL_OPTIONS,
+            'discriminator': imitation_discriminator(discriminator),
+            'discount': discount,
+            'll_buffer': ll_buffer,
+        }
+    )
+
+    generator = stable_baselines3.PPO(
+        OptionsCnnPolicy,
+        options_env,
+        verbose=1,
+        batch_size=generator_batch_size,
+        n_steps=generator_steps,
+        n_epochs=generator_updates_per_round,
+        gamma=1.0,
+        learning_rate=1e-4,
+    )
+
+    for _ in tqdm(range(epochs)):
+        ll_buffer.clear()
+
+        # train generator
+        generator.learn(total_timesteps=generator_total_steps)
+
+        # train discriminator
+        for _ in range(discriminator_updates_per_round):
+            generator_samples = ll_buffer.sample(expert_batch_size)
+            generator_samples = flatten_transitions(generator_samples)
+            discriminator.train_disc(gen_samples=generator_samples)
+
+        generator.save(model_name)
+    
+    return generator
+
+def video(model_name, env):
+    model = stable_baselines3.PPO.load(model_name)
+
+    done = False
+    obs = env.reset()
+    while not done:
+        action, _ = model.predict(obs)
+        obs, _, done, _ = env.step(action)
+
+    env.close(filestr='render/'+model_name)
+
+def evaluate():
+    video_settings = { **env_settings, 'random_skip': False, 'max_episode_steps': 200 }
+    env = TLNRasterizedRouteRandomAgentLocation(**video_settings)
+    env = RenderOptions(env, options=ALL_OPTIONS)
+    video(
+        model_name=model_name,
+        env=env
+    )
+
+# %%
+if __name__ == '__main__':
+    with open("data/NormalizedIntersimpleExpertMu.001N50000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl", "rb") as f:
+        trajectories = pickle.load(f)
+    transitions = rollout.flatten_trajectories(trajectories)
+    train(transitions)

scratch/etienne/intersimple/ppo_speed_image_lowres.py

@@ -0,0 +1,49 @@
+# %%
+from stable_baselines3 import PPO
+from intersim.envs.intersimple import NRasterized, speed_reward
+import functools
+
+model_name = "ppo_speed_image_lowres"
+
+#def reward(state, action, info):
+#    speed = state[2].item()
+#    r = speed if speed < 10 else (10 - 5 * (speed - 10))
+#    return 0.1 * r
+
+env = NRasterized(
+    agent=51,
+    height=36,
+    width=36,
+    m_per_px=2,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+)
+
+# %%
+model = PPO(
+    "CnnPolicy", env,
+    verbose=1,
+)
+model.learn(total_timesteps=100000)
+model.save(model_name)
+
+print('Done training.')
+
+del model # remove to demonstrate saving and loading
+
+# %%
+model = PPO.load(model_name)
+
+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/ppo_speed_image_lowres_random.py

@@ -0,0 +1,42 @@
+# %%
+from stable_baselines3 import PPO
+from intersim.envs.intersimple import NRasterizedRandomAgent, speed_reward
+import functools
+
+model_name = "ppo_speed_image_lowres_random"
+
+env = NRasterizedRandomAgent(
+    height=36,
+    width=36,
+    m_per_px=2,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    )
+)
+
+# %%
+model = PPO(
+    "CnnPolicy", env,
+    verbose=1,
+    batch_size=2048,
+)
+model.learn(total_timesteps=2e5)
+model.save(model_name)
+
+print('Done training.')
+
+del model # remove to demonstrate saving and loading
+
+# %%
+model = PPO.load(model_name)
+
+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,33 @@
+
+import stable_baselines3 as sb3
+from intersim.envs.intersimple import NRasterized
+
+
+def render_env(model_name='gail_image_multiagent_nocollision', agent=51, environment=NRasterized):
+    """
+    Render a video from an model, agent, and environment
+    Args:
+        model_name (str): name of the model
+        agent (int): agent to start the video from
+        environment (gym.Env): gym environment class to render environment on
+    """
+
+    model = sb3.PPO.load(model_name)
+
+    env = environment(stop_on_collision=False, width=36, height=36, m_per_px=2, agent=agent)
+
+    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+'_agent%i'%(agent)) 
+
+if __name__ == '__main__':
+    import fire
+    fire.Fire(render_env)

scratch/etienne/intersimple/train_discrim.py

@@ -0,0 +1,71 @@
+# %%
+import sys
+sys.path.append('../../../')
+
+import pickle
+import imitation.data.rollout as rollout
+import imitation.data.types as types
+import torch
+from gail.envs import TLNRasterizedRouteRandomAgentLocation
+import tempfile
+import pathlib
+from imitation.util import logger
+from stable_baselines3.common.vec_env.dummy_vec_env import DummyVecEnv
+from imitation.algorithms import adversarial
+from src.discriminator import CnnDiscriminator
+import stable_baselines3
+from tqdm import tqdm
+
+with open("data/NormalizedIntersimpleExpertMu.001N50000_TLNRasterizedRouteRandomAgentLocationw70h70mppx1mu.001rskips50.pkl", "rb") as f:
+    trajectories = pickle.load(f)
+transitions = rollout.flatten_trajectories(trajectories)
+
+# %%
+env_settings = {'width': 70, 'height': 70, 'm_per_px': 1, 'mu': 0.001, 'random_skip': True, 'max_episode_steps': 200}
+env = TLNRasterizedRouteRandomAgentLocation(**env_settings)
+
+tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
+tempdir_path = pathlib.Path(tempdir.name)
+logger.configure(tempdir_path / "GAIL/")
+print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
+
+expert_batch_size = 4096
+
+venv = DummyVecEnv([lambda: env])
+discriminator = adversarial.GAIL(
+    expert_data=transitions,
+    expert_batch_size=expert_batch_size,
+    #discrim_kwargs={'discrim_net': NoisyDiscriminator(venv, std=0.25)},
+    disc_opt_cls=torch.optim.RMSprop,
+    disc_opt_kwargs={'lr': 0.0001, 'weight_decay': 0.003},
+    discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
+    venv=venv, # unused
+    gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
+)
+
+expert_data_loader = torch.utils.data.DataLoader(
+    transitions,
+    batch_size=expert_batch_size,
+    collate_fn=types.transitions_collate_fn,
+    shuffle=True,
+    drop_last=True,
+)
+
+gen_data_loader = torch.utils.data.DataLoader(
+    transitions,
+    batch_size=expert_batch_size,
+    collate_fn=types.transitions_collate_fn,
+    shuffle=True,
+    drop_last=True,
+)
+
+# %%
+epochs = 1000
+for i in tqdm(range(epochs)):
+    for expert_samples, gen_samples in zip(expert_data_loader, gen_data_loader):
+        # randomly corrupt actions
+        gen_samples['acts'] = -1 + 2 * torch.rand(*gen_samples['acts'].shape)
+
+        discriminator.train_disc(expert_samples=expert_samples, gen_samples=gen_samples)
+
+    torch.save(discriminator.discrim_net.state_dict(), 'train_discrim.pt')

scratch/johannes/evaluation.py

@@ -0,0 +1,90 @@
+
+def evaluate_policy_simple(
+    model,
+    env: gym.Env,
+    n_eval_episodes: int = 10,
+    deterministic: bool = True,
+    render: bool = False,
+    callback = None,
+    reward_threshold = None,
+    return_episode_rewards: bool = False,
+    warn: bool = True,
+):
+    """
+    Runs policy for ``n_eval_episodes`` episodes and returns average reward.
+    If a vector env is passed in, this divides the episodes to evaluate onto the
+    different elements of the vector env. This static division of work is done to
+    remove bias. See https://github.com/DLR-RM/stable-baselines3/issues/402 for more
+    details and discussion.
+
+    .. note::
+        If environment has not been wrapped with ``Monitor`` wrapper, reward and
+        episode lengths are counted as it appears with ``env.step`` calls. If
+        the environment contains wrappers that modify rewards or episode lengths
+        (e.g. reward scaling, early episode reset), these will affect the evaluation
+        results as well. You can avoid this by wrapping environment with ``Monitor``
+        wrapper before anything else.
+
+    :param model: The RL agent you want to evaluate.
+    :param env: The gym environment or ``VecEnv`` environment.
+    :param n_eval_episodes: Number of episode to evaluate the agent
+    :param deterministic: Whether to use deterministic or stochastic actions
+    :param render: Whether to render the environment or not
+    :param callback: callback function to do additional checks,
+        called after each step. Gets locals() and globals() passed as parameters.
+    :param reward_threshold: Minimum expected reward per episode,
+        this will raise an error if the performance is not met
+    :param return_episode_rewards: If True, a list of rewards and episode lengths
+        per episode will be returned instead of the mean.
+    :param warn: If True (default), warns user about lack of a Monitor wrapper in the
+        evaluation environment.
+    :return: Mean reward per episode, std of reward per episode.
+        Returns ([float], [int]) when ``return_episode_rewards`` is True, first
+        list containing per-episode rewards and second containing per-episode lengths
+        (in number of steps).
+    """
+    episode_rewards = []
+    episode_lengths = []
+
+    episode_counts = 0
+
+    current_rewards = 0
+    current_lengths = 0
+    observations = env.reset()
+    states = None
+    while (episode_counts < n_eval_episodes):
+        actions, states = model.predict(observations, state=states, deterministic=deterministic)
+        observations, rewards, dones, infos = env.step(actions)
+        print(env._env.t)
+        current_rewards += rewards
+        current_lengths += 1
+
+        # unpack values so that the callback can access the local variables
+        reward = rewards
+        done = dones
+        info = infos
+        if info['collision']:
+            print("COLLISION")
+
+        if callback is not None:
+            callback(locals(), globals())
+
+        if dones:
+            episode_rewards.append(current_rewards)
+            episode_lengths.append(current_lengths)
+            episode_counts += 1
+            current_rewards = 0
+            current_lengths = 0
+            if states is not None:
+                states *= 0
+
+        if render:
+            env.render()
+
+    mean_reward = np.mean(episode_rewards)
+    std_reward = np.std(episode_rewards)
+    if reward_threshold is not None:
+        assert mean_reward > reward_threshold, "Mean reward below threshold: " f"{mean_reward:.2f} < {reward_threshold:.2f}"
+    if return_episode_rewards:
+        return episode_rewards, episode_lengths
+    return mean_reward, std_reward

scratch/johannes/intersimple/gail_options_image.py

@@ -0,0 +1,101 @@
+# %%
+import sys
+sys.path.append('../../../')
+
+from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
+from imitation.algorithms import adversarial
+import stable_baselines3
+import torch.utils.data
+import numpy as np
+from intersim.envs.intersimple import NRasterized, speed_reward
+import itertools
+import functools
+from torch.distributions import Categorical
+import gym
+import torch
+import pickle
+import imitation.data.rollout as rollout
+import tempfile
+import pathlib
+from imitation.util import logger
+from stable_baselines3.common.env_util import make_vec_env
+from tqdm import tqdm
+from src.policies.options import OptionsCnnPolicy
+from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
+from src.gail.train import train_discriminator, train_generator
+from src.evaluation.evaluation import Evaluation
+from torch.utils.tensorboard import SummaryWriter
+
+
+model_name = 'gail_options_image'
+env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2}
+
+ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10]] # option 0 is safe fallback
+
+def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, discount=0.99):
+    env = NRasterized(**env_settings)
+    env.discount = discount
+
+    tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
+    tempdir_path = pathlib.Path(tempdir.name)
+    logger.configure(tempdir_path / "GAIL/")
+    print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
+
+    venv = make_vec_env(NRasterized, n_envs=1, env_kwargs=env_settings)
+    discriminator = adversarial.GAIL(
+        expert_data=expert_data,
+        expert_batch_size=expert_batch_size,
+        discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
+        #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
+        venv=venv, # unused
+        gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
+    )
+
+    generator = stable_baselines3.PPO(
+        OptionsCnnPolicy,
+        OptionsEnv(env, options=ALL_OPTIONS),
+        verbose=1,
+        n_steps=generator_steps,
+    )
+
+    # PPO.train requires logger as set up in
+    # PPO._setup_learn (called by PPO.learn)
+    generator._logger = stable_baselines3.common.utils.configure_logger(
+        generator.verbose,
+        generator.tensorboard_log,
+    )
+
+    filestr = os.path.join('out', model_name)
+    writer = SummaryWriter(filestr)
+    ev = Evaluation(filestr, env, expert_data, n_eval_episodes=100)
+    for epoch in tqdm(range(epochs)):
+        train_discriminator(LLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=expert_batch_size)
+        train_generator(HLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=generator_steps)
+
+        metrics = ev.evaluate(epoch, generator, discriminator)
+        for metric, value in metrics.items():
+            writer.add_scalar(metric, value, epoch)
+    
+    return generator
+
+# %%
+if __name__ == '__main__':
+    # %%
+ 
+    with open("scratch/etienne/intersimple/data/NormalizedIntersimpleExpertMu.001_NRasterizedInfoAgent51w36h36mppx2.pkl", "rb") as f:
+        trajectories = pickle.load(f)
+    transitions = rollout.flatten_trajectories(trajectories)
+    generator = train(transitions)
+
+    generator.save(model_name)
+
+    # %%
+    model = stable_baselines3.PPO.load(model_name)
+
+    env = RenderOptions(NRasterized(**env_settings), options=ALL_OPTIONS)
+
+    for s in env.sample_ll(model):
+        if s['dones']:
+            break
+
+    env.close(filestr='render/'+model_name)

scratch/johannes/intersimple/gail_options_image_random.py

@@ -0,0 +1,101 @@
+# %%
+import sys
+sys.path.append('../../../')
+
+from src.discriminator import CnnDiscriminatorFlatAction
+from imitation.algorithms import adversarial
+import stable_baselines3
+import torch.utils.data
+import numpy as np
+from intersim.envs.intersimple import NRasterizedRandomAgent
+import itertools
+from torch.distributions import Categorical
+import gym
+import torch
+import pickle
+import imitation.data.rollout as rollout
+import tempfile
+import pathlib
+from imitation.util import logger
+from stable_baselines3.common.env_util import make_vec_env
+from tqdm import tqdm
+from src.policies.options import OptionsCnnPolicy
+from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
+from src.gail.train import train_discriminator, train_generator
+from src.evaluation.evaluation import Evaluation
+from torch.utils.tensorboard import SummaryWriter
+import os 
+
+model_name = 'gail_options_image_random'
+env_settings = {'width': 36, 'height': 36, 'm_per_px': 2}
+
+ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback
+
+def train(expert_data, epochs=100, expert_batch_size=16, generator_steps=16, discount=0.99):
+    env = NRasterizedRandomAgent(**env_settings)
+    env.discount = discount
+
+    tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
+    tempdir_path = pathlib.Path(tempdir.name)
+    logger.configure(tempdir_path / "GAIL/")
+    print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
+
+    venv = make_vec_env(NRasterizedRandomAgent, n_envs=1, env_kwargs=env_settings)
+    discriminator = adversarial.GAIL(
+        expert_data=expert_data,
+        expert_batch_size=expert_batch_size,
+        discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
+        #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
+        venv=venv, # unused
+        gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
+    )
+
+    generator = stable_baselines3.PPO(
+        OptionsCnnPolicy,
+        OptionsEnv(env, options=ALL_OPTIONS),
+        verbose=1,
+        n_steps=generator_steps,
+    )
+
+    # PPO.train requires logger as set up in
+    # PPO._setup_learn (called by PPO.learn)
+    generator._logger = stable_baselines3.common.utils.configure_logger(
+        generator.verbose,
+        generator.tensorboard_log,
+    )
+    
+    filestr = os.path.join('out', model_name)
+    writer = SummaryWriter(filestr)
+    ev = Evaluation(filestr, env, expert_data, n_eval_episodes=100)
+    for epoch in tqdm(range(epochs)):
+        train_discriminator(LLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=expert_batch_size)
+        train_generator(HLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=generator_steps)
+        generator.save(model_name)
+    
+        metrics = ev.evaluate(epoch, generator, discriminator)
+        for metric, value in metrics.items():
+            writer.add_scalar(metric, value, epoch)
+    
+    return generator
+
+def video(model_name, env):
+    model = stable_baselines3.PPO.load(model_name)
+    env = RenderOptions(env, options=ALL_OPTIONS)
+    for s in env.sample_ll(model):
+        if s['dones']:
+            break
+    env.close(filestr='render/'+model_name)
+
+def evaluate():
+    video(
+        model_name=model_name,
+        env=NRasterizedRandomAgent(**env_settings)
+    )
+
+# %%
+if __name__ == '__main__':
+
+    with open("scratch/etienne/intersimple/data/NormalizedIntersimpleExpertMu.001N10000_NRasterizedRandomAgentInfow36h36mppx2.pkl", "rb") as f:
+        trajectories = pickle.load(f)
+    transitions = rollout.flatten_trajectories(trajectories)
+    train(transitions)

scratch/johannes/raytune_simple.py

@@ -0,0 +1,69 @@
+"""This example demonstrates basic Ray Tune random search and grid search."""
+import time
+
+import ray
+from ray import tune
+
+
+def evaluation_fn(step, width, height):
+    time.sleep(0.1)
+    return (0.1 + width * step / 100)**(-1) + height * 0.1
+
+def easy_objective(config):
+    # Hyperparameters
+    width, height = config["width"], config["height"]
+
+    mydata = ray.get(ray_data)
+    print(mydata)
+
+    for step in range(config["steps"]):
+        # Iterative training function - can be any arbitrary training procedure
+        intermediate_score = evaluation_fn(step, width, height)
+        # Feed the score back back to Tune.
+        tune.report(iterations=step, mean_loss=intermediate_score)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--smoke-test", action="store_true", help="Finish quickly for testing")
+    parser.add_argument(
+        "--server-address",
+        type=str,
+        default=None,
+        required=False,
+        help="The address of server to connect to if using "
+        "Ray Client.")
+    args, _ = parser.parse_known_args()
+    if args.server_address is not None:
+        ray.init(f"ray://{args.server_address}")
+    else:
+        ray.init(configure_logging=False)
+
+    # This will do a grid search over the `activation` parameter. This means
+    # that each of the two values (`relu` and `tanh`) will be sampled once
+    # for each sample (`num_samples`). We end up with 2 * 50 = 100 samples.
+    # The `width` and `height` parameters are sampled randomly.
+    # `steps` is a constant parameter.
+
+    import numpy as np
+    N = 3
+    data = np.random.rand(N,N,N)
+    ray_data = ray.put(data)
+
+
+    analysis = tune.run(
+        easy_objective,
+        metric="mean_loss",
+        mode="min",
+        num_samples=5 if args.smoke_test else 50,
+        config={
+            "steps": 5 if args.smoke_test else 100,
+            "width": tune.uniform(0, 20),
+            "height": tune.uniform(-100, 100),
+            "activation": tune.grid_search(["relu", "tanh"])
+        })
+
+    print("Best hyperparameters found were: ", analysis.best_config)

src/__init__.py

@@ -1,3 +1,3 @@
-from src.expert_data import generate_expert_data, load_expert_data
+from src.data.expert_data import generate_expert_data, load_expert_data
-from src.data_utils import InteractionDatasetSingleAgent
+from src.data.data_utils import InteractionDatasetSingleAgent
-from src.metrics import metrics
+from src.evaluation.metrics import metrics

src/data/__init__.py

@@ -0,0 +1 @@
+from src.data.expert import single_agent_expert, single_agent_demonstrations, multi_agent_demonstrations, load_experts, process_experts 

src/data/data_utils.py

@@ -1,7 +1,7 @@
 import torch
 from torch.utils.data import Dataset
 import numpy as np
-from src.expert_data import load_expert_data
+from src.data.expert_data import load_expert_data
 import os
 opj = os.path.join
 

src/data/expert.py

@@ -0,0 +1,241 @@
+import intersim
+from intersim.envs.intersimple import Intersimple
+from stable_baselines3.common.policies import BasePolicy
+import gym
+import intersim.envs.intersimple
+import pickle
+from tqdm import tqdm
+import copy
+import os
+import numpy as np
+
+from src.util.rollout import rollout_and_save, flatten_trajectories, make_sample_until
+
+class IntersimExpert(BasePolicy):
+    
+    def __init__(self, intersim_env, mu=0, *args, **kwargs):
+        super().__init__(
+            observation_space=gym.spaces.Space(),
+            action_space=gym.spaces.Space(),
+            *args, **kwargs
+        )
+        self._intersim = intersim_env
+        self._mu = mu
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def _predict(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def _action(self):
+        target_t = min(self._intersim._ind + 1, len(self._intersim._svt.simstate) - 1)
+        target_state = self._intersim._svt.simstate[target_t]
+        return self._intersim.target_state(target_state, mu=self._mu)
+
+    def predict(self, *args, **kwargs):
+        return self._action(), None
+
+class IntersimpleExpert(BasePolicy):
+
+    def __init__(self, intersimple_env, mu=0, *args, **kwargs):
+        super().__init__(
+            observation_space=intersimple_env.observation_space,
+            action_space=intersimple_env.action_space,
+            *args, **kwargs
+        )
+        self._intersimple = intersimple_env
+        self._intersim_expert = IntersimExpert(intersimple_env._env, mu=mu)
+    
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def _predict(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def _action(self):
+        return self._intersim_expert._action()[self._intersimple._agent]
+
+    def predict(self, *args, **kwargs):
+        return self._action(), None
+
+class NormalizedIntersimpleExpert(IntersimpleExpert):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def predict(self, *args, **kwargs):
+        action, _ = super().predict(*args, **kwargs)
+        return self._intersimple._normalize(action), None
+
+def save_video(env, expert):
+    env.reset()
+    env.render()
+    done = False
+    while not done:
+        actions, _ = expert.predict()
+        _, _, done, _ = env.step(actions)
+        env.render()
+    env.close()
+
+def load_experts(expert_files, flatten=True):
+    """
+    Load expert trajectories from files and combine their transitions into a single RB
+
+    Args:
+        expert_files (list): list of expert file strings
+        flatten (bool): whether to flatten trajectory info
+    Returns:
+        transitions (list): list of combined expert episode transitions 
+    """
+    transitions = []
+    for file in tqdm(expert_files):
+        with open(file, "rb") as f:
+            new_trajectories = pickle.load(f)
+        transitions += new_trajectories
+    if flatten:
+        transitions = flatten_trajectories(transitions)
+    return transitions
+
+def single_agent_expert(expert='NormalizedIntersimpleExpert', 
+                        env='NRasterizedRouteIncrementingAgent', 
+                        env_args={}, policy_args={}, **kwargs):
+    """
+    Args: 
+        expert (class): class of expert
+        env (class): class of env intersim.envs.intersimple
+        env_args (dict): dictionary of kwargs when instantiating environment class
+        policy_args (dict): dictionary of kwargs when instantiating Expert policy
+        path (str): path to store output
+        min_timesteps (int): min number of timesteps for call to rollout.rollout_and_save
+        min_episodes (int): min number of episodes for call to rollout.rollout_and_save
+        video (bool): whether to save a video of the expert until a single environment instantiation stops
+    """
+    Env = intersim.envs.intersimple.__dict__[env]
+    Expert = globals()[expert]
+    env = Env(**env_args)
+    policy = Expert(env, **policy_args)
+    single_agent_demonstrations(env, policy, **kwargs)
+
+def single_agent_demonstrations(env, policy, 
+                                path=None, min_timesteps=None, 
+                                min_episodes=None, video=False,
+                                env_args={}, policy_args={}):
+    """Rollout and save expert demos.
+    
+    Usage:
+        python -m intersimple.expert <flags>
+    Args: 
+        env (class): intersimple environment
+        policy (BasePolicy): intersimple policy
+        path (str): path to store output
+        min_timesteps (int): min number of timesteps for call to rollout.rollout_and_save
+        min_episodes (int): min number of episodes for call to rollout.rollout_and_save
+        video (bool): whether to save a video of the expert until a single environment instantiation stops
+        env_args (dict): dictionary of kwargs when instantiating environment class
+        policy_args (dict): dictionary of kwargs when instantiating Expert policy
+    """
+    
+    if min_timesteps is None and min_episodes is None:
+        min_episodes = env.nv # one episode per vehicle being controlled in environment (hopefully an incrementing agent environment)
+
+    if video:
+        save_video(env, policy)
+
+    path = path or (policy.__class__.__name__ + '_' + env.__class__.__name__ + '.pkl')
+    suntil = make_sample_until(min_timesteps=min_timesteps,min_episodes=min_episodes)
+    
+    rollout_and_save(
+        path=path,
+        policy=policy,
+        env=env,
+        sample_until=suntil
+    )
+
+def multi_agent_demonstrations(expert='IntersimExpert',path=None, env_args={}, policy_args={}):
+    """
+    Run and save the `intersim' multiagent environment demonstration
+    
+    Args: 
+        expert (class): class of multi-agent expert
+        path (str): path to store output data
+        env_args (dict): dictionary of kwargs when instantiating environment class
+        policy_args (dict): dictionary of kwargs when instantiating Expert policy
+    """
+    if path is None:
+        raise('No path specified')
+
+    env = gym.make('intersim:intersim-v0',**env_args)
+    Expert = globals()[expert] 
+    policy = Expert(env, **policy_args)
+
+    s, done = env.reset(), False
+    env.render(mode='file')
+    while not done:
+        _,_,done,_ = env.step(policy.predict()[0])
+        env.render(mode='file')
+    env.close(filestr=path)
+
+def process_experts(filename:str='expert.pkl', 
+                    locs:list=None, 
+                    tracks:list=None,
+                    env_class:str='NRasterizedRouteIncrementingAgent',
+                    env_args:dict={'width':36,'height':36,'m_per_px':2},
+                    expert_class:str='NormalizedIntersimpleExpert',  
+                    expert_args:dict={'mu':0.001}):
+    """
+    Process all experts in the Interaction Dataset
+    For now, using NormalizedIntersimpleExpert with NRasterizedRouteIncrementingAgent environment
+
+    Args:
+        filename (str): name for track file
+        locs (list): list of location ids
+        tracks (list): list of track numbers
+        env_class (str): class of environment
+        env_args (dict): default environment kwargs
+        expert_class (str): class of expert
+        expert_args (dict): default expert kwargs
+    """
+    locs = locs or intersim.LOCATIONS
+    tracks = tracks or range(intersim.MAX_TRACKS)
+    pbar = tqdm(total=len(locs)*len(tracks))
+    for loc in locs:
+        for track in tracks:
+            
+            iloc = intersim.LOCATIONS.index(loc)
+
+            it_env_args = copy.deepcopy(env_args)
+            env_loc_args = {
+                'loc':iloc,
+                'track':track,
+            }
+            it_env_args.update(env_loc_args)
+            out_folder = os.path.join('expert_data',loc, 'track%04i'%(track))
+            if not os.path.isdir(out_folder):
+                os.makedirs(out_folder) 
+            it_path = os.path.join(out_folder,filename)
+
+            # Multi-Agent demonstrations
+            it_ma_path = os.path.join(out_folder,'joint_expert')
+            multi_agent_demonstrations(
+                expert='IntersimExpert',
+                path=it_ma_path,
+                env_args=env_loc_args,
+                policy_args=expert_args
+            )
+            # Single-Agent POV Demonstrations
+            single_agent_expert(
+                expert=expert_class, 
+                env=env_class, 
+                path=it_path, 
+                env_args=it_env_args, 
+                policy_args=expert_args,
+                )
+            
+            pbar.update(1)
+    pbar.close()
+    
+if __name__=='__main__':
+    import fire
+    fire.Fire(process_experts)
+

src/discriminator/__init__.py

@@ -0,0 +1 @@
+from src.discriminator.discriminator import *

src/discriminator/discriminator.py

@@ -0,0 +1,99 @@
+import torch
+
+# imitation.rewards.discrim_nets.DiscrimNetGAIL is composed of self.discriminator (nn.Module),
+# which gets called with inputs (state, action) when needed.
+
+class CnnDiscriminator(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 + action_size
+
+        self.cnn = torch.nn.Sequential(
+            torch.nn.Conv2d(in_channels, 4, kernel_size=8, stride=4, padding=0),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(4, 8, kernel_size=4, stride=2, padding=0),
+            torch.nn.ReLU(),
+            torch.nn.Flatten(start_dim=-3, end_dim=-1),
+            torch.nn.LazyLinear(512),
+            torch.nn.ReLU(),
+            torch.nn.LazyLinear(1),
+        )
+    
+    @staticmethod
+    def _concatenate(state, action):
+        b, _, h, w = state.shape
+        _, a = action.shape
+        act = action.unsqueeze(-1).unsqueeze(-1).expand((b, a, h, w))
+        sa = torch.cat((state, act), -3)
+        return sa
+
+    def forward(self, state, action):
+        sa = self._concatenate(state, action)
+        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.float())
+        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."""
+
+    def __init__(self, env=None):
+        super().__init__()
+        self.flatten = torch.nn.Flatten(start_dim=1, end_dim=-1)
+        self.mlp = torch.nn.Sequential(
+            torch.nn.LazyLinear(64), # 42 -> 64
+            torch.nn.Tanh(),
+            torch.nn.LazyLinear(64), # 64 -> 64
+            torch.nn.Tanh(),
+            torch.nn.LazyLinear(1),  # 64 -> 1
+        )
+    
+    def forward(self, state, action):
+        flat = self.flatten(state)
+        sa = torch.cat((action, flat), -1)
+        assert sa.ndim == 2
+        return self.mlp(sa).squeeze(1)

src/discriminator/test_discriminator.py

@@ -0,0 +1,45 @@
+from intersim.envs.intersimple import NRasterized
+from discriminator import CnnDiscriminator
+import torch
+
+def test_image_concatenation():
+    env = NRasterized()
+    disc = CnnDiscriminator(env)
+    s = torch.tensor(env.reset()).unsqueeze(0)
+    a = torch.tensor([[0.5]])
+    sa = disc._concatenate(s, a)
+
+    assert s.shape == (1, 5, 200, 200)
+    assert a.shape == (1, 1)
+    assert sa.shape == (1, 6, 200, 200)
+    assert torch.allclose(sa[:, :5], 1.0 * s)
+    assert (sa[:, 5] == a.unsqueeze(-1)).all()
+
+def test_image_concatenation3():
+    env = NRasterized()
+    disc = CnnDiscriminator(env)
+
+    s1 = env.reset()
+    a1 = 0.15
+    s2, _, _, _ = env.step(0.9)
+    a2 = 0.25
+    s3, _, _, _ = env.step(-0.9)
+    a3 = 0.35
+    
+    s = torch.stack([
+        torch.tensor(s1),
+        torch.tensor(s2),
+        torch.tensor(s3)
+    ], axis=0)
+    a = torch.tensor([
+        [a1],
+        [a2],
+        [a3],
+    ])
+    sa = disc._concatenate(s, a)
+
+    assert s.shape == (3, 5, 200, 200)
+    assert a.shape == (3, 1)
+    assert sa.shape == (3, 6, 200, 200)
+    assert torch.allclose(sa[:, :5], 1.0 * s)
+    assert (sa[:, 5] == a.unsqueeze(-1)).all()

src/evaluation/evaluation.py

@@ -0,0 +1,93 @@
+import torch
+import numpy as np
+from stable_baselines3.common.vec_env import VecEnv
+from stable_baselines3.common.evaluation import evaluate_policy
+
+from intersim.envs.intersimple import Intersimple
+from src.evaluation.metrics import nanmean, divergence, visualize_distribution
+import os
+
+
+class Evaluation:
+    def __init__(self, filestr, eval_env, expert_data, n_eval_episodes=10):
+        # if env is a VecEnv, the code needs to be adapted, since the callback will be called after each step, 
+        # so transitions of different envs will be mixed and the total number of episodes could be larger than n_eval_episodes!
+        assert not isinstance(eval_env, VecEnv)
+        self.filestr = filestr
+        self.env = eval_env
+        self.n_eval_episodes = n_eval_episodes
+        self.expert_data = expert_data
+        self.compute_expert_features(expert_data)
+        self.reset()
+
+    def reset(self):
+        self._n_collisions = 0
+        self._trajectories = []
+        self._episode_done = True
+        self._accelerations = []
+
+    def compute_expert_features(self, expert_data):
+        # expert velocities
+        extract_state = lambda info: info['projected_state'][info['agent']]
+        expert_velocities = torch.stack([extract_state(info) for info in expert_data.infos])[:,2]
+        self.expert_velocities = expert_velocities[~torch.isnan(expert_velocities)]
+        # expert accelerations
+        extract_accel = lambda info: info['action_taken'][info['agent']]
+        self.expert_accelerations = torch.cat([extract_accel(info) for info in expert_data.infos])
+
+    def evaluate(self, epoch, generator, discriminator):
+        self.reset()
+        metrics = {}
+        
+        episode_rewards, episode_lengths = evaluate_policy(
+            generator, 
+            self.env,
+            n_eval_episodes=self.n_eval_episodes,
+            callback=self.evaluate_policy_callback, 
+            return_episode_rewards=True
+        )
+
+        collision_rate = self._n_collisions / self.n_eval_episodes
+        metrics['collision_rate'] = collision_rate
+
+        assert len(self._trajectories) >= self.n_eval_episodes
+
+        # velocities produced by generator
+        policy_velocities = torch.cat([torch.stack(t)[:,2] for t in self._trajectories])
+        # if episodes terminate without collisions, then the state is fully nan
+        policy_velocities = policy_velocities[~torch.isnan(policy_velocities)]
+
+        metrics['avg_velocity_loss'] = (self.expert_velocities.mean() - policy_velocities.mean()).item()
+        metrics['velocity_divergence'] = divergence(policy_velocities, self.expert_velocities, type='js')
+        
+
+        # accelerations produced by generator
+        policy_accelerations = torch.tensor(self._accelerations)
+
+        metrics['acceleration_divergence'] = divergence(policy_accelerations, self.expert_accelerations, type='js')
+        visualize_distribution(self.expert_accelerations, policy_accelerations, os.path.join(self.filestr, '_action_viz{:02}'.format(epoch))) 
+
+        print(metrics)
+        return metrics
+
+    def evaluate_policy_callback(self, local_vars, global_vars):
+        venv_i = local_vars['i']
+        info = local_vars['info']
+        done = local_vars['done']
+        _agent = info['agent']
+        env = local_vars['env'].envs[venv_i]
+        assert isinstance(env, Intersimple)
+
+        # Increase collision counter if episode terminated with a collision
+        if info['collision']:
+            assert done
+            self._n_collisions += 1
+
+        # if last episode is done, start new trajectory
+        # this is currently not necessary, only if velocity is to be averaged over individual trajectories first
+        # and then averaging over all trajectories
+        if self._episode_done:
+            self._trajectories.append([])
+        self._trajectories[-1].append(info['projected_state'][_agent])
+        self._accelerations.append(info['action_taken'][_agent])
+        self._episode_done = done

src/gail/options.py

@@ -0,0 +1,165 @@
+import gym
+import torch
+from src.util.collisions import feasible
+import numpy as np
+
+class OptionsEnv(gym.Wrapper):
+    
+    def __init__(self, env, options=[(0, 5), (5, 5), (10, 5)], *args, **kwargs):
+        """option 0 is treated as safe fallback"""
+
+        super().__init__(env, *args, **kwargs)
+        self.options = options
+        num_hl_options = len(self.options)
+        self.action_space = gym.spaces.Discrete(num_hl_options)
+        self.observation_space = gym.spaces.Dict({
+            'obs': env.observation_space,
+            'mask': gym.spaces.Box(low=0, high=1, shape=(num_hl_options,)),
+        })
+
+    def _after_choice(self):
+        pass
+    
+    def _after_step(self):
+        pass
+
+    def _transitions(self):
+        raise NotImplementedError('Use `LLOptions` or `HLOptions` for sampling.')
+    
+    def sample(self, generator):
+        self.done = True
+        while True:
+            self.episode_start = False
+            if self.done:
+                self.s = self.env.reset()
+                self.done = False
+                self.episode_start = True
+
+            self.m = available_actions(self.env, self.options)
+            if not self.m.any():
+                # action 0 is considered safe fallback
+                self.m[0] = True
+            
+            self.ch, self.value, self.log_prob = generator.policy.forward({
+                'obs': torch.tensor(self.s).unsqueeze(0).to(generator.policy.device),
+                'mask': self.m.unsqueeze(0).to(generator.policy.device),
+            })
+            self.plan = list(map(float, generate_plan(self.env, self.ch, self.options)))
+
+            self._after_choice()
+
+            assert not self.done
+            assert self.plan
+            #assert feasible(self.env, self.plan, self.ch)
+
+            while not self.done and self.plan and \
+                (feasible(self.env, safety_plan(self.env, self.plan)) or self.m.sum() == 1):
+                
+                self.a, self.plan = self.plan[0], self.plan[1:]
+                self.a = self.env._normalize(self.a)
+                self.nexts, _, self.done, _ = self.env.step(self.a)
+
+                self._after_step()
+
+                self.s = self.nexts
+            
+            yield from self._transitions()
+
+class LLOptions(OptionsEnv):
+    """Sample low-level (state, action) tuples for discriminator training."""
+    
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.observation_space = self.observation_space['obs']
+
+    def _after_choice(self):
+        self._transition_buffer = []
+
+    def _after_step(self):
+        self._transition_buffer.append({
+            'obs': self.s,
+            'next_obs': self.nexts,
+            'acts': np.array((self.a,)),
+            'dones': np.array(self.done),
+        })
+
+    def _transitions(self):
+        yield from self._transition_buffer
+    
+    def sample_ll(self, policy):
+        return self.sample(policy)
+
+class HLOptions(OptionsEnv):
+    """Sample high-level (state, action, reward) tuples for generator training."""
+    
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _after_choice(self):
+        self.obs = {'obs': np.copy(self.s), 'mask': np.copy(self.m)}
+        self.r = 0
+        self.steps = 0
+
+    def _after_step(self):
+        self.r += self.discount**self.steps * self.discriminator.discrim_net.predict_reward_train(
+            state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()),
+            action=torch.tensor([[self.a]]).to(self.discriminator.discrim_net.device()),
+            next_state=torch.tensor(self.s).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused
+            done=torch.tensor(self.done).unsqueeze(0).to(self.discriminator.discrim_net.device()), # unused
+        )
+        self.steps += 1
+    
+    def _transitions(self):
+        yield {
+            'obs': self.obs,
+            'action': self.ch.cpu(),
+            'reward': self.r,
+            'episode_start': self.episode_start,
+            'value': self.value.detach(),
+            'log_prob': self.log_prob.detach(),
+            'done': self.done,
+        }
+    
+    def sample_hl(self, policy, discriminator):
+        self.discriminator = discriminator
+        return self.sample(policy)
+
+class RenderOptions(LLOptions):
+
+    def _after_step(self):
+        super()._after_step()
+        self.env.render()
+    
+    def close(self, *args, **kwargs):
+        self.env.close(*args, **kwargs)
+
+def safety_plan(env, plan):
+    return np.concatenate((plan, np.array(5 * [env._env._min_acc])), axis=0)
+
+def available_actions(env, options):
+    """Return mask of available actions given current `env` state."""
+    plans = [generate_plan(env, i, options) for i, _ in enumerate(options)]
+    # is emergency braking still possible?
+    plans = list(map(lambda p: safety_plan(env, p), plans))
+
+    T = max(len(p) for p in plans)
+    plans = [np.pad(p, ((0, T-len(p)),), constant_values=np.nan) for p in plans]
+    plans = np.stack(plans, axis=0)
+    
+    valid = feasible(env, plans)
+    return valid
+
+def target_velocity_plan(current_v: float, target_v: float, t: int, dt: float):
+    """Smoothly target a velocity in a given number of steps"""
+    # for now, constant acceleration
+    a = (target_v - current_v)  / (t * dt)
+    return a*np.ones((t,))
+
+def generate_plan(env, i, options):
+    """Generate input profile for high-level action `i`."""
+    assert i < len(options), "Invalid option index {i}"
+    target_v, t = options[i]
+    current_v = env._env.state[env._agent, 1].item() # extract from env
+    plan = target_velocity_plan(current_v, target_v, t, env._env._dt)
+    assert len(plan) == t, "incorrect plan length"
+    return plan

src/gail/test_options.py

@@ -0,0 +1,59 @@
+import pickle
+import imitation.data.rollout as rollout
+from options import LLOptions, OptionsEnv
+from intersim.envs import NRasterized
+import itertools
+import stable_baselines3
+from policy import OptionsCnnPolicy
+from train import flatten_transitions
+import numpy as np
+
+def test_ll_expert_data():
+    with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
+        expert_trajectories = pickle.load(f)
+    expert_transitions = rollout.flatten_trajectories(expert_trajectories)
+
+    env = LLOptions(NRasterized(agent=51, width=36, height=36, m_per_px=2))
+
+    gen_transitions = list(itertools.islice(env.sample_ll(
+        policy=stable_baselines3.PPO(
+            OptionsCnnPolicy,
+            OptionsEnv(env),
+            verbose=1,
+        )
+    ), 10))
+    gen_transitions = flatten_transitions(gen_transitions)
+
+    assert expert_transitions[:10].obs.shape == gen_transitions['obs'].shape
+    assert expert_transitions[:10].next_obs.shape == gen_transitions['next_obs'].shape
+    assert expert_transitions[:10].acts.shape == gen_transitions['acts'].shape
+    assert expert_transitions[:10].dones.shape == gen_transitions['dones'].shape
+
+def test_ll_states():
+    env = NRasterized()
+    policy = stable_baselines3.PPO(
+        OptionsCnnPolicy,
+        OptionsEnv(env),
+        verbose=1,
+    )
+    llenv = LLOptions(env)
+    transitions = list(itertools.islice(llenv.sample_ll(policy=policy), 100))
+
+    env2 = NRasterized()
+    s2 = env2.reset()
+    for i, t in enumerate(transitions):
+        assert i == 0 or np.array_equal(t['obs'], transitions[i-1]['next_obs'])
+        assert np.array_equal(t['obs'], s2)
+        assert t['acts'].shape == (1,)
+
+        nexts2, _, done2, _ = env2.step(t['acts'])
+        assert np.array_equal(t['next_obs'], nexts2)
+        assert np.array_equal(t['dones'], done2)
+
+        if done2:
+            break
+
+        s2 = nexts2
+
+def test_hl_transitions():
+    pass

src/gail/train.py

@@ -0,0 +1,37 @@
+import numpy as np
+import itertools
+
+def flatten_transitions(transitions):
+    return {
+        'obs': np.stack(list(t['obs'] for t in transitions), axis=0),
+        'next_obs': np.stack(list(t['next_obs'] for t in transitions), axis=0),
+        'acts': np.stack(list(t['acts'] for t in transitions), axis=0),
+        'dones': np.stack(list(t['dones'] for t in transitions), axis=0),
+    }
+
+def train_discriminator(env, generator, discriminator, num_samples, n_updates=1):
+    transitions = list(itertools.islice(env.sample_ll(generator), num_samples))
+    generator_samples = flatten_transitions(transitions)
+    for _ in range(n_updates):
+        discriminator.train_disc(gen_samples=generator_samples)
+
+def train_generator(env, generator, discriminator, num_samples):
+    generator_samples = list(itertools.islice(env.sample_hl(generator, discriminator), num_samples+1))
+    
+    generator.rollout_buffer.reset()
+    for s in generator_samples[:-1]:
+        generator.rollout_buffer.add(
+            obs=s['obs'],
+            action=s['action'],
+            reward=s['reward'],
+            episode_start=s['episode_start'],
+            value=s['value'],
+            log_prob=s['log_prob'],
+        )
+    
+    generator.rollout_buffer.compute_returns_and_advantage(
+        last_values=generator_samples[-1]['value'],
+        dones=generator_samples[-1]['done'],
+    )
+    
+    generator.train()

src/policies/__init__.py

@@ -1 +1,2 @@
 from src.policies.policy import IntersimPolicy, IntersimStateNet, IntersimStateActionNet, generate_transforms
+from src.policies.options import OptionsCnnPolicy

src/policies/options.py

@@ -0,0 +1,110 @@
+from stable_baselines3.common.policies import ActorCriticPolicy, ActorCriticCnnPolicy
+from torch.distributions import Categorical
+
+import gym
+import torch
+import torch.nn as nn
+from stable_baselines3.common.torch_layers import BaseFeaturesExtractor
+from stable_baselines3.common.preprocessing import is_image_space
+
+class CustomCNN(BaseFeaturesExtractor):
+    """
+    Smaller version of `stable_baselines3.common.torch_layers.NatureCNN`
+
+    :param observation_space:
+    :param features_dim: Number of features extracted.
+        This corresponds to the number of unit for the last layer.
+    """
+
+    def __init__(self, observation_space: gym.spaces.Box, features_dim: int = 512):
+        super().__init__(observation_space, features_dim)
+        # We assume CxHxW images (channels first)
+        # Re-ordering will be done by pre-preprocessing or wrapper
+        assert is_image_space(observation_space, check_channels=False), (
+            "You should use CustomCNN "
+            f"only with images not with {observation_space}\n"
+            "(you are probably using `CnnPolicy` instead of `MlpPolicy` or `MultiInputPolicy`)\n"
+            "If you are using a custom environment,\n"
+            "please check it using our env checker:\n"
+            "https://stable-baselines3.readthedocs.io/en/master/common/env_checker.html"
+        )
+        n_input_channels = observation_space.shape[0]
+        self.cnn = nn.Sequential(
+            nn.Conv2d(n_input_channels, 4, kernel_size=8, stride=4, padding=0),
+            nn.ReLU(),
+            nn.Conv2d(4, 8, kernel_size=4, stride=2, padding=0),
+            nn.ReLU(),
+            nn.Flatten(),
+        )
+
+        # Compute shape by doing one forward pass
+        with torch.no_grad():
+            n_flatten = self.cnn(torch.as_tensor(observation_space.sample()[None]).float()).shape[1]
+
+        self.linear = nn.Sequential(nn.Linear(n_flatten, features_dim), nn.ReLU())
+
+    def forward(self, observations: torch.Tensor) -> torch.Tensor:
+        return self.linear(self.cnn(observations))
+
+class OptionsCnnPolicy(ActorCriticPolicy):
+    """
+    Class for high-level options policy (generator)
+    """
+    def __init__(self, observation_space, *args, eps=0, **kwargs):
+        super().__init__(observation_space, *args, **kwargs)
+        self.cnn_policy = ActorCriticCnnPolicy(observation_space['obs'], *args, features_extractor_class=CustomCNN, **kwargs)
+        self.eps = eps
+
+    def _prior_distribution(self, s):
+        """
+        Return prior distribution over high-level options (before masking)
+        Args:
+            s (torch.tensor): observation
+        Returns:
+            values (torch.tensor): values from critic
+            dist (torch.distributions): prior distribution over actions
+        """
+        latent_pi, latent_vf, latent_sde = self.cnn_policy._get_latent(s)
+        distribution = self.cnn_policy._get_action_dist_from_latent(latent_pi, latent_sde)
+        values = self.cnn_policy.value_net(latent_vf)
+        return values, distribution.distribution
+
+    def forward(self, obs):
+        """
+        Will mask invalid states before making action selections
+        Args:
+            obs: dict with keys:
+                obs (torch.tensor): (*,o)  true observations
+                mask (torch.tensor): (*,m) mask over valid actions
+        Returns:
+            ch (torch.tensor): (*,a) sampled actions
+            values (torch.tensor): (*,) predicted value at observation
+            log_probs (torch.tensor): (*,) log probabilities of selected actions
+        """
+        s, m = obs['obs'], obs['mask']
+        values, prior = self._prior_distribution(s)
+        posterior = Categorical((prior.probs + self.eps) * m)
+        ch = posterior.sample()
+        return ch, values, posterior.log_prob(ch)
+    
+    def _predict(self, obs, deterministic=False):
+        action, _, _ = self.forward(obs)
+        return action
+
+    def evaluate_actions(self, obs, ch):
+        """
+        Evaluate particular actions
+        Args:
+            obs: dict with keys:
+                obs (torch.tensor): (*,o) true observations
+                mask (torch.tensor): (*,m) masks over valid actions
+            ch (torch.tensor): (*,a) selected actions
+        Returns:
+            values (torch.tensor): (*,) predicted value at observation
+            log_probs (torch.tensor): (*,) log probabilities of selected actions
+            ent (torch.tensor): (*,) entropy of each distribution over actions
+        """
+        s, m = obs['obs'], obs['mask']
+        values, prior = self._prior_distribution(s)
+        posterior = Categorical((prior.probs + self.eps) * m)
+        return values, posterior.log_prob(ch), posterior.entropy() # additional values used by PPO.train

src/util/__init__.py

@@ -0,0 +1,2 @@
+from src.util.render_env import *
+from src.util.collisions import feasible

src/util/collisions.py

@@ -0,0 +1,154 @@
+import torch
+import numpy as np
+from intersim.collisions import state_to_polygon
+
+def feasible(env, plan, method='exact'):
+    """Check if input profile is feasible given current `env` state."""
+    # zero pad plan - Take (B, T) or (T,) np plan and convert it to (B, T, nv, 1) torch.Tensor
+    plan = torch.tensor(plan)
+    plan = plan.reshape(-1, plan.shape[-1])
+    full_plan = torch.zeros(*plan.shape, env._env._nv, 1)
+    full_plan[:, :, env._agent, 0] = plan
+
+    # check_future_collisions_fast takes in B-list and outputs (B,) bool tensor
+    if method=='circle':
+        valid = check_future_collisions_fast(env, full_plan) 
+    elif method=='ncircles':
+        valid = check_future_collisions_ncircles(env, full_plan)
+    elif method=='exact':
+        valid = check_future_collisions_exact(env, full_plan)
+    else:
+        raise NotImplementedError('Invalid collision-checking method')
+
+    return valid
+
+def check_future_collisions_ncircles(env, actions, n_circles:int=2):
+    """Checks whether `env._agent` would collide with other agents assuming `actions` as input.
+    
+    Vehicles are (over-)approximated by multiple circles.
+
+    Args:
+        env (gym.Env): current environment state
+        actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles
+    Returns:
+        feasible (torch.Tensor): tensor of shape (B,) indicating whether the respective action profiles are collision-free
+    """
+    assert n_circles >= 2
+    B, (T, nv, _) = len(actions), actions[0].shape
+
+    states = env._env.propagate_action_profile_vectorized(actions)
+    assert states.shape == (B, T, nv, 5)
+    centers = states[:, :, :, :2]
+    psi = states[:, :, :, 3]
+    lon = torch.stack([psi.cos(), psi.sin()],dim=-1) # (B, T, nv, 2)
+
+    # offset between [-env._env.lengths+env._env.widths/2, env._env.lengths/2-env._env.widths/2]
+    back = (-env._env._lengths/2+env._env._widths/2).unsqueeze(-1) # (nv, 1)
+    length = (env._env._lengths-env._env._widths).unsqueeze(-1) # (nv, 1)
+    diff_d =  back + length*(torch.arange(n_circles)/(n_circles-1)).unsqueeze(0)  # (nv, n_circles) 
+    assert diff_d.shape == (nv, n_circles)
+    
+    offsets = diff_d[None, None, :, :, None] *  lon[:, :, :, None, :]
+    assert offsets.shape == (B, T, nv, n_circles, 2)
+
+    expanded_centers=centers.unsqueeze(-2) + offsets #(B, T, nv, n_circles, 2)
+    assert expanded_centers.shape == (B, T, nv, n_circles, 2)
+    agent_centers = expanded_centers[:,:,env._agent:env._agent+1,:,:] #(B, T, 1, n_circles, 2)
+    ds = expanded_centers.reshape((B, T, nv*n_circles, 1, 2)) - agent_centers #(B, T, nv*nc,1, 2) - (B, T, 1, nc, 2) = (B, T, nv*nc, nc, 2) 
+
+    distance = (ds**2).sum(-1).sqrt().reshape((B, T, nv, n_circles, n_circles)) # (B, T, nv, nc, nc)
+    distance = torch.where(distance.isnan(), np.inf*torch.ones_like(distance), distance) # only collide with spawned agents
+    distance[:, :, env._agent] = np.inf # cannot collide with itself
+    assert distance.shape == (B, T, nv, n_circles, n_circles)
+
+    radius = env._env._widths*np.sqrt(2) / 2
+    min_distance = radius[env._agent] + radius
+    min_distance = min_distance[None, None, :, None, None]
+    assert min_distance.shape == (1, 1, nv, 1, 1)
+
+    return (distance > min_distance).all(-1).all(-1).all(-1).all(-1)
+
+def check_future_collisions_circle(env, actions):
+    """Compute collision information for circular vehicle approximations
+
+    Args:
+        env (gym.Env): current environment state
+        actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles
+    Returns:
+        states (torch.Tensor): tensor of shape (B, T, nv, 5) of future states based on the action profiles
+        collision_tensor (torch.Tensor): tensor of shape (B, T, nv) of bools indicating which plan collides with which vehicles in which time frame
+            false: colliding, true: not colliding
+    """
+    B, (T, nv, _) = len(actions), actions[0].shape
+
+    states = env._env.propagate_action_profile_vectorized(actions)
+    assert states.shape == (B, T, nv, 5)
+
+    distance = ((states[:, :, :, :2] - states[:, :, env._agent:env._agent+1, :2])**2).sum(-1).sqrt()
+    distance = torch.where(distance.isnan(), np.inf*torch.ones_like(distance), distance) # only collide with spawned agents
+    distance[:, :, env._agent] = np.inf # cannot collide with itself
+    assert distance.shape == (B, T, nv)
+
+    radius = (env._env._lengths**2 + env._env._widths**2).sqrt() / 2
+    min_distance = radius[env._agent] + radius
+    min_distance = min_distance.unsqueeze(0).unsqueeze(0)
+    assert min_distance.shape == (1, 1, nv)
+
+    collision_tensor = distance > min_distance
+    assert collision_tensor.shape == (B, T, nv)
+    return states, collision_tensor
+
+def check_future_collisions_fast(env, actions):
+    """Checks whether `env._agent` would collide with other agents assuming `actions` as input.
+    
+    Vehicles are (over-)approximated by single circles.
+
+    Args:
+        env (gym.Env): current environment state
+        actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles
+    Returns:
+        feasible (torch.Tensor): tensor of shape (B,) indicating whether the respective action profiles are collision-free
+    """
+    _, collision_tensor = check_future_collisions_circle(env, actions)
+    return collision_tensor.all(-1).all(-1)
+
+def check_future_collisions_exact(env, actions):
+    """
+        Checks whether `env._agent` would collide with other agents assuming `actions` as input.
+
+    Args:
+        env (gym.Env): current environment state
+        actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles
+    Returns:
+        feasible (torch.Tensor): tensor of shape (B,) indicating whether the respective action profiles are collision-free
+    """
+    # First check with simple circle collision check
+    states, collision_tensor = check_future_collisions_circle(env, actions)
+    (B, T, nv, _) = states.shape
+    # For those that have colliding circles, check exactly
+    colliding_mask = ~collision_tensor
+    
+    ego_states = states[:, :, env._agent:env._agent+1, :].expand(states.shape)
+    assert ego_states.shape == states.shape
+
+    # get dimensions
+    lengths = env._env._lengths.expand(states.shape[:3])
+    widths = env._env._widths.expand(states.shape[:3])
+    ego_lengths = lengths[:, :, env._agent:env._agent+1].expand(lengths.shape)
+    ego_widths = widths[:, :, env._agent:env._agent+1].expand(widths.shape)
+    assert lengths.shape == widths.shape == ego_lengths.shape == ego_widths.shape == (B, T, nv)
+    
+    # For every collision instance between ego and other vehicle, check whether rectangles intersect
+    exact_collisions = torch.zeros_like(collision_tensor[colliding_mask])
+    for i, (ego_state, ego_length, ego_width, other_state, other_length, other_width) in enumerate(zip(
+        ego_states[colliding_mask], ego_lengths[colliding_mask], ego_widths[colliding_mask],
+        states[colliding_mask], lengths[colliding_mask], widths[colliding_mask]
+    )):
+        assert ego_state.shape == other_state.shape == (5,)
+        assert ego_length.shape == ego_width.shape == other_length.shape == other_width.shape == ()
+        p_ego = state_to_polygon(ego_state, ego_length, ego_width)
+        p_other = state_to_polygon(other_state, other_length, other_width)
+        exact_collisions[i] = p_ego.intersects(p_other)
+
+    collision_tensor[colliding_mask] = ~exact_collisions
+    return collision_tensor.all(-1).all(-1) 

src/util/render_env.py

@@ -0,0 +1,40 @@
+import stable_baselines3 as sb3
+import intersim 
+from src.gail.options import RenderOptions
+from tqdm import tqdm
+
+def render_env(model_name='gail_image_multiagent_nocollision', env='NRasterizedRoute', max_frames=600, options=False, options_list=None,
+    **env_kwargs):
+    """
+    Render a video from an model, agent, and environment
+    Args:
+        model_name (str): name of the model
+        environment (str): gym environment class to render environment on
+    """
+
+    model = sb3.PPO.load(model_name)
+    Env = intersim.envs.intersimple.__dict__[env]
+    
+    print(f'Rendering environment with \'{model_name}\' policy')
+    if not options:
+        env = Env(**env_kwargs)
+        obs = env.reset()
+        for i in tqdm(range(max_frames)):
+            action, _states = model.predict(obs)
+            obs, rewards, done, info = env.step(action)
+            env.render(mode='post')
+            if done:
+                break
+    else:
+        assert options_list, "No option list specified"
+        env = RenderOptions(Env(**env_kwargs), options=options_list)
+        with tqdm(total=max_frames) as pbar:
+            for i, s in enumerate(env.sample_ll(model)):
+                pbar.update(1)
+                if s['dones'] or i >= max_frames:
+                    break
+    env.close(filestr='render/'+model_name) 
+
+if __name__ == '__main__':
+    import fire
+    fire.Fire(render_env)

src/util/rollout.py

@@ -0,0 +1,206 @@
+# Borrowed heavily from https://github.com/HumanCompatibleAI/imitation/tree/master/src/imitation/data
+
+import pickle
+import numpy as np
+import logging
+import os
+import pathlib
+from typing import Optional, List, Dict
+
+def generate_trajectories(
+    policy,
+    env,
+    sample_until,
+    rng: Optional[np.random.RandomState] = None, # np.random to shuffle
+) -> List[dict]:
+    """Generate trajectory dictionaries from a policy and an environment.
+
+    Args:
+      policy: a stable_baselines3 policy or algorithm trained on the gym environment
+      env: The environment to interact with.
+      sample_until: A function determining the termination condition.
+          It takes a sequence of trajectories, and returns a bool.
+          Most users will want to use one of `min_episodes` or `min_timesteps`.
+      rng: used for shuffling trajectories.
+
+    Returns:
+      Sequence of trajectories, satisfying `sample_until`. 
+    """
+    trajectories = []
+    while not sample_until(trajectories):
+        
+        # sample a trajectory
+        ob, done = env.reset(), False
+        ob_list, rew_list, act_list, info_list = [], [], [], []
+        while not done:
+            act, _ = policy.predict(ob)
+            next_ob, rew, done, info = env.step(act) # ignore infos
+            ob_list.append(ob)
+            act_list.append(act)
+            rew_list.append(rew)
+            info_list.append(info)
+            ob = next_ob
+        ob_list.append(ob)
+
+        traj = {
+            'obs':np.stack(ob_list),
+            'acts':np.stack(act_list),
+            'rews':np.stack(rew_list),
+            'infos': info_list,
+            'terminal': True
+        }
+        trajectories.append(traj)
+    
+    # Shuffle trajectories
+    if rng:
+        rng.shuffle(trajectories)
+
+    # Sanity checks.
+    for trajectory in trajectories:
+        n_steps = len(trajectory['acts'])
+        # extra 1 for the end
+        exp_obs = (n_steps + 1,) + env.observation_space.shape
+        real_obs = trajectory['obs'].shape
+        assert real_obs == exp_obs, f"expected shape {exp_obs}, got {real_obs}"
+        exp_act = (n_steps,) + env.action_space.shape
+        real_act = trajectory['acts'].shape
+        assert real_act == exp_act, f"expected shape {exp_act}, got {real_act}"
+        exp_rew = (n_steps,)
+        real_rew = trajectory['rews'].shape
+        assert real_rew == exp_rew, f"expected shape {exp_rew}, got {real_rew}"
+
+    return trajectories
+
+
+def flatten_trajectories(trajectory_list: list) -> list:
+    """
+    Turn a list of trajectories into a (longer) list of transitions with appropriate fields
+
+    Args:
+        trajectory_list (list): list of trajectory dicts with keys:
+            obs (np.ndarray): (T, *O) tensor of all observations in T-step trajectory
+            acts (np.ndarray): (T-1, *A) tensor of all actions in T-step trajectory
+            infos (list[dict]): (T-1)-length list of all information dictionaries
+            terminal (bool): True if the trajectory ends at the last step
+            rews (np.ndarray): (T-1, 1) tensor of rewards along trajectory
+    Returns:
+        transition_list (list): list of all transition dicts with keys:
+            obs (np.ndarray): (*O) tensor of single-step observation
+            acts (np.ndarray): (*A) tensor of single-step action
+            infos (dict): single-step information dictionary
+            next_obs (np.ndarray): (*O) tensor of next observation
+            rews (np.ndarray): (1) tensor of single-step reward
+            dones (bool): whether state is terminal
+    """
+    transition_list = []
+    for traj in trajectory_list:
+        T = traj['obs'].shape[0]
+        if traj['infos']:
+            infos = traj['infos'] 
+        else: 
+            infos = [{}] * T
+        for i in range(T-1):
+            transition_list.append({
+               'obs': traj['obs'][i],
+               'acts': traj['acts'][i],
+               'next_obs': traj['obs'][i+1],
+               'rews': traj['rews'][i],
+               'dones': (i==T-2),
+               'infos': infos[i],
+            })
+    return transition_list
+
+
+def make_sample_until(min_timesteps: Optional[int]=None, min_episodes: Optional[int]=None):
+    """Returns a termination condition sampling for a number of timesteps and episodes.
+    
+    Args:
+        min_timesteps: Sampling will not stop until there are at least this many
+            timesteps.
+        min_episodes: Sampling will not stop until there are at least this many
+            episodes.
+
+    Returns:
+        A termination condition which given a list of trajectories returns true if the condition is met.
+
+    Raises:
+        ValueError if neither of n_timesteps and n_episodes are set, or if either are
+            non-positive.
+    """
+    if min_timesteps is None and min_episodes is None:
+        raise ValueError(
+            "At least one of min_timesteps and min_episodes needs to be non-None"
+        )
+
+    conditions = []
+    if min_timesteps is not None:
+        if min_timesteps <= 0:
+            raise ValueError(
+                f"min_timesteps={min_timesteps} if provided must be positive"
+            )
+
+        def timestep_cond(trajectories):
+            if len(trajectories) == 0:
+                return False
+            timesteps = sum(len(t['obs']) - 1 for t in trajectories)
+            return timesteps >= min_timesteps   
+        conditions.append(timestep_cond)
+
+    if min_episodes is not None:
+        if min_episodes <= 0:
+            raise ValueError(
+                f"min_episodes={min_episodes} if provided must be positive"
+            )
+        conditions.append(lambda trajectories: len(trajectories) >= min_episodes)
+
+    def sample_until(trajs: List[dict]) -> bool:
+        for cond in conditions:
+            if not cond(trajs):
+                return False
+        return True
+
+    return sample_until
+
+def rollout_and_save(
+    path: str,
+    policy,
+    env,
+    sample_until,
+    *,
+    exclude_infos: bool = True,
+    **kwargs,
+) -> None:
+    """Generate policy rollouts and save them to a pickled list of trajectories.
+
+    The `.infos` field of each Trajectory is set to `None` to save space.
+
+    Args:
+      path: Rollouts are saved to this path.
+      policy: a stable_baselines3 policy or algorithm trained on the gym environment
+      env: The environment to interact with.
+      sample_until: End condition for rollout sampling.
+      exclude_infos: If True, then exclude `infos` from pickle by setting
+        this field to None. Excluding `infos` can save a lot of space during
+        pickles.
+      **kwargs: Passed through to `generate_trajectories`.
+    """
+    trajs = generate_trajectories(policy, env, sample_until, **kwargs)
+    if exclude_infos:
+        [traj.update(infos=None) for traj in trajs]
+    save(path, trajs)
+
+def save(path: str, trajectories: List[dict]) -> None:
+    """Save a sequence of Trajectories to disk.
+
+    Args:
+        path: Trajectories are saved to this path.
+        trajectories: The trajectories to save.
+    """
+    p = pathlib.Path(path)
+    p.parent.mkdir(parents=True, exist_ok=True)
+    tmp_path = f"{path}.tmp"
+    with open(tmp_path, "wb") as f:
+        pickle.dump(trajectories, f)
+    # Ensure atomic write
+    os.replace(tmp_path, path)
+    logging.info(f"Dumped demonstrations to {path}.")

tests/test_discriminator.py

@@ -0,0 +1,45 @@
+from intersim.envs.intersimple import NRasterized
+from src.discriminator import CnnDiscriminator
+import torch
+
+def test_image_concatenation():
+    env = NRasterized()
+    disc = CnnDiscriminator(env)
+    s = torch.tensor(env.reset()).unsqueeze(0)
+    a = torch.tensor([[0.5]])
+    sa = disc._concatenate(s, a)
+
+    assert s.shape == (1, 5, 200, 200)
+    assert a.shape == (1, 1)
+    assert sa.shape == (1, 6, 200, 200)
+    assert torch.allclose(sa[:, :5], 1.0 * s)
+    assert (sa[:, 5] == a.unsqueeze(-1)).all()
+
+def test_image_concatenation3():
+    env = NRasterized()
+    disc = CnnDiscriminator(env)
+
+    s1 = env.reset()
+    a1 = 0.15
+    s2, _, _, _ = env.step(0.9)
+    a2 = 0.25
+    s3, _, _, _ = env.step(-0.9)
+    a3 = 0.35
+    
+    s = torch.stack([
+        torch.tensor(s1),
+        torch.tensor(s2),
+        torch.tensor(s3)
+    ], axis=0)
+    a = torch.tensor([
+        [a1],
+        [a2],
+        [a3],
+    ])
+    sa = disc._concatenate(s, a)
+
+    assert s.shape == (3, 5, 200, 200)
+    assert a.shape == (3, 1)
+    assert sa.shape == (3, 6, 200, 200)
+    assert torch.allclose(sa[:, :5], 1.0 * s)
+    assert (sa[:, 5] == a.unsqueeze(-1)).all()