making the expert demonstration processor go through all agents in order when producing a single (default) trajectory file, using a randomized agent environment in optionsgail, starting function to process and store all expert data

scratch/arec/intersimple/data/expert.py

@@ -0,0 +1,147 @@
+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
+    """
+    import pdb
+    pdb.set_trace()
+
+    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,8 @@
+#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'

scratch/arec/intersimple/data/process_all_experts.py

@@ -0,0 +1,40 @@
+import tqdm
+import expert
+import copy
+import sys, os
+
+def process_all_experts(env_args={}, policy_args={}):
+    """
+    Process all experts in the Interaction Dataset
+    For now, using NormalizedIntersimpleExpert with NRasterizedIncrementingAgent environment
+
+    Args:
+        env_args (dict): default environment kwargs
+        policy_args (dict): default policy kwargs
+    """
+
+    for loc in LOCATIONS:
+        for track in TRACKS:
+            
+            it_env_args = copy.deepcopy(env_args)
+            it_env_args.update({
+                'loc':loc,
+                'track':track,
+            })
+
+            it_path = 'newpathname'
+
+            expert.demonstrations(
+                expert='NormalizedIntersimpleExpert', 
+                env='NRasterizedIncrementingAgent', 
+                path=it_path, 
+                env_args=it_env_args, 
+                policy_args=policy_args,
+                )
+    
+
+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,386 @@
+# %%
+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):
+
+    def __init__(self, observation_space, *args, **kwargs):
+        super().__init__(observation_space['obs'], *args, **kwargs)
+
+    def _prior_distribution(self, s):
+        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):
+        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):
+        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):
+    
+    def __init__(self, env, *args, **kwargs):
+        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):
+        self.done = True
+        while True:
+            self.episode_start = False
+            if self.done:
+                self.s = self.env.reset()
+                self.m = available_actions(self.env)
+                self.done = False
+                self.episode_start = True
+
+            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.nextm = available_actions(self.env)
+
+                self._after_step()
+
+                self.s = self.nexts
+                self.m = self.nextm
+            
+            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.r = 0
+        self.steps = 0
+
+    def _after_step(self):
+        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 {
+            '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):
+        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 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 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
+    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):
+    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=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_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/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)