Use ray tune

Merge different methods to sample the policy and collect transitions

.gitignore

@@ -1,3 +1,7 @@
+*.pkl
+*.pt
+*.zip
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

scratch/etienne/intersimple/data/generate.sh

@@ -1,3 +1,4 @@
-#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'

scratch/etienne/intersimple/gail/discriminator.py

@@ -1,5 +1,8 @@
 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."""
 
@@ -23,11 +26,16 @@ class CnnDiscriminator(torch.nn.Module):
             torch.nn.LazyLinear(1), # 512 -> 1
         )
     
-    def forward(self, state, action):
+    @staticmethod
+    def _concatenate(state, action):
         b, _, h, w = state.shape
         _, a = action.shape
-        act_layer = action.unsqueeze(-1).unsqueeze(-1).expand((b, a, h, w))
+        act = action.unsqueeze(-1).unsqueeze(-1).expand((b, a, h, w))
-        sa = torch.cat((act_layer, state), -3)
+        sa = torch.cat((state, act), -3)
+        return sa
+
+    def forward(self, state, action):
+        sa = self._concatenate(state, action)
         return self.cnn(sa).squeeze()
 
 class MlpDiscriminator(torch.nn.Module):

scratch/etienne/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/etienne/intersimple/gail_flat.py

@@ -10,15 +10,22 @@ from imitation.algorithms import adversarial, bc
 from imitation.data import rollout
 from imitation.util import logger
 
-from intersim.envs.intersimple import IntersimpleReward
+from intersimple.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_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.
@@ -36,20 +43,50 @@ 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.
+def training_function(config, checkpoint_dir=None):
-# 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/")
+
+    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=220,
+        expert_batch_size=config['expert_batch_size'],
-    #n_disc_updates_per_round=32,
+        n_disc_updates_per_round=config['n_disc_updates_per_round'],
         discrim_kwargs={'discrim_net': MlpDiscriminator()},
-    gen_algo=sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=4096),
+        gen_algo=generator,
     )
-gail_trainer.train(total_timesteps=80000)
+
-gail_trainer.gen_algo.save(model_name)
+    def callback(epoch):
+        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)
+        tune.report(progress=np.mean(episode_rewards))
+
+        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=400000, callback=callback)
+
+analysis = tune.run(
+    training_function,
+    config = {
+        'expert_batch_size': tune.randint(1, 220), #220,
+        'n_disc_updates_per_round': tune.randint(2, 100), #16,
+        'n_steps': tune.randint(1, 10000), #4096,
+    },
+    resources_per_trial={
+        'gpu': 1,
+    }
+)
+
+print('Best config', analysis.get_best_config(metric='progress', mode='max'))
 
 #del gail_trainer
 

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_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,297 @@
+# %%
+from gail.discriminator import CnnDiscriminator
+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
+
+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, 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
+
+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 sample(env, generator, discriminator, level: str):
+    """
+        Sample low-level (state, action, next_state) tuples for discriminator training or 
+        high-level (state, action, reward) tuples for generator training.
+    """
+    done = True
+    while True:
+        episode_start = False
+        if done:
+            s = env.reset()
+            m = available_actions(env)
+            done = False
+            episode_start = True
+
+        obs = {'obs': s, 'mask': m}
+        ch, value, log_prob = generator.policy.predict({
+            'obs': torch.tensor(s).unsqueeze(0).to(generator.policy.device),
+            'mask': torch.tensor(m).unsqueeze(0).to(generator.policy.device),
+        })
+        plan = list(map(float, generate_plan(env, ch)))
+
+        assert not done
+        assert plan
+        assert feasible(env, plan, ch), f'Infeasible hl action {ch}'
+
+        r = 0
+        discount = 1
+        while not done and plan and feasible(env, plan, ch):
+            a, plan = env._normalize(plan[0]), plan[1:]
+            if level == 'high':
+                r += discount * discriminator.discrim_net.discriminator(
+                    torch.tensor(s).unsqueeze(0).to(discriminator.discrim_net.device()),
+                    torch.tensor([[a]]).to(discriminator.discrim_net.device()),
+                )
+                discount *= env.discount
+
+            nexts, _, done, _ = env.step(a)
+            m = available_actions(env)
+
+            if level == 'low':
+                yield {
+                    'obs': s,
+                    'next_obs': nexts,
+                    'acts': np.array((a,)),
+                    'dones': np.array(done),
+                }
+            s = nexts
+
+        if level == 'high':
+            yield {
+                'obs': obs,
+                'option': ch,
+                'reward': r.detach(),
+                'episode_start': episode_start,
+                'value': value.detach(),
+                'log_prob': log_prob.detach(),
+                'done': done,
+            }
+
+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(sample(env, generator, None, 'low'), 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(sample(env, generator, discriminator, 'high'), num_samples+1))
+    
+    generator.rollout_buffer.reset()
+    for s in generator_samples[:-1]:
+        generator.rollout_buffer.add(
+            obs=s['obs'],
+            action=s['option'].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()
+
+class OptionsEnv(gym.Wrapper):
+    
+    def __init__(self, env):
+        super().__init__(env)
+        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 train(expert_data, epochs=10, expert_batch_size=32, generator_steps=2048, 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': 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 range(epochs):
+        train_discriminator(env, generator, discriminator, num_samples=expert_batch_size)
+        train_generator(env, 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_steps=200)
+
+    generator.save(model_name)
+
+    # %%
+    model = stable_baselines3.PPO.load(model_name)
+
+    env = NRasterized(**env_settings)
+
+    for transition in sample(env, generator, None, 'low'):
+        env.render()
+        if transition['dones']:
+            break
+
+    env.close(filestr='render/'+model_name)
+
+# %% Tests
+
+def test_ll_transitions_vs_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 = NRasterized(agent=51, width=36, height=36, m_per_px=2)
+
+    gen_transitions = list(itertools.islice(sample(
+        env=NRasterized(**env_settings),
+        generator=stable_baselines3.PPO(
+            OptionsCnnPolicy,
+            OptionsEnv(env),
+            verbose=1,
+        ),
+        discriminator=None,
+        level='low'
+    ), 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_hl_transitions():
+    pass

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)