Scratch for horner scheme

Better way would probably be to rewrite flat_grad and reparam

.gitignore

@@ -2,6 +2,7 @@
 *.pt
 *.zip
 **/ray/*
+**/runs/*
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

evaluate_models.sh

@@ -0,0 +1,42 @@
+# eval_main inputs
+# locations: List[Tuple[int,int]]= [(0,0)],
+# method: str='expert', 
+# policy_file: str='', 
+# policy_kwargs: dict={},
+# env: str='NRasterizedRouteIncrementingAgent', 
+# env_kwargs: dict={},
+# seed: int=0
+
+# expert
+python -m src.eval_main
+
+# idm
+python -m src.eval_main --method=idm
+
+# behavior cloning
+python -m src.eval_main --method=bc --policy_file='checkpoints/bc-intersimple-setobs2.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True}' --seed=0
+python -m src.eval_main --method=bc --policy_file='checkpoints/bc-intersimple-setobs2.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True}' --seed=1
+python -m src.eval_main --method=bc --policy_file='checkpoints/bc-intersimple-setobs2.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True}' --seed=2
+python -m src.eval_main --method=bc --policy_file='checkpoints/bc-intersimple-setobs2.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True}' --seed=3
+python -m src.eval_main --method=bc --policy_file='checkpoints/bc-intersimple-setobs2.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True}' --seed=4
+python -m src.evaluation.utils load_and_average out/bc
+
+# GAIL
+python -m src.eval_main --method=gail --policy_file='checkpoints/gail-intersimple-setobs2-03-02-22.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True}' --seed=0
+python -m src.eval_main --method=gail --policy_file='checkpoints/gail-intersimple-setobs2-03-02-22.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True}' --seed=1
+python -m src.eval_main --method=gail --policy_file='checkpoints/gail-intersimple-setobs2-03-02-22.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True}' --seed=2
+python -m src.eval_main --method=gail --policy_file='checkpoints/gail-intersimple-setobs2-03-02-22.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True}' --seed=3
+python -m src.eval_main --method=gail --policy_file='checkpoints/gail-intersimple-setobs2-03-02-22.pt' --env='NormalizedContinuousEvalEnv' --env_kwargs='{stop_on_collision:True}' --seed=4
+python -m src.evaluation.utils load_and_average out/gail
+
+# options GAIL
+python -m src.eval_main --method=ogail --policy_file='checkpoints/gail-options-setobs2-Feb15_18-49-05.pt' --env='NormalizedOptionsEvalEnv' --env_kwargs='{stop_on_collision:True}'
+
+# options GAIL-PPO
+python -m src.eval_main --method=ogail-ppo --policy_file='checkpoints/gail-ppo-options-setobs2-Feb15_22-05-38.pt' --env='NormalizedOptionsEvalEnv' --env_kwargs='{stop_on_collision:True}'
+
+# SHAIL
+python -m src.eval_main --method=sgail --policy_file='checkpoints/sgail-options-setobs2.pt' --env='NormalizedSafeOptionsEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}'
+
+# SHAIL-PPO
+python -m src.eval_main --method=sgail-ppo --policy_file='checkpoints/sgail-ppo-options-setobs2-17-02-2022.pt' --env='NormalizedSafeOptionsEvalEnv' --env_kwargs='{stop_on_collision:True,max_episode_steps:1000}'

scratch/arec/intersimple/test_model.py

@@ -1,148 +0,0 @@
-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/ppo_speed_lidar.py

@@ -0,0 +1,50 @@
+# %%
+from stable_baselines3 import PPO
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+
+model_name = "ppo_speed_lidar"
+
+#def reward(state, action, info):
+#    speed = state[2].item()
+#    r = speed if speed < 10 else (10 - 5 * (speed - 10))
+#    return 0.1 * r
+
+env = IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+)
+
+# %%
+model = PPO(
+    "MlpPolicy", env,
+    learning_rate=1e-4,
+    verbose=1,
+    tensorboard_log='runs/'
+)
+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_lidar_random.py

@@ -0,0 +1,49 @@
+# %%
+from stable_baselines3 import PPO
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+
+model_name = "ppo_speed_lidar_random"
+
+#def reward(state, action, info):
+#    speed = state[2].item()
+#    r = speed if speed < 10 else (10 - 5 * (speed - 10))
+#    return 0.1 * r
+
+env = IntersimpleLidarFlatRandom(
+    n_rays=5,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+)
+
+# %%
+model = PPO(
+    "MlpPolicy", env,
+    learning_rate=1e-4,
+    verbose=1,
+    tensorboard_log='runs/'
+)
+model.learn(total_timesteps=1000000)
+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/trpo_speed_lidar.py

@@ -0,0 +1,52 @@
+# %%
+from sb3_contrib import TRPO
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+
+model_name = "trpo_speed_lidar"
+
+#def reward(state, action, info):
+#    speed = state[2].item()
+#    r = speed if speed < 10 else (10 - 5 * (speed - 10))
+#    return 0.1 * r
+
+env = IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+)
+
+# %%
+model = TRPO(
+    "MlpPolicy", env,
+    learning_rate=1e-4,
+    verbose=1,
+    tensorboard_log='runs/',
+    #use_sde=True,
+    #sde_sample_freq=4,
+)
+model.learn(total_timesteps=1000000)
+model.save(model_name)
+
+print('Done training.')
+
+del model # remove to demonstrate saving and loading
+
+# %%
+model = TRPO.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/trpo/experiments/bc-intersimple-setobs2.py

@@ -0,0 +1,78 @@
+# %%
+import torch
+from core.policy import SetPolicy
+from tqdm import tqdm
+
+expert_data = torch.load('intersimple-expert-data-setobs2.pt')
+states, actions, _, dones = expert_data
+
+policy = SetPolicy(actions.shape[-1])
+
+policy = policy.cuda()
+optim = torch.optim.Adam(policy.parameters(), lr=1e-4)
+states = states[~dones].cuda()
+actions = actions[~dones].cuda()
+
+for _ in tqdm(range(10000)):
+    optim.zero_grad()
+    loss = -policy.log_prob(policy(states), actions).mean()
+    loss.backward()
+    optim.step()
+
+    print('Loss', loss)
+
+torch.save(policy.state_dict(), 'bc-intersimple-setobs2.pt')
+
+# %%
+import numpy as np
+from core.policy import SetPolicy
+from util.wrappers import Setobs, TransformObservation, CollisionPenaltyWrapper
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+
+policy = SetPolicy(actions.shape[-1])
+policy.load_state_dict(torch.load('bc-intersimple-setobs2.pt'))
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+env = Setobs(
+    TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
+        n_rays=5,
+        reward=functools.partial(
+            speed_reward,
+            collision_penalty=0
+        ),
+        stop_on_collision=False,
+    ), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
+)
+
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    #action, _ = policy.predict(torch.tensor(obs))
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action)
+    env.render(mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()
+
+# %%

scratch/etienne/trpo/experiments/gail-intersimple-minobs.py

@@ -0,0 +1,74 @@
+import gym
+from core.gail import gail, Buffer
+from core.value import Value
+from core.policy import Policy
+from core.discriminator import Discriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, Minobs
+import numpy as np
+from gym.wrappers import TransformObservation
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [Minobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(30)]
+env_fn = lambda i: envs[i]
+
+policy = Policy(env_fn(0).action_space.shape[0])
+
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4)
+
+discriminator = Discriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3)
+
+expert_data = torch.load('intersimple-expert-data-minobs.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=100,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=4000,
+    rollout_episodes=30,
+    rollout_steps=200,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+)
+
+torch.save(policy.state_dict(), 'gail-intersimple-minobs.pt')

scratch/etienne/trpo/experiments/gail-intersimple-minobs2.py

@@ -0,0 +1,100 @@
+# %%
+import gym
+from core.gail import gail, Buffer
+from core.value import Value
+from core.policy import Policy
+from core.discriminator import Discriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, Minobs
+import numpy as np
+from gym.wrappers import TransformObservation
+from core.reparam_module import ReparamPolicy
+from torch.utils.tensorboard import SummaryWriter
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [Minobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
+    n_rays=5,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(50)]
+env_fn = lambda i: envs[i]
+
+policy = Policy(env_fn(0).action_space.shape[0])
+
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
+
+discriminator = Discriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-5)
+
+expert_data = torch.load('intersimple-expert-data-minobs2.pt')
+expert_data = Buffer(*expert_data)
+
+# %%
+value, policy = gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=500,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=800,
+    rollout_episodes=50,
+    rollout_steps=200,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+    logger=SummaryWriter(comment='minobs2'),
+)
+
+torch.save(policy.state_dict(), 'gail-intersimple-minobs2.pt')
+
+# %%
+policy = Policy(env_fn(0).action_space.shape[0])
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy = ReparamPolicy(policy)
+policy.load_state_dict(torch.load('gail-intersimple-minobs2.pt'))
+
+env = env_fn(0)
+env.random_skip = False
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    #action, _ = policy.predict(torch.tensor(obs))
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action)
+    env.render(mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()
+
+# %%

scratch/etienne/trpo/experiments/gail-intersimple-normobs.py

@@ -0,0 +1,74 @@
+import gym
+from core.gail import gail, Buffer
+from core.value import Value
+from core.policy import Policy
+from core.discriminator import Discriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper
+import numpy as np
+from gym.wrappers import TransformObservation
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)) for _ in range(30)]
+env_fn = lambda i: envs[i]
+
+policy = Policy(env_fn(0).action_space.shape[0])
+
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4)
+
+discriminator = Discriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-4)
+
+expert_data = torch.load('intersimple-expert-data-normobs.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=10,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=4000,
+    rollout_episodes=30,
+    rollout_steps=100,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+)
+
+torch.save(policy.state_dict(), 'gail-intersimple-normobs.pt')

scratch/etienne/trpo/experiments/gail-intersimple-setobs.py

@@ -0,0 +1,74 @@
+import gym
+from core.gail import gail, Buffer
+from core.value import SetValue
+from core.policy import SetPolicy
+from core.discriminator import DeepsetDiscriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, Setobs
+import numpy as np
+from gym.wrappers import TransformObservation
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [Setobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(30)]
+env_fn = lambda i: envs[i]
+
+policy = SetPolicy(env_fn(0).action_space.shape[0])
+
+value = SetValue()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4)
+
+discriminator = DeepsetDiscriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3)
+
+expert_data = torch.load('intersimple-expert-data-setobs.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=500,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=4000,
+    rollout_episodes=30,
+    rollout_steps=200,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+)
+
+torch.save(policy.state_dict(), 'gail-intersimple-setobs.pt')

scratch/etienne/trpo/experiments/gail-intersimple-setobs2-recurrent.py

@@ -0,0 +1,97 @@
+# %%
+import gym
+from core.gail import gail, Buffer
+from core.value import SetValue
+from core.policy import SetPolicy
+from core.discriminator import RecurrentDiscriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, Setobs
+import numpy as np
+from gym.wrappers import TransformObservation
+from core.reparam_module import ReparamPolicy
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [Setobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
+    n_rays=5,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(50)]
+env_fn = lambda i: envs[i]
+
+policy = SetPolicy(env_fn(0).action_space.shape[0])
+
+value = SetValue()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
+
+discriminator = RecurrentDiscriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-5)
+
+expert_data = torch.load('intersimple-expert-data-setobs2.pt')
+expert_data = Buffer(*expert_data)
+
+# %%
+value, policy = gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=500,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=800,
+    rollout_episodes=50,
+    rollout_steps=200,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+)
+
+torch.save(policy.state_dict(), 'gail-intersimple-setobs-recurrent.pt')
+
+# %%
+policy = SetPolicy(env_fn(0).action_space.shape[0])
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy = ReparamPolicy(policy)
+policy.load_state_dict(torch.load('gail-intersimple-setobs-recurrent.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    #action, _ = policy.predict(torch.tensor(obs))
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action)
+    env.render(mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()
+
+# %%

scratch/etienne/trpo/experiments/gail-intersimple-setobs2.py

@@ -0,0 +1,101 @@
+# %%
+import gym
+from core.gail import gail, Buffer
+from core.value import SetValue
+from core.policy import SetPolicy
+from core.discriminator import DeepsetDiscriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, Setobs
+import numpy as np
+from gym.wrappers import TransformObservation
+from core.reparam_module import ReparamPolicy
+from torch.utils.tensorboard import SummaryWriter
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [Setobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
+    n_rays=5,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+    random_skip=True,
+), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(50)]
+env_fn = lambda i: envs[i]
+
+policy = SetPolicy(env_fn(0).action_space.shape[0])
+
+value = SetValue()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
+
+discriminator = DeepsetDiscriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-5)
+
+expert_data = torch.load('intersimple-expert-data-setobs2.pt')
+expert_data = Buffer(*expert_data)
+
+# %%
+value, policy = gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=500,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=800,
+    rollout_episodes=50,
+    rollout_steps=200,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+    logger=SummaryWriter(comment='setobs2-batchaug'),
+)
+
+torch.save(policy.state_dict(), 'gail-intersimple-setobs2.pt')
+
+# %%
+policy = SetPolicy(env_fn(0).action_space.shape[0])
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy = ReparamPolicy(policy)
+policy.load_state_dict(torch.load('gail-intersimple-setobs2.pt'))
+
+env = env_fn(0)
+env.random_skip = False
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    #action, _ = policy.predict(torch.tensor(obs))
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action)
+    env.render(mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()
+
+# %%

scratch/etienne/trpo/experiments/gail-intersimple.py

@@ -0,0 +1,54 @@
+import gym
+from core.gail import gail, Buffer
+from core.value import Value
+from core.policy import Policy
+from core.discriminator import Discriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper
+
+envs = [CollisionPenaltyWrapper(IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100) for _ in range(30)]
+env_fn = lambda i: envs[i]
+
+policy = Policy(env_fn(0).action_space.shape[0])
+
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4)
+
+discriminator = Discriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-4)
+
+expert_data = torch.load('intersimple-expert-data.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=10,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=4000,
+    rollout_episodes=30,
+    rollout_steps=100,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+)
+
+torch.save(policy.state_dict(), 'gail-intersimple.pt')

scratch/etienne/trpo/experiments/gail-options-minobs.py

@@ -0,0 +1,97 @@
+import gym
+from options.options import gail
+from core.gail import Buffer
+from core.value import Value
+from core.policy import DiscretePolicy
+from core.discriminator import Discriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, TransformObservation, Minobs
+import numpy as np
+from options.options import OptionsEnv
+from torch.utils.tensorboard import SummaryWriter
+from core.reparam_module import ReparamPolicy
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [OptionsEnv(Minobs(
+    TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
+        n_rays=5,
+        agent=51,
+        reward=functools.partial(
+            speed_reward,
+            collision_penalty=0
+        ),
+        stop_on_collision=False,
+    ), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
+), options=[(0, 5), (1, 5), (2, 5), (4, 5), (8, 5)]) for _ in range(60)]
+
+env_fn = lambda i: envs[i]
+policy = DiscretePolicy(env_fn(0).action_space.n)
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4)
+
+discriminator = Discriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3)
+
+expert_data = torch.load('intersimple-expert-data-minobs.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=100,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=50,
+    rollout_episodes=60,
+    rollout_steps=60,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+    logger=SummaryWriter(comment='-options-minobs'),
+)
+
+torch.save(policy.state_dict(), 'gail-options-minobs.pt')
+
+# %%
+policy = DiscretePolicy(env_fn(0).action_space.n)
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy = ReparamPolicy(policy)
+policy.load_state_dict(torch.load('gail-options-minobs.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    #action, _ = policy.predict(torch.tensor(obs))
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action, render_mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()

scratch/etienne/trpo/experiments/gail-options-setobs.py

@@ -0,0 +1,97 @@
+import gym
+from options.options import gail
+from core.gail import Buffer
+from core.value import SetValue
+from core.policy import SetDiscretePolicy
+from core.discriminator import DeepsetDiscriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, TransformObservation, Setobs
+import numpy as np
+from options.options import OptionsEnv
+from torch.utils.tensorboard import SummaryWriter
+from core.reparam_module import ReparamPolicy
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [OptionsEnv(Setobs(
+    TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
+        n_rays=5,
+        agent=51,
+        reward=functools.partial(
+            speed_reward,
+            collision_penalty=0
+        ),
+        stop_on_collision=False,
+    ), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
+), options=[(0, 5), (1, 5), (2, 5), (4, 5), (8, 5)]) for _ in range(60)]
+
+env_fn = lambda i: envs[i]
+policy = SetDiscretePolicy(env_fn(0).action_space.n)
+value = SetValue()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4)
+
+discriminator = DeepsetDiscriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-3)
+
+expert_data = torch.load('intersimple-expert-data-setobs.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=100,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=150,
+    rollout_episodes=60,
+    rollout_steps=60,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+    logger=SummaryWriter(comment='gail-options-setobs'),
+)
+
+torch.save(policy.state_dict(), 'gail-options-setobs.pt')
+
+# %%
+policy = SetDiscretePolicy(env_fn(0).action_space.n)
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy = ReparamPolicy(policy)
+policy.load_state_dict(torch.load('gail-options-setobs.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    #action, _ = policy.predict(torch.tensor(obs))
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action, render_mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()

scratch/etienne/trpo/experiments/gail-options-setobs2.py

@@ -0,0 +1,104 @@
+# %%
+import gym
+from options.options import gail
+from core.gail import Buffer
+from core.value import SetValue
+from core.policy import SetDiscretePolicy
+from core.discriminator import DeepsetDiscriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, TransformObservation, Setobs
+import numpy as np
+from options.options import OptionsEnv
+from torch.utils.tensorboard import SummaryWriter
+from core.reparam_module import ReparamPolicy
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [OptionsEnv(Setobs(
+    TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
+        n_rays=5,
+        reward=functools.partial(
+            speed_reward,
+            collision_penalty=0
+        ),
+        stop_on_collision=False,
+    ), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
+), options=[(0, 5), (1, 5), (2, 5), (4, 5), (6, 5), (8, 5)]) for _ in range(60)]
+
+env_fn = lambda i: envs[i]
+policy = SetDiscretePolicy(env_fn(0).action_space.n)
+value = SetValue()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4)
+
+discriminator = DeepsetDiscriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-4)
+
+expert_data = torch.load('intersimple-expert-data-setobs2.pt')
+expert_data = Buffer(*expert_data)
+
+# %%
+def callback(epoch, value, policy):
+    if not epoch % 10:
+        torch.save(policy.state_dict(), f'gail-options-setobs2-{epoch}.pt')
+        torch.save(value.state_dict(), f'gail-options-setobs2-value-{epoch}.pt')
+
+value, policy = gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=100,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=300,
+    rollout_episodes=60,
+    rollout_steps=60,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+    logger=SummaryWriter(comment='gail-options-setobs2'),
+    callback=callback,
+)
+
+torch.save(policy.state_dict(), 'gail-options-setobs2.pt')
+
+# %%
+policy = SetDiscretePolicy(env_fn(0).action_space.n)
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy = ReparamPolicy(policy)
+policy.load_state_dict(torch.load('gail-options-setobs2.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    #action, _ = policy.predict(torch.tensor(obs))
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action, render_mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()

scratch/etienne/trpo/experiments/gail-pendulum.py

@@ -0,0 +1,39 @@
+import gym
+from core.gail import gail, Buffer
+from core.value import Value
+from core.policy import Policy
+from core.discriminator import Discriminator
+import torch.optim
+
+env_fn = lambda _: gym.make('Pendulum-v0')
+policy = Policy(env_fn(0).action_space.shape[0])
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
+discriminator = Discriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3)
+
+expert_data = torch.load('trpo-pendulum-expert-data.pt')
+expert_data = Buffer(*expert_data)
+
+gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=10,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=100,
+    rollout_episodes=20,
+    rollout_steps=250,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+)
+
+
+torch.save(policy.state_dict(), 'gail-pendulum.pt')

scratch/etienne/trpo/experiments/gail-ppo-intersimple-minobs.py

@@ -0,0 +1,75 @@
+import gym
+from core.gail import gail_ppo, Buffer
+from core.value import Value
+from core.policy import Policy
+from core.discriminator import Discriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, Minobs
+import numpy as np
+from gym.wrappers import TransformObservation
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [Minobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(30)]
+env_fn = lambda i: envs[i]
+
+policy = Policy(env_fn(0).action_space.shape[0])
+pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
+
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3, weight_decay=1e-3)
+
+discriminator = Discriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3)
+
+expert_data = torch.load('intersimple-expert-data-minobs.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail_ppo(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=100,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=4000,
+    rollout_episodes=30,
+    rollout_steps=100,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+)
+
+torch.save(policy.state_dict(), 'gail-ppo-intersimple-minobs.pt')

scratch/etienne/trpo/experiments/gail-ppo-intersimple-normobs.py

@@ -0,0 +1,75 @@
+import gym
+from core.gail import gail_ppo, Buffer
+from core.value import Value
+from core.policy import Policy
+from core.discriminator import Discriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper
+import numpy as np
+from gym.wrappers import TransformObservation
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)) for _ in range(30)]
+env_fn = lambda i: envs[i]
+
+policy = Policy(env_fn(0).action_space.shape[0])
+pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
+
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
+
+discriminator = Discriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3)
+
+expert_data = torch.load('intersimple-expert-data-normobs.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail_ppo(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=100,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=4000,
+    rollout_episodes=30,
+    rollout_steps=100,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+)
+
+torch.save(policy.state_dict(), 'gail-ppo-intersimple-normobs.pt')

scratch/etienne/trpo/experiments/gail-ppo-intersimple-setobs2.py

@@ -0,0 +1,102 @@
+# %%
+import gym
+from core.gail import gail_ppo, Buffer
+from core.value import SetValue
+from core.policy import SetPolicy
+from core.discriminator import DeepsetDiscriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, Setobs
+import numpy as np
+from gym.wrappers import TransformObservation
+from core.reparam_module import ReparamPolicy
+from torch.utils.tensorboard import SummaryWriter
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [Setobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
+    n_rays=5,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+    random_skip=True,
+), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(50)]
+env_fn = lambda i: envs[i]
+
+policy = SetPolicy(env_fn(0).action_space.shape[0])
+pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
+
+value = SetValue()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
+
+discriminator = DeepsetDiscriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-5)
+
+expert_data = torch.load('intersimple-expert-data-setobs2.pt')
+expert_data = Buffer(*expert_data)
+
+# %%
+value, policy = gail_ppo(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=500,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=800,
+    rollout_episodes=50,
+    rollout_steps=200,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+    logger=SummaryWriter(comment='-ppo-setobs2'),
+)
+
+torch.save(policy.state_dict(), 'gail-ppo-intersimple-setobs2.pt')
+
+# %%
+policy = SetPolicy(env_fn(0).action_space.shape[0])
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy = ReparamPolicy(policy)
+policy.load_state_dict(torch.load('gail-ppo-intersimple-setobs2.pt'))
+
+env = env_fn(0)
+env.random_skip = False
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    #action, _ = policy.predict(torch.tensor(obs))
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action)
+    env.render(mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()
+
+# %%

scratch/etienne/trpo/experiments/gail-ppo-intersimple.py

@@ -0,0 +1,55 @@
+import gym
+from core.gail import gail_ppo, Buffer
+from core.value import Value
+from core.policy import Policy
+from core.discriminator import Discriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper
+
+envs = [CollisionPenaltyWrapper(IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100) for _ in range(30)]
+env_fn = lambda i: envs[i]
+
+policy = Policy(env_fn(0).action_space.shape[0])
+pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
+
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
+
+discriminator = Discriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=3e-4)
+
+expert_data = torch.load('intersimple-expert-data.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail_ppo(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=10,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=4000,
+    rollout_episodes=30,
+    rollout_steps=100,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+)
+
+torch.save(policy.state_dict(), 'gail-ppo-intersimple.pt')

scratch/etienne/trpo/experiments/gail-ppo-options-minobs.py

@@ -0,0 +1,96 @@
+import gym
+from options.options import gail_ppo, Buffer
+from core.value import Value
+from core.policy import DiscretePolicy
+from core.discriminator import Discriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, TransformObservation, Minobs
+import numpy as np
+from options.options import OptionsEnv
+from torch.utils.tensorboard import SummaryWriter
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [OptionsEnv(Minobs(
+    TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
+        n_rays=5,
+        agent=51,
+        reward=functools.partial(
+            speed_reward,
+            collision_penalty=0
+        ),
+        stop_on_collision=False,
+    ), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
+), options=[(0, 5), (1, 5), (2, 5), (4, 5), (8, 5)]) for _ in range(60)]
+
+env_fn = lambda i: envs[i]
+
+policy = DiscretePolicy(env_fn(0).action_space.n)
+pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
+
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
+
+discriminator = Discriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3)
+
+expert_data = torch.load('intersimple-expert-data-minobs.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail_ppo(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=100,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=50,
+    rollout_episodes=60,
+    rollout_steps=60,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+    logger=SummaryWriter(comment='gail-ppo-options-minobs'),
+)
+
+torch.save(policy.state_dict(), 'gail-ppo-options-minobs.pt')
+
+# %%
+policy = DiscretePolicy(env_fn(0).action_space.n)
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy.load_state_dict(torch.load('gail-ppo-options-minobs.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action, render_mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()

scratch/etienne/trpo/experiments/gail-ppo-options-setobs.py

@@ -0,0 +1,96 @@
+import gym
+from options.options import gail_ppo, Buffer
+from core.value import SetValue
+from core.policy import SetDiscretePolicy
+from core.discriminator import DeepsetDiscriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, TransformObservation, Setobs
+import numpy as np
+from options.options import OptionsEnv
+from torch.utils.tensorboard import SummaryWriter
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [OptionsEnv(Setobs(
+    TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
+        n_rays=5,
+        agent=51,
+        reward=functools.partial(
+            speed_reward,
+            collision_penalty=0
+        ),
+        stop_on_collision=False,
+    ), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
+), options=[(0, 5), (1, 5), (2, 5), (4, 5), (8, 5)]) for _ in range(60)]
+
+env_fn = lambda i: envs[i]
+
+policy = SetDiscretePolicy(env_fn(0).action_space.n)
+pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
+
+value = SetValue()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
+
+discriminator = DeepsetDiscriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-3)
+
+expert_data = torch.load('intersimple-expert-data-setobs.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail_ppo(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=100,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=150,
+    rollout_episodes=60,
+    rollout_steps=60,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+    logger=SummaryWriter(comment='gail-ppo-options-setobs'),
+)
+
+torch.save(policy.state_dict(), 'gail-ppo-options-setobs.pt')
+
+# %%
+policy = SetDiscretePolicy(env_fn(0).action_space.n)
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy.load_state_dict(torch.load('gail-ppo-options-setobs.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action, render_mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()

scratch/etienne/trpo/experiments/gail-ppo-options-setobs2.py

@@ -0,0 +1,103 @@
+# %%
+import gym
+from options.options import gail_ppo, Buffer
+from core.value import SetValue
+from core.policy import SetDiscretePolicy
+from core.discriminator import DeepsetDiscriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, TransformObservation, Setobs
+import numpy as np
+from options.options import OptionsEnv
+from torch.utils.tensorboard import SummaryWriter
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [OptionsEnv(Setobs(
+    TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
+        n_rays=5,
+        reward=functools.partial(
+            speed_reward,
+            collision_penalty=0
+        ),
+        stop_on_collision=False,
+    ), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
+), options=[(0, 5), (1, 5), (2, 5), (4, 5), (6, 5), (8, 5)]) for _ in range(60)]
+
+env_fn = lambda i: envs[i]
+
+policy = SetDiscretePolicy(env_fn(0).action_space.n)
+pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
+
+value = SetValue()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
+
+discriminator = DeepsetDiscriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-4)
+
+expert_data = torch.load('intersimple-expert-data-setobs2.pt')
+expert_data = Buffer(*expert_data)
+
+# %%
+def callback(epoch, value, policy):
+    if not epoch % 10:
+        torch.save(policy.state_dict(), f'gail-ppo-options-setobs2-{epoch}.pt')
+        torch.save(value.state_dict(), f'gail-ppo-options-setobs2-value-{epoch}.pt')
+
+value, policy = gail_ppo(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=100,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=200,
+    rollout_episodes=60,
+    rollout_steps=60,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+    logger=SummaryWriter(comment='gail-ppo-options-setobs2'),
+    callback=callback,
+)
+
+torch.save(policy.state_dict(), 'gail-ppo-options-setobs2.pt')
+
+# %%
+policy = SetDiscretePolicy(env_fn(0).action_space.n)
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy.load_state_dict(torch.load('gail-ppo-options-setobs2.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action, render_mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()

scratch/etienne/trpo/experiments/intersimple-expert-rollout-minobs.py

@@ -0,0 +1,54 @@
+import torch
+import functools
+from core.sampling import rollout_sb3
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+from intersim.expert import NormalizedIntersimpleExpert
+from util.wrappers import CollisionPenaltyWrapper, Minobs
+import numpy as np
+from gym.wrappers import TransformObservation
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+env = IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+)
+
+policy = NormalizedIntersimpleExpert(env, mu=0.001)
+
+env = Minobs(TransformObservation(
+    CollisionPenaltyWrapper(
+        env,
+        collision_distance=6, collision_penalty=100
+    ), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)
+))
+expert_data = rollout_sb3(env, policy, n_episodes=64, max_steps_per_episode=200)
+
+states, actions, rewards, dones = expert_data
+print(f'Expert mean episode length {(~dones).sum() / states.shape[0]}')
+print(f'Expert mean reward per episode {rewards[~dones].sum() / states.shape[0]}')
+print(f'Observation mean', states[~dones].mean(0))
+print(f'Observation std', states[~dones].std(0))
+
+torch.save(expert_data, 'intersimple-expert-data-minobs.pt')

scratch/etienne/trpo/experiments/intersimple-expert-rollout-minobs2.py

@@ -0,0 +1,53 @@
+import torch
+import functools
+from core.sampling import rollout_sb3
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+from intersim.expert import NormalizedIntersimpleExpert
+from util.wrappers import CollisionPenaltyWrapper, Minobs
+import numpy as np
+from gym.wrappers import TransformObservation
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+env = IntersimpleLidarFlatRandom(
+    n_rays=5,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+)
+
+policy = NormalizedIntersimpleExpert(env, mu=0.001)
+
+env = Minobs(TransformObservation(
+    CollisionPenaltyWrapper(
+        env,
+        collision_distance=6, collision_penalty=100
+    ), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)
+))
+expert_data = rollout_sb3(env, policy, n_episodes=2048, max_steps_per_episode=200)
+
+states, actions, rewards, dones = expert_data
+print(f'Expert mean episode length {(~dones).sum() / states.shape[0]}')
+print(f'Expert mean reward per episode {rewards[~dones].sum() / states.shape[0]}')
+print(f'Observation mean', states[~dones].mean(0))
+print(f'Observation std', states[~dones].std(0))
+
+torch.save(expert_data, 'intersimple-expert-data-minobs2.pt')

scratch/etienne/trpo/experiments/intersimple-expert-rollout-normobs.py

@@ -0,0 +1,54 @@
+import torch
+import functools
+from core.sampling import rollout_sb3
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+from intersim.expert import NormalizedIntersimpleExpert
+from util.wrappers import CollisionPenaltyWrapper
+import numpy as np
+from gym.wrappers import TransformObservation
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+env = IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+)
+
+policy = NormalizedIntersimpleExpert(env, mu=0.001)
+
+env = TransformObservation(
+    CollisionPenaltyWrapper(
+        env,
+        collision_distance=6, collision_penalty=100
+    ), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)
+)
+expert_data = rollout_sb3(env, policy, n_episodes=64, max_steps_per_episode=200)
+
+states, actions, rewards, dones = expert_data
+print(f'Expert mean episode length {(~dones).sum() / states.shape[0]}')
+print(f'Expert mean reward per episode {rewards[~dones].sum() / states.shape[0]}')
+print(f'Observation mean', states[~dones].mean(0))
+print(f'Observation std', states[~dones].std(0))
+
+torch.save(expert_data, 'intersimple-expert-data-normobs.pt')

scratch/etienne/trpo/experiments/intersimple-expert-rollout-setobs.py

@@ -0,0 +1,54 @@
+import torch
+import functools
+from core.sampling import rollout_sb3
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+from intersim.expert import NormalizedIntersimpleExpert
+from util.wrappers import CollisionPenaltyWrapper, Setobs
+import numpy as np
+from gym.wrappers import TransformObservation
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+env = IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+)
+
+policy = NormalizedIntersimpleExpert(env, mu=0.001)
+
+env = Setobs(TransformObservation(
+    CollisionPenaltyWrapper(
+        env,
+        collision_distance=6, collision_penalty=100
+    ), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)
+))
+expert_data = rollout_sb3(env, policy, n_episodes=64, max_steps_per_episode=200)
+
+states, actions, rewards, dones = expert_data
+print(f'Expert mean episode length {(~dones).sum() / states.shape[0]}')
+print(f'Expert mean reward per episode {rewards[~dones].sum() / states.shape[0]}')
+print(f'Observation mean', states[~dones].mean(0))
+print(f'Observation std', states[~dones].std(0))
+
+torch.save(expert_data, 'intersimple-expert-data-setobs.pt')

scratch/etienne/trpo/experiments/intersimple-expert-rollout-setobs2.py

@@ -0,0 +1,53 @@
+import torch
+import functools
+from core.sampling import rollout_sb3
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+from intersim.expert import NormalizedIntersimpleExpert
+from util.wrappers import CollisionPenaltyWrapper, Setobs
+import numpy as np
+from gym.wrappers import TransformObservation
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+env = IntersimpleLidarFlatRandom(
+    n_rays=5,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+)
+
+policy = NormalizedIntersimpleExpert(env, mu=0.001)
+
+env = Setobs(TransformObservation(
+    CollisionPenaltyWrapper(
+        env,
+        collision_distance=6, collision_penalty=100
+    ), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)
+))
+expert_data = rollout_sb3(env, policy, n_episodes=2048, max_steps_per_episode=200)
+
+states, actions, rewards, dones = expert_data
+print(f'Expert mean episode length {(~dones).sum() / states.shape[0]}')
+print(f'Expert mean reward per episode {rewards[~dones].sum() / states.shape[0]}')
+print(f'Observation mean', states[~dones].mean(0))
+print(f'Observation std', states[~dones].std(0))
+
+torch.save(expert_data, 'intersimple-expert-data-setobs2.pt')

scratch/etienne/trpo/experiments/intersimple-expert-rollout.py

@@ -0,0 +1,25 @@
+import torch
+import functools
+from core.sampling import rollout_sb3
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+from intersim.expert import NormalizedIntersimpleExpert
+from util.wrappers import CollisionPenaltyWrapper
+
+env = CollisionPenaltyWrapper(IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+), collision_distance=6, collision_penalty=100)
+policy = NormalizedIntersimpleExpert(env.env, mu=0.001)
+
+expert_data = rollout_sb3(env, policy, n_episodes=64, max_steps_per_episode=200)
+
+states, actions, rewards, dones = expert_data
+print(f'Expert mean episode length {(~dones).sum() / states.shape[0]}')
+print(f'Expert mean reward per episode {rewards[~dones].sum() / states.shape[0]}')
+
+torch.save(expert_data, 'intersimple-expert-data.pt')

scratch/etienne/trpo/experiments/ppo-intersimple-minobs.py

@@ -0,0 +1,63 @@
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+
+from core.ppo import ppo
+from core.value import Value
+from core.policy import Policy
+import torch.optim
+import numpy as np
+from gym.wrappers import TransformObservation
+
+from util.wrappers import Minobs
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [Minobs(TransformObservation(IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(30)]
+
+env_fn = lambda i: envs[i]
+policy = Policy(env_fn(0).action_space.shape[0])
+value = Value()
+pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
+
+value, policy = ppo(
+    env_fn=env_fn,
+    value=value,
+    policy=policy,
+    epochs=50,
+    rollout_episodes=30,
+    rollout_steps=200,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+    v_opt=v_opt,
+    v_iters=1000,
+)
+
+torch.save(policy.state_dict(), 'ppo-intersimple.pt')

scratch/etienne/trpo/experiments/ppo-intersimple-minobs2.py

@@ -0,0 +1,62 @@
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+
+from core.ppo import ppo
+from core.value import Value
+from core.policy import Policy
+import torch.optim
+import numpy as np
+from gym.wrappers import TransformObservation
+
+from util.wrappers import Minobs
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [Minobs(TransformObservation(IntersimpleLidarFlatRandom(
+    n_rays=5,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=1000
+    ),
+), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(30)]
+
+env_fn = lambda i: envs[i]
+policy = Policy(env_fn(0).action_space.shape[0])
+value = Value()
+pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3, weight_decay=1e-3)
+
+value, policy = ppo(
+    env_fn=env_fn,
+    value=value,
+    policy=policy,
+    epochs=50,
+    rollout_episodes=30,
+    rollout_steps=100,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+    v_opt=v_opt,
+    v_iters=1000,
+)
+
+torch.save(policy.state_dict(), 'ppo-intersimple.pt')

scratch/etienne/trpo/experiments/ppo-intersimple-normobs.py

@@ -0,0 +1,61 @@
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+
+from core.ppo import ppo
+from core.value import Value
+from core.policy import Policy
+import torch.optim
+import numpy as np
+from gym.wrappers import TransformObservation
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [TransformObservation(IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=10
+    ),
+), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)) for _ in range(30)]
+
+env_fn = lambda i: envs[i]
+policy = Policy(env_fn(0).action_space.shape[0])
+value = Value()
+pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
+
+value, policy = ppo(
+    env_fn=env_fn,
+    value=value,
+    policy=policy,
+    epochs=50,
+    rollout_episodes=30,
+    rollout_steps=100,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+    v_opt=v_opt,
+    v_iters=1000,
+)
+
+torch.save(policy.state_dict(), 'ppo-intersimple.pt')

scratch/etienne/trpo/experiments/ppo-intersimple.py

@@ -0,0 +1,41 @@
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+
+from core.ppo import ppo
+from core.value import Value
+from core.policy import Policy
+import torch.optim
+
+envs = [IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=10
+    ),
+) for _ in range(30)]
+
+env_fn = lambda i: envs[i]
+policy = Policy(env_fn(0).action_space.shape[0])
+value = Value()
+pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
+
+value, policy = ppo(
+    env_fn=env_fn,
+    value=value,
+    policy=policy,
+    epochs=50,
+    rollout_episodes=30,
+    rollout_steps=100,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+    v_opt=v_opt,
+    v_iters=1000,
+)
+
+torch.save(policy.state_dict(), 'ppo-intersimple.pt')

scratch/etienne/trpo/experiments/ppo-options-minobs.py

@@ -0,0 +1,89 @@
+# %%
+import gym
+from core.sampling import rollout
+from core.ppo import ppo
+from core.value import Value
+from core.policy import DiscretePolicy
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+import numpy as np
+from util.wrappers import CollisionPenaltyWrapper, TransformObservation
+
+from util.wrappers import Minobs
+from options.options import OptionsEnv
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [OptionsEnv(Minobs(
+    TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
+        n_rays=5,
+        agent=51,
+        reward=functools.partial(
+            speed_reward,
+            collision_penalty=0
+        ),
+        stop_on_collision=False,
+    ), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
+), options=[(0, 5), (5, 5), (10, 5)]) for _ in range(50)]
+
+env_fn = lambda i: envs[i]
+policy = DiscretePolicy(env_fn(0).action_space.n)
+value = Value()
+pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
+
+# %%
+value, policy = ppo(
+    env_fn=env_fn,
+    value=value,
+    policy=policy,
+    epochs=50,
+    rollout_episodes=30,
+    rollout_steps=20,
+    gamma=0.99,
+    gae_lambda=0.95,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+    v_opt=v_opt,
+    v_iters=1000,
+)
+
+torch.save(policy.state_dict(), 'ppo-options-minobs.pt')
+
+# %%
+policy = DiscretePolicy(env_fn(0).action_space.n)
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy.load_state_dict(torch.load('ppo-options-minobs.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    #action, _ = policy.predict(torch.tensor(obs))
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action, render_mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()
+
+# %%

scratch/etienne/trpo/experiments/ppo-pendulum.py

@@ -0,0 +1,27 @@
+import gym
+from core.ppo import ppo
+from core.value import Value
+from core.policy import Policy
+import torch.optim
+
+env_fn = lambda _: gym.make('Pendulum-v0')
+policy = Policy(env_fn(0).action_space.shape[0])
+value = Value()
+pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
+
+ppo(
+    env_fn=env_fn,
+    value=value,
+    policy=policy,
+    epochs=300,
+    rollout_episodes=100,
+    rollout_steps=200,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+    v_opt=v_opt,
+    v_iters=1000,
+)

scratch/etienne/trpo/experiments/readme.md

@@ -0,0 +1,7 @@
+|                     | TRPO | PPO   | GAIL  | GAIL PPO | WGAIL | WGAIL PPO |
+|---------------------|------|-------|-------|----------|-------|-----------|
+| Pendulum            | -120 | -1000 | -120  | -1000    | -120  | -1000     |
+| intersimple-minobs  | +1@30|       | +6@26 | +1@20    | -7000@26, -2000@60 | -6000@30, -5000@60 |
+| intersimple-setobs  |      |       | -200@20 |        |       |           |
+| intersimple-minobs2 |      |       | -1500@800 |      |       |           |
+| intersimple-setobs2 |      |       | -500@800 | -750@800 | -1300@800 | -2500@600, unstable |

scratch/etienne/trpo/experiments/requirements.txt

@@ -0,0 +1,3 @@
+torch
+stable-baselines3
+gym

scratch/etienne/trpo/experiments/sgail-options-setobs2.py

@@ -0,0 +1,111 @@
+# %%
+import sys
+sys.path.append('../../../../')
+
+import gym
+from src.safe_options.options import gail
+from src.core.gail import Buffer
+from src.core.value import SetValue
+from src.safe_options.policy import SetMaskedDiscretePolicy
+from src.core.discriminator import DeepsetDiscriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+from src.util.wrappers import CollisionPenaltyWrapper, TransformObservation, Setobs
+import numpy as np
+from src.safe_options.options import SafeOptionsEnv
+from torch.utils.tensorboard import SummaryWriter
+from src.core.reparam_module import ReparamPolicy
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [SafeOptionsEnv(Setobs(
+    TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
+        n_rays=5,
+        reward=functools.partial(
+            speed_reward,
+            collision_penalty=0
+        ),
+        stop_on_collision=True,
+    ), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
+), options=[(0, 5), (1, 5), (2, 5), (4, 5), (6, 5), (8, 5), (10, 5)], safe_actions_collision_method='circle', abort_unsafe_collision_method='circle') for _ in range(60)]
+
+env_fn = lambda i: envs[i]
+policy = SetMaskedDiscretePolicy(env_fn(0).action_space.n)
+value = SetValue()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4)
+
+discriminator = DeepsetDiscriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-4)
+
+expert_data = torch.load('intersimple-expert-data-setobs2.pt')
+expert_data = Buffer(*expert_data)
+
+# %%
+def callback(epoch, value, policy):
+    if not epoch % 10:
+        torch.save(policy.state_dict(), f'sgail-options-setobs2-{epoch}.pt')
+        torch.save(value.state_dict(), f'sgail-options-setobs2-value-{epoch}.pt')
+
+value, policy = gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=100,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=300,
+    rollout_episodes=60,
+    rollout_steps=60,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+    logger=SummaryWriter(comment='sgail-options-setobs2'),
+    callback=callback,
+)
+
+torch.save(policy.state_dict(), 'sgail-options-setobs2.pt')
+
+# %%
+policy = SetMaskedDiscretePolicy(env_fn(0).action_space.n)
+policy(torch.zeros(env_fn(0).observation_space['observation'].shape), torch.zeros(env_fn(0).observation_space['safe_actions'].shape))
+policy = ReparamPolicy(policy)
+policy.load_state_dict(torch.load('sgail-options-setobs2.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    action = policy.sample(policy(
+        torch.tensor(obs['observation'], dtype=torch.float32),
+        torch.tensor(obs['safe_actions'], dtype=torch.float32),
+    ))
+    obs, reward, done, _ = env.step(action, render_mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()
+
+# %%

scratch/etienne/trpo/experiments/sgail-ppo-options-setobs2.py

@@ -0,0 +1,110 @@
+# %%
+import sys
+sys.path.append('../../../../')
+
+import gym
+from src.safe_options.options import gail_ppo, Buffer
+from src.core.value import SetValue
+from src.safe_options.policy import SetMaskedDiscretePolicy
+from src.core.discriminator import DeepsetDiscriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+from src.util.wrappers import CollisionPenaltyWrapper, TransformObservation, Setobs
+import numpy as np
+from src.safe_options.options import SafeOptionsEnv
+from torch.utils.tensorboard import SummaryWriter
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [SafeOptionsEnv(Setobs(
+    TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
+        n_rays=5,
+        reward=functools.partial(
+            speed_reward,
+            collision_penalty=0
+        ),
+        stop_on_collision=True,
+    ), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
+), options=[(0, 5), (1, 5), (2, 5), (4, 5), (6, 5), (8, 5), (10, 5)], safe_actions_collision_method='circle', abort_unsafe_collision_method='circle') for _ in range(60)]
+
+env_fn = lambda i: envs[i]
+
+policy = SetMaskedDiscretePolicy(env_fn(0).action_space.n)
+pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
+
+value = SetValue()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
+
+discriminator = DeepsetDiscriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-4)
+
+expert_data = torch.load('intersimple-expert-data-setobs2.pt')
+expert_data = Buffer(*expert_data)
+
+# %%
+def callback(epoch, value, policy):
+    if not epoch % 10:
+        torch.save(policy.state_dict(), f'sgail-ppo-options-setobs2-{epoch}.pt')
+        torch.save(value.state_dict(), f'sgail-ppo-options-setobs2-value-{epoch}.pt')
+
+value, policy = gail_ppo(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=100,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=200,
+    rollout_episodes=60,
+    rollout_steps=60,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+    logger=SummaryWriter(comment='sgail-ppo-options-setobs2'),
+    callback=callback,
+)
+
+torch.save(policy.state_dict(), 'sgail-ppo-options-setobs2.pt')
+
+# %%
+policy = SetMaskedDiscretePolicy(env_fn(0).action_space.n)
+policy(torch.zeros(env_fn(0).observation_space['observation'].shape), torch.zeros(env_fn(0).observation_space['safe_actions'].shape))
+policy.load_state_dict(torch.load('sgail-ppo-options-setobs2.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    action = policy.sample(policy(
+        torch.tensor(obs['observation'], dtype=torch.float32),
+        torch.tensor(obs['safe_actions'], dtype=torch.float32),
+    ))
+    obs, reward, done, _ = env.step(action, render_mode='post')
+    print('step', i, 'reward', reward, 'safe actions', obs['safe_actions'])
+    if done:
+        break
+env.close()
+# %%

scratch/etienne/trpo/experiments/trpo-intersimple-minobs.py

@@ -0,0 +1,88 @@
+# %%
+import gym
+from core.sampling import rollout
+from core.trpo import trpo
+from core.value import Value
+from core.policy import Policy
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+import numpy as np
+from gym.wrappers import TransformObservation
+from util.wrappers import CollisionPenaltyWrapper
+from core.reparam_module import ReparamPolicy
+
+from util.wrappers import Minobs
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [Minobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(50)]
+
+env_fn = lambda i: envs[i]
+policy = Policy(env_fn(0).action_space.shape[0])
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
+
+# %%
+value, policy = trpo(
+    env_fn=env_fn,
+    value=value,
+    policy=policy,
+    epochs=50,
+    rollout_episodes=30,
+    rollout_steps=100,
+    gamma=0.99,
+    gae_lambda=0.95,
+    delta=0.01,
+    backtrack_coeff=0.9,
+    backtrack_iters=50,
+    v_opt=v_opt,
+    v_iters=1000,
+    cg_damping=0.1,
+)
+
+torch.save(policy.state_dict(), 'trpo-intersimple-minobs.pt')
+
+# %%
+policy = Policy(env_fn(0).action_space.shape[0])
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy = ReparamPolicy(policy)
+policy.load_state_dict(torch.load('trpo-intersimple-minobs.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    #action, _ = policy.predict(torch.tensor(obs))
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action)
+    env.render(mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()

scratch/etienne/trpo/experiments/trpo-intersimple-minobs2.py

@@ -0,0 +1,87 @@
+# %%
+import gym
+from core.sampling import rollout
+from core.trpo import trpo
+from core.value import Value
+from core.policy import Policy
+import torch.optim
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+import numpy as np
+from gym.wrappers import TransformObservation
+from util.wrappers import CollisionPenaltyWrapper
+from core.reparam_module import ReparamPolicy
+
+from util.wrappers import Minobs
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [Minobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
+    n_rays=5,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(50)]
+
+env_fn = lambda i: envs[i]
+policy = Policy(env_fn(0).action_space.shape[0])
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
+
+# %%
+value, policy = trpo(
+    env_fn=env_fn,
+    value=value,
+    policy=policy,
+    epochs=200,
+    rollout_episodes=30,
+    rollout_steps=100,
+    gamma=0.99,
+    gae_lambda=0.95,
+    delta=0.01,
+    backtrack_coeff=0.9,
+    backtrack_iters=50,
+    v_opt=v_opt,
+    v_iters=1000,
+    cg_damping=0.1,
+)
+
+torch.save(policy.state_dict(), 'trpo-intersimple-minobs2.pt')
+
+# %%
+policy = Policy(env_fn(0).action_space.shape[0])
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy = ReparamPolicy(policy)
+policy.load_state_dict(torch.load('trpo-intersimple-minobs2.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    #action, _ = policy.predict(torch.tensor(obs))
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action)
+    env.render(mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()

scratch/etienne/trpo/experiments/trpo-intersimple-normobs.py

@@ -0,0 +1,62 @@
+import gym
+from core.sampling import rollout
+from core.trpo import trpo
+from core.value import Value
+from core.policy import Policy
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+import numpy as np
+from gym.wrappers import TransformObservation
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [TransformObservation(IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=10
+    ),
+), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)) for _ in range(50)]
+
+env_fn = lambda i: envs[i]
+policy = Policy(env_fn(0).action_space.shape[0])
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
+
+value, policy = trpo(
+    env_fn=env_fn,
+    value=value,
+    policy=policy,
+    epochs=4000,
+    rollout_episodes=30,
+    rollout_steps=100,
+    gamma=0.99,
+    gae_lambda=0.95,
+    delta=0.01,
+    backtrack_coeff=0.9,
+    backtrack_iters=50,
+    v_opt=v_opt,
+    v_iters=1000,
+    cg_damping=0.1,
+)
+
+#rollout(env_fn, policy, n_episodes=9, max_steps_per_episode=200, render=True)

scratch/etienne/trpo/experiments/trpo-intersimple-setobs.py

@@ -0,0 +1,90 @@
+# %%
+import gym
+from core.sampling import rollout
+from core.trpo import trpo
+from core.value import SetValue
+from core.policy import DeepSetPolicy
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+import numpy as np
+from gym.wrappers import TransformObservation
+from util.wrappers import CollisionPenaltyWrapper
+from core.reparam_module import ReparamPolicy
+
+from util.wrappers import Setobs
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [Setobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(50)]
+
+env_fn = lambda i: envs[i]
+policy = DeepSetPolicy(env_fn(0).action_space.shape[0])
+value = SetValue()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
+
+# %%
+value, policy = trpo(
+    env_fn=env_fn,
+    value=value,
+    policy=policy,
+    epochs=150,
+    rollout_episodes=30,
+    rollout_steps=100,
+    gamma=0.99,
+    gae_lambda=0.95,
+    delta=0.01,
+    backtrack_coeff=0.9,
+    backtrack_iters=50,
+    v_opt=v_opt,
+    v_iters=1000,
+    cg_damping=0.1,
+)
+
+torch.save(policy.state_dict(), 'trpo-intersimple-setobs.pt')
+
+# %%
+policy = DeepSetPolicy(env_fn(0).action_space.shape[0])
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy = ReparamPolicy(policy)
+policy.load_state_dict(torch.load('trpo-intersimple-setobs.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    #action, _ = policy.predict(torch.tensor(obs))
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action)
+    env.render(mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()
+
+# %%

scratch/etienne/trpo/experiments/trpo-intersimple-setobs2.py

@@ -0,0 +1,87 @@
+# %%
+import gym
+from core.sampling import rollout
+from core.trpo import trpo
+from core.value import DeepSetValue
+from core.policy import DeepSetPolicy
+import torch.optim
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+import numpy as np
+from gym.wrappers import TransformObservation
+from util.wrappers import CollisionPenaltyWrapper
+from core.reparam_module import ReparamPolicy
+
+from util.wrappers import Setobs
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [Setobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
+    n_rays=5,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(50)]
+
+env_fn = lambda i: envs[i]
+policy = DeepSetPolicy(env_fn(0).action_space.shape[0])
+value = DeepSetValue()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
+
+# %%
+value, policy = trpo(
+    env_fn=env_fn,
+    value=value,
+    policy=policy,
+    epochs=200,
+    rollout_episodes=30,
+    rollout_steps=100,
+    gamma=0.99,
+    gae_lambda=0.95,
+    delta=0.01,
+    backtrack_coeff=0.9,
+    backtrack_iters=50,
+    v_opt=v_opt,
+    v_iters=1000,
+    cg_damping=0.1,
+)
+
+torch.save(policy.state_dict(), 'trpo-intersimple-setobs2.pt')
+
+# %%
+policy = DeepSetPolicy(env_fn(0).action_space.shape[0])
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy = ReparamPolicy(policy)
+policy.load_state_dict(torch.load('trpo-intersimple-setobs2.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    #action, _ = policy.predict(torch.tensor(obs))
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action)
+    env.render(mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()

scratch/etienne/trpo/experiments/trpo-intersimple.py

@@ -0,0 +1,42 @@
+import gym
+from core.sampling import rollout
+from core.trpo import trpo
+from core.value import Value
+from core.policy import Policy
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+
+envs = [IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=10
+    ),
+) for _ in range(50)]
+
+env_fn = lambda i: envs[i]
+policy = Policy(env_fn(0).action_space.shape[0])
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
+
+value, policy = trpo(
+    env_fn=env_fn,
+    value=value,
+    policy=policy,
+    epochs=4000,
+    rollout_episodes=30,
+    rollout_steps=100,
+    gamma=0.99,
+    gae_lambda=0.95,
+    delta=0.01,
+    backtrack_coeff=0.9,
+    backtrack_iters=50,
+    v_opt=v_opt,
+    v_iters=1000,
+    cg_damping=0.1,
+)
+
+#rollout(env_fn, policy, n_episodes=9, max_steps_per_episode=200, render=True)

scratch/etienne/trpo/experiments/trpo-options-minobs.py

@@ -0,0 +1,91 @@
+# %%
+import gym
+from core.sampling import rollout
+from core.trpo import trpo
+from core.value import Value
+from core.policy import DiscretePolicy
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+import numpy as np
+from util.wrappers import CollisionPenaltyWrapper, TransformObservation
+from core.reparam_module import ReparamPolicy
+
+from util.wrappers import Minobs
+from options.options import OptionsEnv
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [OptionsEnv(Minobs(
+    TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
+        n_rays=5,
+        agent=51,
+        reward=functools.partial(
+            speed_reward,
+            collision_penalty=0
+        ),
+        stop_on_collision=False,
+    ), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
+), options=[(0, 5), (5, 5), (10, 5)]) for _ in range(50)]
+
+env_fn = lambda i: envs[i]
+policy = DiscretePolicy(env_fn(0).action_space.n)
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
+
+# %%
+value, policy = trpo(
+    env_fn=env_fn,
+    value=value,
+    policy=policy,
+    epochs=50,
+    rollout_episodes=30,
+    rollout_steps=20,
+    gamma=0.99,
+    gae_lambda=0.95,
+    delta=0.01,
+    backtrack_coeff=0.9,
+    backtrack_iters=50,
+    v_opt=v_opt,
+    v_iters=1000,
+    cg_damping=0.1,
+)
+
+torch.save(policy.state_dict(), 'trpo-options-minobs.pt')
+
+# %%
+policy = DiscretePolicy(env_fn(0).action_space.n)
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy = ReparamPolicy(policy)
+policy.load_state_dict(torch.load('trpo-options-minobs.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    #action, _ = policy.predict(torch.tensor(obs))
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action, render_mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()
+
+# %%

scratch/etienne/trpo/experiments/trpo-pendulum-rollout.py

@@ -0,0 +1,17 @@
+import gym
+from core.gail import gail
+from core.reparam_module import ReparamPolicy
+from core.value import Value
+from core.policy import Policy
+from core.discriminator import Discriminator
+import torch.optim
+from core.sampling import rollout
+
+env_fn = lambda _: gym.make('Pendulum-v0')
+policy = Policy(env_fn(0).action_space.shape[0])
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy = ReparamPolicy(policy)
+policy.load_state_dict(torch.load('trpo-pendulum.pt'))
+
+expert_data = rollout(env_fn, policy, n_episodes=20, max_steps_per_episode=200)
+torch.save(expert_data, 'trpo-pendulum-expert-data.pt')

scratch/etienne/trpo/experiments/trpo-pendulum.py

@@ -0,0 +1,30 @@
+import gym
+from gym.wrappers import TransformObservation
+from core.trpo import trpo
+from core.value import Value
+from core.policy import Policy
+import torch.optim
+
+env_fn = lambda _: TransformObservation(gym.make('Pendulum-v0'), lambda obs: obs)
+policy = Policy(env_fn(0).action_space.shape[0])
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3, weight_decay=1e-4)
+
+value, policy = trpo(
+    env_fn=env_fn,
+    value=value,
+    policy=policy,
+    epochs=100,
+    rollout_episodes=20,
+    rollout_steps=250,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+    v_opt=v_opt,
+    v_iters=1000,
+)
+
+
+torch.save(policy.state_dict(), 'trpo-pendulum.pt')

scratch/etienne/trpo/experiments/trpo-walker.py

@@ -0,0 +1,26 @@
+import gym
+from core.trpo import trpo
+from core.value import Value
+from core.policy import Policy
+import torch.optim
+
+env_fn = lambda _: gym.make('BipedalWalker-v3')
+policy = Policy(env_fn(0).action_space.shape[0])
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-2)
+
+trpo(
+    env_fn=env_fn,
+    value=value,
+    policy=policy,
+    epochs=1000,
+    rollout_episodes=20,
+    rollout_steps=250,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+    v_opt=v_opt,
+    v_iters=1000,
+)

scratch/etienne/trpo/experiments/vec-env.ipynb

@@ -0,0 +1,346 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from stable_baselines3.common.env_util import make_vec_env\n",
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "env = make_vec_env('Pendulum-v0', n_envs=6)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(6, 3)"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "obs = env.reset()\n",
+    "obs.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1\n",
+      "2\n",
+      "3\n",
+      "4\n",
+      "5\n",
+      "6\n",
+      "7\n",
+      "8\n",
+      "9\n",
+      "10\n",
+      "11\n",
+      "12\n",
+      "13\n",
+      "14\n",
+      "15\n",
+      "16\n",
+      "17\n",
+      "18\n",
+      "19\n",
+      "20\n",
+      "21\n",
+      "22\n",
+      "23\n",
+      "24\n",
+      "25\n",
+      "26\n",
+      "27\n",
+      "28\n",
+      "29\n",
+      "30\n",
+      "31\n",
+      "32\n",
+      "33\n",
+      "34\n",
+      "35\n",
+      "36\n",
+      "37\n",
+      "38\n",
+      "39\n",
+      "40\n",
+      "41\n",
+      "42\n",
+      "43\n",
+      "44\n",
+      "45\n",
+      "46\n",
+      "47\n",
+      "48\n",
+      "49\n",
+      "50\n",
+      "51\n",
+      "52\n",
+      "53\n",
+      "54\n",
+      "55\n",
+      "56\n",
+      "57\n",
+      "58\n",
+      "59\n",
+      "60\n",
+      "61\n",
+      "62\n",
+      "63\n",
+      "64\n",
+      "65\n",
+      "66\n",
+      "67\n",
+      "68\n",
+      "69\n",
+      "70\n",
+      "71\n",
+      "72\n",
+      "73\n",
+      "74\n",
+      "75\n",
+      "76\n",
+      "77\n",
+      "78\n",
+      "79\n",
+      "80\n",
+      "81\n",
+      "82\n",
+      "83\n",
+      "84\n",
+      "85\n",
+      "86\n",
+      "87\n",
+      "88\n",
+      "89\n",
+      "90\n",
+      "91\n",
+      "92\n",
+      "93\n",
+      "94\n",
+      "95\n",
+      "96\n",
+      "97\n",
+      "98\n",
+      "99\n",
+      "100\n",
+      "101\n",
+      "102\n",
+      "103\n",
+      "104\n",
+      "105\n",
+      "106\n",
+      "107\n",
+      "108\n",
+      "109\n",
+      "110\n",
+      "111\n",
+      "112\n",
+      "113\n",
+      "114\n",
+      "115\n",
+      "116\n",
+      "117\n",
+      "118\n",
+      "119\n",
+      "120\n",
+      "121\n",
+      "122\n",
+      "123\n",
+      "124\n",
+      "125\n",
+      "126\n",
+      "127\n",
+      "128\n",
+      "129\n",
+      "130\n",
+      "131\n",
+      "132\n",
+      "133\n",
+      "134\n",
+      "135\n",
+      "136\n",
+      "137\n",
+      "138\n",
+      "139\n",
+      "140\n",
+      "141\n",
+      "142\n",
+      "143\n",
+      "144\n",
+      "145\n",
+      "146\n",
+      "147\n",
+      "148\n",
+      "149\n",
+      "150\n",
+      "151\n",
+      "152\n",
+      "153\n",
+      "154\n",
+      "155\n",
+      "156\n",
+      "157\n",
+      "158\n",
+      "159\n",
+      "160\n",
+      "161\n",
+      "162\n",
+      "163\n",
+      "164\n",
+      "165\n",
+      "166\n",
+      "167\n",
+      "168\n",
+      "169\n",
+      "170\n",
+      "171\n",
+      "172\n",
+      "173\n",
+      "174\n",
+      "175\n",
+      "176\n",
+      "177\n",
+      "178\n",
+      "179\n",
+      "180\n",
+      "181\n",
+      "182\n",
+      "183\n",
+      "184\n",
+      "185\n",
+      "186\n",
+      "187\n",
+      "188\n",
+      "189\n",
+      "190\n",
+      "191\n",
+      "192\n",
+      "193\n",
+      "194\n",
+      "195\n",
+      "196\n",
+      "197\n",
+      "198\n",
+      "199\n",
+      "200\n"
+     ]
+    }
+   ],
+   "source": [
+    "dones = [False]\n",
+    "i = 0\n",
+    "while not any(dones):\n",
+    "    i += 1\n",
+    "    print(i)\n",
+    "    _, _, dones, _ = env.step(np.zeros((6, 1)))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([ True,  True,  True,  True,  True,  True])"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "dones"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "_, _, dones, _ = env.step(np.zeros((6, 1)))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([False, False, False, False, False, False])"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "dones"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "6c7a4ac80dd345f83235e10baa3acc437d966916e1cc075a45b91bb9cc030938"
+  },
+  "kernelspec": {
+   "display_name": "Python 3.9.7 64-bit ('.venv': venv)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.7"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

scratch/etienne/trpo/experiments/wgail-intersimple-minobs.py

@@ -0,0 +1,76 @@
+import gym
+from core.gail import gail, Buffer
+from core.value import Value
+from core.policy import Policy
+from core.discriminator import Discriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, Minobs
+import numpy as np
+from gym.wrappers import TransformObservation
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [Minobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(30)]
+env_fn = lambda i: envs[i]
+
+policy = Policy(env_fn(0).action_space.shape[0])
+
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
+
+discriminator = Discriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-2)
+
+expert_data = torch.load('intersimple-expert-data-minobs.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=100,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=4000,
+    rollout_episodes=30,
+    rollout_steps=200,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+    wasserstein=True,
+    wasserstein_c=1.,
+)
+
+torch.save(policy.state_dict(), 'wgail-intersimple-minobs.pt')

scratch/etienne/trpo/experiments/wgail-intersimple-minobs2.py

@@ -0,0 +1,75 @@
+import gym
+from core.gail import gail, Buffer
+from core.value import Value
+from core.policy import Policy
+from core.discriminator import Discriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, Minobs
+import numpy as np
+from gym.wrappers import TransformObservation
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [Minobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
+    n_rays=5,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(50)]
+env_fn = lambda i: envs[i]
+
+policy = Policy(env_fn(0).action_space.shape[0])
+
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
+
+discriminator = Discriminator()
+disc_opt = torch.optim.RMSprop(discriminator.parameters(), lr=1e-3, weight_decay=1e-5)
+
+expert_data = torch.load('intersimple-expert-data-minobs2.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=500,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=4000,
+    rollout_episodes=50,
+    rollout_steps=200,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+    wasserstein=True,
+    wasserstein_c=0.1,
+)
+
+torch.save(policy.state_dict(), 'wgail-intersimple-minobs2.pt')

scratch/etienne/trpo/experiments/wgail-intersimple-setobs2.py

@@ -0,0 +1,77 @@
+import gym
+from core.gail import gail, Buffer
+from core.value import SetValue
+from core.policy import SetPolicy
+from core.discriminator import DeepsetDiscriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, Setobs
+import numpy as np
+from gym.wrappers import TransformObservation
+from torch.utils.tensorboard import SummaryWriter
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [Setobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
+    n_rays=5,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(50)]
+env_fn = lambda i: envs[i]
+
+policy = SetPolicy(env_fn(0).action_space.shape[0])
+
+value = SetValue()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
+
+discriminator = DeepsetDiscriminator()
+disc_opt = torch.optim.RMSprop(discriminator.parameters(), lr=1e-3)
+
+expert_data = torch.load('intersimple-expert-data-setobs2.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=500,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=800,
+    rollout_episodes=50,
+    rollout_steps=200,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+    wasserstein=True,
+    wasserstein_c=100.,
+    logger=SummaryWriter(comment='-wgail-setobs2'),
+)
+
+torch.save(policy.state_dict(), 'wgail-intersimple-setobs2.pt')

scratch/etienne/trpo/experiments/wgail-intersimple.py

@@ -0,0 +1,56 @@
+import gym
+from core.gail import gail, Buffer
+from core.value import Value
+from core.policy import Policy
+from core.discriminator import Discriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper
+
+envs = [CollisionPenaltyWrapper(IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100) for _ in range(30)]
+env_fn = lambda i: envs[i]
+
+policy = Policy(env_fn(0).action_space.shape[0])
+
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4)
+
+discriminator = Discriminator()
+disc_opt = torch.optim.RMSprop(discriminator.parameters(), lr=1e-4)
+
+expert_data = torch.load('intersimple-expert-data.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=10,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=4000,
+    rollout_episodes=30,
+    rollout_steps=100,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+    wasserstein=True,
+    wasserstein_c=1.,
+)
+
+torch.save(policy.state_dict(), 'wgail-intersimple.pt')

scratch/etienne/trpo/experiments/wgail-options-setobs.py

@@ -0,0 +1,101 @@
+# %%
+import gym
+from options.options import gail
+from core.gail import Buffer
+from core.value import SetValue
+from core.policy import SetDiscretePolicy
+from core.discriminator import DeepsetDiscriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, TransformObservation, Setobs
+import numpy as np
+from options.options import OptionsEnv
+from torch.utils.tensorboard import SummaryWriter
+from core.reparam_module import ReparamPolicy
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [OptionsEnv(Setobs(
+    TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
+        n_rays=5,
+        agent=51,
+        reward=functools.partial(
+            speed_reward,
+            collision_penalty=0
+        ),
+        stop_on_collision=False,
+    ), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
+), options=[(0, 5), (1, 5), (2, 5), (4, 5), (8, 5)]) for _ in range(60)]
+
+env_fn = lambda i: envs[i]
+policy = SetDiscretePolicy(env_fn(0).action_space.n)
+value = SetValue()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4)
+
+discriminator = DeepsetDiscriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-3)
+
+expert_data = torch.load('intersimple-expert-data-setobs.pt')
+expert_data = Buffer(*expert_data)
+
+# %%
+value, policy = gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=100,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=150,
+    rollout_episodes=60,
+    rollout_steps=60,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+    wasserstein=True,
+    wasserstein_c=1.,
+    logger=SummaryWriter(comment='wgail-options-setobs'),
+)
+
+torch.save(policy.state_dict(), 'wgail-options-setobs.pt')
+
+# %%
+policy = SetDiscretePolicy(env_fn(0).action_space.n)
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy = ReparamPolicy(policy)
+policy.load_state_dict(torch.load('wgail-options-setobs.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    #action, _ = policy.predict(torch.tensor(obs))
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action, render_mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()

scratch/etienne/trpo/experiments/wgail-options-setobs2.py

@@ -0,0 +1,100 @@
+# %%
+import gym
+from options.options import gail
+from core.gail import Buffer
+from core.value import SetValue
+from core.policy import SetDiscretePolicy
+from core.discriminator import DeepsetDiscriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, TransformObservation, Setobs
+import numpy as np
+from options.options import OptionsEnv
+from torch.utils.tensorboard import SummaryWriter
+from core.reparam_module import ReparamPolicy
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [OptionsEnv(Setobs(
+    TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
+        n_rays=5,
+        reward=functools.partial(
+            speed_reward,
+            collision_penalty=0
+        ),
+        stop_on_collision=False,
+    ), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
+), options=[(0, 5), (1, 5), (2, 5), (4, 5), (6, 5), (8, 5)]) for _ in range(60)]
+
+env_fn = lambda i: envs[i]
+policy = SetDiscretePolicy(env_fn(0).action_space.n)
+value = SetValue()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4)
+
+discriminator = DeepsetDiscriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-4)
+
+expert_data = torch.load('intersimple-expert-data-setobs2.pt')
+expert_data = Buffer(*expert_data)
+
+# %%
+value, policy = gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=100,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=200,
+    rollout_episodes=60,
+    rollout_steps=60,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+    wasserstein=True,
+    wasserstein_c=1.,
+    logger=SummaryWriter(comment='wgail-options-setobs2'),
+)
+
+torch.save(policy.state_dict(), 'wgail-options-setobs2.pt')
+
+# %%
+policy = SetDiscretePolicy(env_fn(0).action_space.n)
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy = ReparamPolicy(policy)
+policy.load_state_dict(torch.load('wgail-options-setobs2.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    #action, _ = policy.predict(torch.tensor(obs))
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action, render_mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()

scratch/etienne/trpo/experiments/wgail-pendulum.py

@@ -0,0 +1,40 @@
+import gym
+from core.gail import gail, Buffer
+from core.value import Value
+from core.policy import Policy
+from core.discriminator import Discriminator
+import torch.optim
+
+env_fn = lambda _: gym.make('Pendulum-v0')
+policy = Policy(env_fn(0).action_space.shape[0])
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
+discriminator = Discriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-3)
+
+expert_data = torch.load('trpo-pendulum-expert-data.pt')
+expert_data = Buffer(*expert_data)
+
+gail(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=10,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=100,
+    rollout_episodes=20,
+    rollout_steps=200,
+    gamma=0.99,
+    gae_lambda=0.9,
+    delta=0.01,
+    backtrack_coeff=0.8,
+    backtrack_iters=10,
+    wasserstein=True,
+    wasserstein_c=100.,
+)
+
+torch.save(policy.state_dict(), 'gail-pendulum.pt')

scratch/etienne/trpo/experiments/wgail-ppo-intersimple-minobs.py

@@ -0,0 +1,77 @@
+import gym
+from core.gail import gail_ppo, Buffer
+from core.value import Value
+from core.policy import Policy
+from core.discriminator import Discriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, Minobs
+import numpy as np
+from gym.wrappers import TransformObservation
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [Minobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(30)]
+env_fn = lambda i: envs[i]
+
+policy = Policy(env_fn(0).action_space.shape[0])
+pi_opt = torch.optim.RMSprop(policy.parameters(), lr=3e-4)
+
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3, weight_decay=1e-3)
+
+discriminator = Discriminator()
+disc_opt = torch.optim.RMSprop(discriminator.parameters(), lr=1e-3, weight_decay=1e-5)
+
+expert_data = torch.load('intersimple-expert-data-minobs.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail_ppo(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=100,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=4000,
+    rollout_episodes=30,
+    rollout_steps=100,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+    wasserstein=True,
+    wasserstein_c=1.,
+)
+
+torch.save(policy.state_dict(), 'wgail-ppo-intersimple-minobs.pt')

scratch/etienne/trpo/experiments/wgail-ppo-intersimple-setobs2.py

@@ -0,0 +1,78 @@
+import gym
+from core.gail import gail_ppo, Buffer
+from core.value import SetValue
+from core.policy import SetPolicy
+from core.discriminator import DeepsetDiscriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, Setobs
+import numpy as np
+from gym.wrappers import TransformObservation
+from torch.utils.tensorboard import SummaryWriter
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [Setobs(TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
+    n_rays=5,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))) for _ in range(50)]
+env_fn = lambda i: envs[i]
+
+policy = SetPolicy(env_fn(0).action_space.shape[0])
+pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
+
+value = SetValue()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-4, weight_decay=1e-3)
+
+discriminator = DeepsetDiscriminator()
+disc_opt = torch.optim.RMSprop(discriminator.parameters(), lr=1e-3)
+
+expert_data = torch.load('intersimple-expert-data-setobs2.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail_ppo(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=500,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=800,
+    rollout_episodes=50,
+    rollout_steps=200,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+    wasserstein=True,
+    wasserstein_c=100.,
+    logger=SummaryWriter(comment='-wgail-ppo-setobs2'),
+)
+
+torch.save(policy.state_dict(), 'wgail-ppo-intersimple-setobs2.pt')

scratch/etienne/trpo/experiments/wgail-ppo-intersimple.py

@@ -0,0 +1,57 @@
+import gym
+from core.gail import gail_ppo, Buffer
+from core.value import Value
+from core.policy import Policy
+from core.discriminator import Discriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from collision_penalty import CollisionPenaltyWrapper
+
+envs = [CollisionPenaltyWrapper(IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+), collision_distance=6, collision_penalty=100) for _ in range(30)]
+env_fn = lambda i: envs[i]
+
+policy = Policy(env_fn(0).action_space.shape[0])
+pi_opt = torch.optim.RMSprop(policy.parameters(), lr=3e-4)
+
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
+
+discriminator = Discriminator()
+disc_opt = torch.optim.RMSprop(discriminator.parameters(), lr=3e-4)
+
+expert_data = torch.load('intersimple-expert-data.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail_ppo(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=10,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=4000,
+    rollout_episodes=30,
+    rollout_steps=100,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+    wasserstein=True,
+    wasserstein_c=1.,
+)
+
+torch.save(policy.state_dict(), 'wgail-ppo-intersimple.pt')

scratch/etienne/trpo/experiments/wgail-ppo-options-setobs.py

@@ -0,0 +1,98 @@
+import gym
+from options.options import gail_ppo, Buffer
+from core.value import SetValue
+from core.policy import SetDiscretePolicy
+from core.discriminator import DeepsetDiscriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, TransformObservation, Setobs
+import numpy as np
+from options.options import OptionsEnv
+from torch.utils.tensorboard import SummaryWriter
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [OptionsEnv(Setobs(
+    TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlat(
+        n_rays=5,
+        agent=51,
+        reward=functools.partial(
+            speed_reward,
+            collision_penalty=0
+        ),
+        stop_on_collision=False,
+    ), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
+), options=[(0, 5), (1, 5), (2, 5), (4, 5), (8, 5)]) for _ in range(60)]
+
+env_fn = lambda i: envs[i]
+
+policy = SetDiscretePolicy(env_fn(0).action_space.n)
+pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
+
+value = SetValue()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
+
+discriminator = DeepsetDiscriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-3)
+
+expert_data = torch.load('intersimple-expert-data-setobs.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail_ppo(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=100,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=150,
+    rollout_episodes=60,
+    rollout_steps=60,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+    wasserstein=True,
+    wasserstein_c=1.,
+    logger=SummaryWriter(comment='wgail-ppo-options-setobs'),
+)
+
+torch.save(policy.state_dict(), 'wgail-ppo-options-setobs.pt')
+
+# %%
+policy = SetDiscretePolicy(env_fn(0).action_space.n)
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy.load_state_dict(torch.load('wgail-ppo-options-setobs.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action, render_mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()

scratch/etienne/trpo/experiments/wgail-ppo-options-setobs2.py

@@ -0,0 +1,97 @@
+import gym
+from options.options import gail_ppo, Buffer
+from core.value import SetValue
+from core.policy import SetDiscretePolicy
+from core.discriminator import DeepsetDiscriminator
+import torch.optim
+from intersim.envs import IntersimpleLidarFlatRandom
+from intersim.envs.intersimple import speed_reward
+import functools
+from util.wrappers import CollisionPenaltyWrapper, TransformObservation, Setobs
+import numpy as np
+from options.options import OptionsEnv
+from torch.utils.tensorboard import SummaryWriter
+
+obs_min = np.array([
+    [-1000, -1000, 0, -np.pi, -1e-1, 0.],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+    [0, -np.pi, -20, -20, -np.pi, -1e-1],
+]).reshape(-1)
+
+obs_max = np.array([
+    [1000, 1000, 20, np.pi, 1e-1, 0.],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+    [50, np.pi, 20, 20, np.pi, 1e-1],
+]).reshape(-1)
+
+envs = [OptionsEnv(Setobs(
+    TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
+        n_rays=5,
+        reward=functools.partial(
+            speed_reward,
+            collision_penalty=0
+        ),
+        stop_on_collision=False,
+    ), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
+), options=[(0, 5), (1, 5), (2, 5), (4, 5), (6, 5), (8, 5)]) for _ in range(60)]
+
+env_fn = lambda i: envs[i]
+
+policy = SetDiscretePolicy(env_fn(0).action_space.n)
+pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
+
+value = SetValue()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
+
+discriminator = DeepsetDiscriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-4)
+
+expert_data = torch.load('intersimple-expert-data-setobs2.pt')
+expert_data = Buffer(*expert_data)
+
+value, policy = gail_ppo(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=100,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=200,
+    rollout_episodes=60,
+    rollout_steps=60,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+    wasserstein=True,
+    wasserstein_c=1.,
+    logger=SummaryWriter(comment='wgail-ppo-options-setobs2'),
+)
+
+torch.save(policy.state_dict(), 'wgail-ppo-options-setobs2.pt')
+
+# %%
+policy = SetDiscretePolicy(env_fn(0).action_space.n)
+policy(torch.zeros(env_fn(0).observation_space.shape))
+policy.load_state_dict(torch.load('wgail-ppo-options-setobs2.pt'))
+
+env = env_fn(0)
+obs = env.reset()
+env.render(mode='post')
+for i in range(300):
+    action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+    obs, reward, done, _ = env.step(action, render_mode='post')
+    print('step', i, 'reward', reward)
+    if done:
+        break
+env.close()

scratch/etienne/trpo/experiments/wgail-ppo-pendulum.py

@@ -0,0 +1,44 @@
+import gym
+from core.gail import gail_ppo, Buffer
+from core.value import Value
+from core.policy import Policy
+from core.discriminator import Discriminator
+import torch.optim
+
+env_fn = lambda _: gym.make('Pendulum-v0')
+
+policy = Policy(env_fn(0).action_space.shape[0])
+pi_opt = torch.optim.RMSprop(policy.parameters(), lr=3e-4)
+
+value = Value()
+v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
+
+discriminator = Discriminator()
+disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-3)
+
+expert_data = torch.load('trpo-pendulum-expert-data.pt')
+expert_data = Buffer(*expert_data)
+
+gail_ppo(
+    env_fn=env_fn,
+    expert_data=expert_data,
+    discriminator=discriminator,
+    disc_opt=disc_opt,
+    disc_iters=10,
+    policy=policy,
+    value=value,
+    v_opt=v_opt,
+    v_iters=1000,
+    epochs=100,
+    rollout_episodes=20,
+    rollout_steps=200,
+    gamma=0.99,
+    gae_lambda=0.9,
+    clip_ratio=0.2,
+    pi_opt=pi_opt,
+    pi_iters=100,
+    wasserstein=True,
+    wasserstein_c=100.,
+)
+
+torch.save(policy.state_dict(), 'gail-pendulum.pt')

scratch/etienne/trpo/sb3/sb3-ppo-intersimple-nocollision.py

@@ -0,0 +1,55 @@
+from stable_baselines3 import PPO
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from gym import Wrapper
+
+model_name = "ppo_speed_lidar_nocollision"
+
+env = IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+)
+
+class CollisionPenaltyWrapper(Wrapper):
+
+    def __init__(self, env, collision_distance, collision_penalty, *args, **kwargs):
+        super().__init__(env, *args, **kwargs)
+        self.penalty = collision_penalty
+        self.distance = collision_distance
+    
+    def step(self, action):
+        obs, reward, done, info = super().step(action)
+        reward = -self.penalty if (obs.reshape(-1, 6)[1:, 0] < self.distance).any() else reward
+
+        self.env._rewards.pop()
+        self.env._rewards.append(reward)
+
+        return obs, reward, done, info
+
+env = CollisionPenaltyWrapper(env, collision_distance=6, collision_penalty=100)
+
+model = PPO(
+    "MlpPolicy", env,
+    learning_rate=1e-4,
+    verbose=1,
+)
+model.learn(total_timesteps=100000)
+model.save(model_name)
+
+model = PPO.load(model_name)
+obs = env.reset()
+env.render(mode='post')
+for i in range(200):
+    action, _ = model.predict(obs)
+    obs, reward, done, _ = env.step(action)
+    env.render(mode='post')
+    print('step', i, 'front distance', obs.reshape(-1, 6)[3, 0], 'reward', reward)
+    if done:
+        break
+env.close()

scratch/etienne/trpo/sb3/sb3-ppo-intersimple-nocollision2.py

@@ -0,0 +1,58 @@
+from stable_baselines3 import PPO
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+from gym import Wrapper
+
+model_name = "ppo_speed_lidar_nocollision"
+
+env = IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+    stop_on_collision=False,
+)
+
+class CollisionPenaltyWrapper(Wrapper):
+
+    def __init__(self, env, collision_distance, collision_penalty, last_reward_weight, *args, **kwargs):
+        super().__init__(env, *args, **kwargs)
+        self.penalty = collision_penalty
+        self.distance = collision_distance
+        self.last_reward = -collision_penalty
+        self.last_reward_weight = last_reward_weight
+    
+    def step(self, action):
+        obs, reward, done, info = super().step(action)
+        reward = -self.penalty if (obs.reshape(-1, 6)[1:, 0] < self.distance).any() else reward
+        reward = self.last_reward_weight * self.last_reward + (1 - self.last_reward_weight) * self.last_reward
+
+        self.env._rewards.pop()
+        self.env._rewards.append(reward)
+
+        return obs, reward, done, info
+
+env = CollisionPenaltyWrapper(env, collision_distance=6, collision_penalty=10, last_reward_weight=0.9)
+
+model = PPO(
+    "MlpPolicy", env,
+    learning_rate=1e-4,
+    verbose=1,
+)
+model.learn(total_timesteps=100000)
+model.save(model_name)
+
+model = PPO.load(model_name)
+obs = env.reset()
+env.render(mode='post')
+for i in range(200):
+    action, _ = model.predict(obs)
+    obs, reward, done, _ = env.step(action)
+    env.render(mode='post')
+    print('step', i, 'front distance', obs.reshape(-1, 6)[3, 0], 'reward', reward)
+    if done:
+        break
+env.close()

scratch/etienne/trpo/sb3/sb3-ppo-intersimple-rollout.py

@@ -0,0 +1,28 @@
+import sys
+sys.path.append('..')
+
+from stable_baselines3 import PPO
+from core.sampling import rollout_sb3
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+import torch
+from util.wrappers import CollisionPenaltyWrapper
+
+model = PPO.load('sb3-ppo-intersimple')
+env = CollisionPenaltyWrapper(IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+), collision_distance=6, collision_penalty=100)
+
+expert_data = rollout_sb3(env, model, n_episodes=200, max_steps_per_episode=200)
+
+states, actions, rewards, dones = expert_data
+print(f'Expert mean episode length {(~dones).sum() / states.shape[0]}')
+print(f'Expert mean reward per episode {rewards[~dones].sum() / states.shape[0]}')
+
+torch.save(expert_data, 'sb3-ppo-intersimple-expert-data.pt')

scratch/etienne/trpo/sb3/sb3-ppo-intersimple.py

@@ -0,0 +1,23 @@
+from stable_baselines3 import PPO
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+
+env = IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+)
+
+model = PPO(
+    "MlpPolicy", env,
+    learning_rate=1e-4,
+    verbose=1,
+    use_sde=False,
+    sde_sample_freq=4,
+)
+model.learn(total_timesteps=100000)
+model.save('sb3-ppo-intersimple')

scratch/etienne/trpo/sb3/sb3-ppo-pendulum.py

@@ -0,0 +1,6 @@
+from stable_baselines3 import PPO
+from stable_baselines3.common.env_util import make_vec_env
+
+env = make_vec_env("Pendulum-v0", n_envs=4)
+model = PPO("MlpPolicy", env, verbose=1)
+model.learn(total_timesteps=250000)

scratch/etienne/trpo/sb3/sb3-trpo-intersimple.py

@@ -0,0 +1,16 @@
+from sb3_contrib import TRPO
+from intersim.envs import IntersimpleLidarFlat
+from intersim.envs.intersimple import speed_reward
+import functools
+
+env = IntersimpleLidarFlat(
+    n_rays=5,
+    agent=51,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=0
+    ),
+)
+
+model = TRPO("MlpPolicy", env, use_sde=False, sde_sample_freq=4, verbose=1)
+model.learn(total_timesteps=250000)

scratch/etienne/trpo/sb3/sb3-trpo-pendulum.py

@@ -0,0 +1,8 @@
+from sb3_contrib import TRPO
+import gym
+from gym.wrappers import TransformObservation
+
+env = TransformObservation(gym.make('Pendulum-v0'), lambda obs: obs)
+
+model = TRPO("MlpPolicy", env, verbose=1)
+model.learn(total_timesteps=250000)

scratch/johannes/horner_scheme.py

@@ -0,0 +1,90 @@
+# %%
+
+import numpy as np
+import torch
+from timeit import default_timer as timer
+
+# %%
+
+def powerseries(x, deg):
+    return torch.stack([x**i for i in range(deg+1)],dim=-1)
+
+def improved_powerseries(x, deg):
+    r = torch.ones(*x.shape, deg+1, dtype=torch.float64)
+    for i in range(1,deg+1):
+        r[:, :, i] = r[:, :, i-1] * x
+    return r
+
+def horner_scheme(x, poly):
+    deg = poly.shape[-1]
+    nsteps = x.shape[-1]
+    r = poly[:, -1:].repeat(1, nsteps)
+    for i in range(2, deg+1):
+        r *= x
+        r += poly[:, -i:1-i]
+    return r
+
+# %%
+
+nv = 151
+delta = 10
+n = 20
+
+state_s = torch.rand((nv, 1))
+nan_idx = np.random.choice([True, False], 151)
+state_s[nan_idx] = np.nan
+
+# %%
+
+n_coef = 21
+xpoly = torch.rand((nv, n_coef),dtype=torch.float64)
+ypoly = torch.rand((nv, n_coef),dtype=torch.float64)
+
+ds = delta * torch.arange(1,n+1).repeat(nv,1)
+
+s = ds + state_s
+s = s.type(torch.float64)
+
+smax = s[:, 0]
+smax = smax.unsqueeze(-1)
+
+
+start = timer()
+for _ in range(100):
+    deg = xpoly.shape[-1] - 1
+    expand_sims = powerseries(s, deg) # (nv, n, deg+1)
+    # print(expand_sims.shape)
+    y = (ypoly.unsqueeze(1) * expand_sims).sum(dim=-1)
+    x = (xpoly.unsqueeze(1) * expand_sims).sum(dim=-1)
+end = timer()
+print("Powerseries: {}".format((end-start)*1))
+
+start = timer()
+for _ in range(100):
+    deg = xpoly.shape[-1] - 1
+    expand_sims = improved_powerseries(s, deg) # (nv, n, deg+1)
+    # print(expand_sims.shape)
+    yp = (ypoly.unsqueeze(1) * expand_sims).sum(dim=-1)
+    xp = (xpoly.unsqueeze(1) * expand_sims).sum(dim=-1)
+end = timer()
+print("Improved Powerseries: {}".format((end-start)*1))
+
+start = timer()
+for _ in range(100):
+    x_horner = horner_scheme(s, xpoly)
+    y_horner = horner_scheme(s, ypoly)
+end = timer()
+print("Horner: {}".format((end-start)*1))
+
+start = timer()
+for _ in range(100):
+    x_max = horner_scheme(smax, xpoly)
+    y_max = horner_scheme(smax, ypoly)
+end = timer()
+# print("Horner smax: {}".format((end-start)*1))
+
+assert np.all(np.isclose(xp,x)[~nan_idx])
+assert np.all(np.isclose(yp,y)[~nan_idx])
+assert np.all(np.isclose(x_horner,x)[~nan_idx])
+assert np.all(np.isclose(y_horner,y)[~nan_idx])
+# %%

scratch/johannes/intersimple/idm.py

@@ -0,0 +1,42 @@
+# %%
+import torch
+from src.baselines.rule_policies import IDMRulePolicy
+from tqdm import tqdm
+
+from intersim.envs import NRasterizedIncrementingAgent, NRasterizedRandomAgent, NRasterized
+from intersim.envs.intersimple import speed_reward
+import functools
+
+
+env = NRasterizedIncrementingAgent(
+    # agent = 4,
+    reward=functools.partial(
+        speed_reward,
+        collision_penalty=1000
+    ),
+    stop_on_collision=True,
+)
+policy = IDMRulePolicy(env)
+
+colliding_agents = []
+
+for agent in range(151):
+    print("Start agent", agent)
+    obs = env.reset()
+    env.render(mode='post')
+    for i in range(300):
+        action, _ = policy.predict(torch.tensor(obs))
+        # action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+        obs, reward, done, _ = env.step(action)
+        env.render(mode='post')
+        # print('step', i, 'reward', reward)
+        if done:
+            if reward < -500:
+                collising_agents.append(agent)
+                print("  Collision")
+            break
+    env.close(filestr='idm/agent_{}'.format(agent))
+
+print(len(colliding_agents), "colliding_agents")
+print(colliding_agents)
+# %%

src/__init__.py

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

src/baselines/__init__.py

@@ -0,0 +1 @@
+from src.baselines.rule_policies import IDMRulePolicy, PControllerPolicy

src/baselines/rule_policies.py

@@ -0,0 +1,274 @@
+from stable_baselines3.common.base_class import BaseAlgorithm
+from intersim.envs.intersimple import Intersimple, NormalizedActionSpace
+from typing import Tuple, Optional, List
+import numpy as np
+
+class PControllerPolicy(BaseAlgorithm):
+
+
+    def __init__(self, env):
+        """
+        Initialize policy with pointer to environment it will run on
+        """
+        assert isinstance(env, Intersimple), 'Environment is not an intersimple environment'
+        self._env = env
+        self.target_v = 8.94 # m/s
+        self.attn_weight = 20
+
+    # BaseAlgorithm abstract methods
+    def _setup_model(self):
+        return None
+    def learn(self, *args, **kwargs):
+        return self
+
+    def predict(self, observation: np.ndarray, *args, **kwargs):
+        """
+        Generate action, state from observation
+
+        (But actually generate next action from underlying environment state)
+
+        Args:
+            observation (np.ndarray): instantaneous observation from environment
+
+        Returns
+            action (np.ndarray): action for controlled agent to take
+            state (np.ndarray): hidden state for use in next prediction (null)
+        """
+        agent = self._env._agent
+        ego_state = self._env._env.projected_state[agent].numpy() # (5,) tensor
+
+
+        # relative_state = np.delete(self._env._env.relative_state[agent].numpy(), agent, axis=0) #(nv-1, 6) tensor
+
+        # calculate front and left distances from ego
+
+        # calculate relative speed in direction of position difference vector
+
+        # calculate angle alpha and distance d of vehicle i from ego heading
+
+        # attn[i] ~= exp( -(alpha[i])^2 - .01 * d[i] - .1 * vrel[i]
+
+
+        # Proportional controller
+        # action = (self.target_v  - self.attn_weight * attn.sum()) - ego_state[2]
+        action = self.target_v  - ego_state[2]
+        return action, None
+
+class IDMRulePolicy(BaseAlgorithm):
+    """
+    IDMRulePolicy returns action predictions based on an IDM policy.
+
+    The front car is chosen as the closer of:
+        - closest car within a 45 degree half angle cone of the ego's heading
+        - ''' after propagating the environment forward by `t_future' seconds with
+            current headings and velocities
+
+    """
+
+    def __init__(self, env: Intersimple,
+        target_speed:float= 8.94,
+        t_future:List[float]=[0., 1., 2., 3.],
+        half_angle:float=60.):
+        """
+        Initialize policy with pointer to environment it will run on and target speed
+
+        Args:
+            env (Intersimple): intersimple environment which IDM runs on
+            target_speed (float): target speed in roundabout (default: 8.94=20 mph)
+            t_future (List[float]): list of future time at which to compare closest vehicle
+            half_angle (float): half angle to look inside for closest vehicle
+        """
+
+        self._env = env
+        self.t_future = t_future
+        self.half_angle = half_angle
+
+        # Default IDM parameters
+        assert target_speed>0, 'negative target speed'
+        self.v_max = target_speed
+        self.a_max = np.array([3.]) # nominal acceleration
+        self.tau = 0.5 # desired time headway
+        self.b_pref = 2.5 # preferred deceleration
+        self.d_min = 1 #minimum spacing
+
+        # for np.remainder nan warnings
+        np.seterr(invalid='ignore')
+
+    # BaseAlgorithm abstract methods
+    def _setup_model(self):
+        return None
+    def learn(self, *args, **kwargs):
+        return self
+
+    def predict(self, observation:np.ndarray,
+        *args, **kwargs) -> Tuple[np.ndarray, None]:
+        """
+        Predict action, state from observation
+
+        (But actually generate next action from underlying environment state)
+
+        Args:
+            observation (np.ndarray): instantaneous observation from environment
+
+        Returns
+            action (np.ndarray): action for controlled agent to take
+            state (None): None (hidden state for a recurrent policy)
+        """
+        return self.forward(observation, *args, **kwargs), None
+
+    def forward(self, *args, **kwargs) -> np.ndarray:
+        """
+        Generate action from underlying environment
+
+        Returns
+            action (np.ndarray): action for controlled agent to take
+        """
+        agent = self._env._agent
+        state = self._env._env.state.numpy()
+        full_state = self._env._env.projected_state.numpy() #(nv, 5)
+        ego_state = full_state[agent] # (5,)
+        v_ego = ego_state[2]
+        # xy = full_state[:,0:2] # (nv, 2)
+        v = full_state[:,2:3] # (nv, 1)
+        # psi = full_state[:,3:4] # (nv, 1)
+
+        length = 20
+        step = 0.5
+        x, y = self._env._env._generate_paths(delta=step, n=length/step, is_distance=True)
+        heading = to_circle(np.arctan2(np.diff(y), np.diff(x)))
+        velocities = state[:,1]
+
+        # something like this could be done to also take future proximity of vehicles to ego path into account
+        # time_horizon = np.array(range(3))
+        # predictions = state[:,0:1] + np.outer(state[:,1], time_horizon)
+
+        paths = np.stack([x[:,:-1],y[:,:-1], heading], axis=1) # (nv x 3 x (path_length-1))
+        ego_path = paths[agent:agent+1] # (1 x 3 x path_length-1)
+
+        # (x,y,phi) of all vehicles
+        poses = np.expand_dims(full_state[:, [0,1,3]], 2) # (nv x 3 x 1)
+
+        diff = ego_path - poses
+        diff[:, 2, :] = to_circle(diff[:, 2, :])
+
+        # Test if position and heading angle are close for some point on the future vehicle track
+        max_pos_error = 1
+        pos_close = np.sum(diff[:, 0:2, :]**2, 1) <= max_pos_error**2 # (nv x path_length-1)
+        max_deg_error = 20
+        heading_close = np.abs(diff[:, 2, :]) <= 20 * np.pi / 180 # (nv x path_length-1)
+        # For all vehicles get the path points where they are close to the ego path
+        close = np.logical_and(pos_close, heading_close) # (nv x path_length-1)
+        close[agent, :] = False # exclude ego agent
+
+        leader = agent
+        min_idx = np.Inf
+        # Determine vehicle that is closest to ego in terms of path coordinate
+        for veh_id in range(len(close)):
+            path_idx = np.nonzero(close[veh_id])[0]
+            # veh_id is never close to agent
+            if len(path_idx) == 0:
+                continue
+            # first path index where veh_id is close to agent
+            elif path_idx[0] < min_idx:
+                leader = veh_id
+                min_idx = path_idx[0]
+
+        # alternative vectorized code
+        # def findfirst(a):
+        #     idx = np.argwhere(a)
+        #     if len(idx) == 0:
+        #         return np.NaN
+        #     else:
+        #         return float(idx[0]) # float conversion, to get a numpy array of dtype=float64
+        # d = np.apply_along_axis(findfirst, 1, close) # (nv)
+        # if np.all(np.isnan(d)):
+        #     leader = agent
+        # else:
+        #     leader = np.nanargmin(d)
+        #     path_idx = d[leader]
+        #     min_idx = np.sqrt(np.sum(diff[leader, 0:2, path_idx]**2))
+
+
+        if leader != agent:
+            # distance along ego path to point with closest distance
+            d = step * min_idx
+            # add distance from ego path point with closest distance to actual vehicle position
+            d += np.sqrt(np.sum(diff[leader, 0:2, min_idx]**2))
+
+            # Update environment interaction graph with leader
+            self._env._env._graph._neighbor_dict={agent:[leader]}
+
+            delta_v = v_ego - v[leader, 0]
+            d_des = self.d_min + self.tau * v_ego + v_ego * delta_v / (2* (self.a_max*self.b_pref)**0.5 )
+            d_des = max(d_des, self.d_min)
+        else:
+            d = np.Inf
+            d_des = self.d_min
+            self._env._env._graph._neighbor_dict={}
+
+        assert (d_des>= self.d_min)
+        action = self.a_max*(1 - (v_ego/self.v_max)**4 - (d_des/d)**2)
+
+        # normalize action to range if env is a NormalizedActionSpace
+        if isinstance(self._env, NormalizedActionSpace):
+            action = self._env._normalize(action)
+
+        assert action.shape==(1,)
+        return action
+
+
+    def get_ego_dr(self, agent:int, xy: np.ndarray,
+        v: np.ndarray, psi: np.ndarray) -> Tuple[float, float, Optional[int]]:
+        """
+        Return distance and relative speed of closest car within half angle from heading
+
+        Args:
+            agent (int): agent index
+            xy (np.ndarray): (nv, 2) x and y positions
+            v (np.ndarray): (nv, 1) velocity
+            psi (np.ndarray): (nv, 1) heading angle
+
+        Returns:
+            d (float): distance to closest vehicle in cone
+            r (float): relative speed between the two vehicles
+            i (Optional[int]): index of closest vehicle, or None
+        """
+        nv, nxy = xy.shape
+        nv2, nvel = v.shape
+        nv3, npsi = psi.shape
+        assert nv==nv2==nv3
+        assert nxy==2
+        assert nvel==npsi==1
+
+        dxys = xy - xy[agent] # (nv, 2)
+        ds = np.linalg.norm(dxys,axis=1) # (nv,)
+        df = (dxys*np.hstack((np.cos(psi),np.sin(psi)))).sum(-1) # (nv, )
+        dl = (dxys*np.hstack((-np.sin(psi), np.cos(psi)))).sum(-1) # (nv, )
+        alpha = to_circle(np.arctan2(dl, df))
+
+        val_idx = np.arange(nv)[(np.abs(alpha) < self.half_angle*np.pi/180) & (np.arange(nv) != agent)]
+
+        if len(val_idx)==0:
+            i = None
+            d = float('inf')
+            r = float('inf')
+        else:
+            idx = np.argmin(ds[val_idx]) # closest car which meets requirements
+            i = int(val_idx[idx])
+            d = ds[i]
+            r = v[i,0]-v[agent,0]
+
+        return d, r, i
+
+def to_circle(x: np.ndarray) -> np.ndarray:
+    """
+    Casts x (in rad) to [-pi, pi)
+
+    Args:
+        x (np.ndarray): (*) input angle (radians)
+
+    Returns:
+        y (np.ndarray): (*) x cast to [-pi, pi)
+    """
+    y = np.remainder(x + np.pi, 2*np.pi) - np.pi
+    return y

src/core/discriminator.py

@@ -0,0 +1,74 @@
+import torch
+import torch.nn as nn
+
+class Discriminator(nn.Module):
+    
+    def __init__(self):
+        super().__init__()
+        self.nn = nn.Sequential(
+            nn.LazyLinear(50),
+            nn.Tanh(),
+            nn.LazyLinear(50),
+            nn.Tanh(),
+            nn.LazyLinear(1),
+        )
+    
+    def forward(self, states, actions):
+        return self.nn(torch.cat((states, actions), dim=-1)).squeeze(-1)
+
+class DeepsetDiscriminator(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.elem = nn.Sequential(
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(10),
+        )
+        self.glob = nn.Sequential(
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(1),
+        )
+    
+    def forward(self, states, actions):
+        actions = actions.unsqueeze(-2)
+        actions = actions.expand(*actions.shape[:-2], states.shape[-2], actions.shape[-1])
+        sa = torch.cat((states, actions), dim=-1)
+        return self.glob(self.elem(sa).sum(-2)).squeeze(-1)
+
+class RecurrentDiscriminator(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.state_dim = 10
+        self.state = nn.Sequential(
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(self.state_dim),
+        )
+        self.glob = nn.Sequential(
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(1),
+        )
+    
+    def forward(self, states, actions):
+        actions = actions.unsqueeze(-2)
+        batch_size = actions.shape[:-2]
+        set_size = states.shape[-2]
+        action_dim = actions.shape[-1]
+        actions = actions.expand(*batch_size, set_size, action_dim)
+        sa = torch.cat((states, actions), dim=-1)
+
+        state = torch.zeros((*batch_size, self.state_dim))
+        for i in range(set_size):
+            state = state + self.state(torch.cat((state, sa[..., i, :]), dim=-1))
+
+        return self.glob(state).squeeze(-1)

src/core/gail.py

@@ -0,0 +1,125 @@
+import torch
+import torch.nn.functional as F
+from dataclasses import dataclass
+from src.core.reparam_module import ReparamPolicy
+from src.core.sampling import rollout
+from src.core.trpo import trpo_step
+from src.core.ppo import ppo_step
+from tqdm import tqdm
+
+class TerminalLogger:
+    def add_scalar(self, key, scalar, i=None):
+        if i is not None:
+            print('Iteration', i, end=' ')
+        print(key, scalar)
+
+@dataclass
+class Buffer:
+    states: torch.Tensor
+    actions: torch.Tensor
+    rewards: torch.Tensor
+    dones: torch.Tensor
+
+def roll_buffer(buffer, *args, **kwargs):
+    return Buffer(
+        torch.roll(buffer.states, *args, **kwargs),
+        torch.roll(buffer.actions, *args, **kwargs),
+        torch.roll(buffer.rewards, *args, **kwargs),
+        torch.roll(buffer.dones, *args, **kwargs),
+    )
+
+def gail(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, value,
+         v_opt, v_iters, epochs, rollout_episodes, rollout_steps, gamma,
+         gae_lambda, delta, backtrack_coeff, backtrack_iters, cg_iters=10, cg_damping=0.1, wasserstein=False, wasserstein_c=None, logger=TerminalLogger()):
+
+    policy(torch.zeros(env_fn(0).observation_space.shape))
+    policy = ReparamPolicy(policy)
+
+    logger.add_scalar('expert/mean_episode_length', (~expert_data.dones).sum() / expert_data.states.shape[0])
+    logger.add_scalar('expert/mean_reward_per_episode', expert_data.rewards[~expert_data.dones].sum() / expert_data.states.shape[0])
+
+    for epoch in tqdm(range(epochs)):
+        generator_data = Buffer(*rollout(env_fn, policy, rollout_episodes, rollout_steps))
+
+        logger.add_scalar('gen/mean_episode_length', (~generator_data.dones).sum() / generator_data.states.shape[0], epoch)
+        logger.add_scalar('gen/mean_reward_per_episode', generator_data.rewards[~generator_data.dones].sum() / generator_data.states.shape[0], epoch)
+
+        discriminator, loss = train_discriminator(expert_data, generator_data, discriminator, disc_opt, disc_iters, wasserstein, wasserstein_c)
+        if wasserstein:
+            generator_data.rewards = discriminator(generator_data.states, generator_data.actions)
+        else:
+            generator_data.rewards = -F.logsigmoid(discriminator(generator_data.states, generator_data.actions))
+        logger.add_scalar('disc/final_loss', loss, epoch)
+        logger.add_scalar('disc/mean_reward_per_episode', generator_data.rewards[~generator_data.dones].sum() / generator_data.states.shape[0], epoch)
+
+        value, policy = trpo_step(value, policy, generator_data.states, generator_data.actions, generator_data.rewards, generator_data.dones, gamma, gae_lambda, delta, backtrack_coeff, backtrack_iters, v_opt, v_iters, cg_iters, cg_damping)
+        expert_data = roll_buffer(expert_data, shifts=-3, dims=0)
+    
+    return value, policy
+
+def gail_ppo(env_fn, expert_data, discriminator, disc_opt, disc_iters, policy, value,
+         v_opt, v_iters, epochs, rollout_episodes, rollout_steps, gamma,
+         gae_lambda, clip_ratio, pi_opt, pi_iters, target_kl=None, max_grad_norm=None, wasserstein=False, wasserstein_c=None, logger=TerminalLogger()):
+
+    logger.add_scalar('expert/mean_episode_length', (~expert_data.dones).sum() / expert_data.states.shape[0])
+    logger.add_scalar('expert/mean_reward_per_episode', expert_data.rewards[~expert_data.dones].sum() / expert_data.states.shape[0])
+
+    for epoch in range(epochs):
+        generator_data = Buffer(*rollout(env_fn, policy, rollout_episodes, rollout_steps))
+
+        logger.add_scalar('gen/mean_episode_length', (~generator_data.dones).sum() / generator_data.states.shape[0], epoch)
+        logger.add_scalar('gen/mean_reward_per_episode', generator_data.rewards[~generator_data.dones].sum() / generator_data.states.shape[0], epoch)
+
+        discriminator, loss = train_discriminator(expert_data, generator_data, discriminator, disc_opt, disc_iters, wasserstein, wasserstein_c)
+        if wasserstein:
+            generator_data.rewards = discriminator(generator_data.states, generator_data.actions)
+        else:
+            generator_data.rewards = -F.logsigmoid(discriminator(generator_data.states, generator_data.actions))
+        logger.add_scalar('disc/final_loss', loss, epoch)
+        logger.add_scalar('disc/mean_reward_per_episode', generator_data.rewards[~generator_data.dones].sum() / generator_data.states.shape[0], epoch)
+
+        value, policy = ppo_step(value, policy, generator_data.states, generator_data.actions, generator_data.rewards, generator_data.dones, clip_ratio, gamma, gae_lambda, pi_opt, pi_iters, v_opt, v_iters, target_kl, max_grad_norm)
+        expert_data = roll_buffer(expert_data, shifts=-3, dims=0)
+    
+    return value, policy
+
+def train_discriminator(expert_data, generator_data, discriminator, disc_opt, disc_iters, wasserstein, wasserstein_c=None):
+
+    n_expert_samples = (~expert_data.dones).sum()
+    n_generator_samples = (~generator_data.dones).sum()
+    n_samples = torch.minimum(n_expert_samples, n_generator_samples)
+
+    gen_states = generator_data.states[~generator_data.dones][:n_samples]
+    gen_actions = generator_data.actions[~generator_data.dones][:n_samples]
+    exp_states = expert_data.states[~expert_data.dones][:n_samples]
+    exp_actions = expert_data.actions[~expert_data.dones][:n_samples]
+
+    states = torch.cat((exp_states, gen_states), dim=0).detach()
+    actions = torch.cat((exp_actions, gen_actions), dim=0).detach()
+    labels = torch.cat((torch.zeros(n_samples), torch.ones(n_samples))).detach()
+
+    # print('Batch augmentation on')
+    # random_states = torch.rand_like(gen_states)
+    # random_actions = torch.rand_like(gen_actions)
+    # states = torch.cat((exp_states, gen_states, random_states), dim=0).detach()
+    # actions = torch.cat((exp_actions, gen_actions, random_actions), dim=0).detach()
+    # labels = torch.cat((torch.zeros(n_samples), torch.ones(n_samples), torch.ones(n_samples))).detach()
+
+    for _ in range(disc_iters):
+        disc_opt.zero_grad()
+        pred = discriminator(states, actions)
+
+        if wasserstein:
+            loss = -(pred * (1 - labels) - pred * labels).mean()
+        else:
+            loss = F.binary_cross_entropy(torch.sigmoid(pred), labels)
+        
+        loss.backward()
+        disc_opt.step()
+
+        if wasserstein_c is not None:
+            with torch.no_grad():
+                for param in discriminator.parameters():
+                    param.clamp_(-wasserstein_c, wasserstein_c)
+    
+    return discriminator, loss

src/core/optimization.py

@@ -0,0 +1,39 @@
+import torch
+
+def conjugate_gradient(A, b, max_iters, res_tol=1e-10):
+    x = torch.zeros_like(b)
+    r = b - A(x)
+    p = r
+
+    rTr = r.T @ r
+
+    for _ in range(max_iters):
+        Ap = A(p)
+        alpha = rTr / (p.T @ Ap)
+        x = x + alpha * p
+
+        r = r - alpha * Ap
+        if torch.norm(r) < res_tol:
+            break
+
+        rTrnew = r.T @ r
+        beta = rTrnew / rTr
+        p = r + beta * p
+        rTr = rTrnew
+
+    return x
+
+def line_search(f, x0, dx, g0, alpha, condition, max_steps=10, c1=0.1):
+    assert 0 < alpha < 1
+    
+    f0 = f(x0)
+    for _ in range(max_steps):
+        x = x0 + dx
+
+        if (f(x) > f0 + c1 * g0.T @ dx) and condition(x):
+            return x
+        
+        dx *= alpha
+    
+    print('Line search failed, returning x0')
+    return x0

src/core/policy.py

@@ -0,0 +1,111 @@
+import torch
+import torch.nn as nn
+from torch.distributions import Independent, Normal, Categorical
+from torch.distributions.kl import kl_divergence
+
+class BasePolicy(nn.Module):
+    
+    def __init__(self, action_dim):
+        super().__init__()
+        self.action_dim = action_dim
+
+    def torch_dist(self, dist):
+        return Independent(Normal(dist[..., :self.action_dim], dist[..., self.action_dim:].exp()), 1)
+    
+    def sample(self, dist):
+        return self.torch_dist(dist).sample()
+    
+    def predict(self, observations, state=None, episode_start=None, deterministic=True):
+        observations = torch.tensor(observations)
+        return self._predict(self.forward(observations), state, episode_start, deterministic)
+    
+    def _predict(self, dist, state=None, episode_start=None, deterministic=True):
+        if deterministic:
+            actions = dist[..., :self.action_dim]
+        else:
+            actions = self.sample(dist)
+        return actions, None
+
+    def log_prob(self, dist, actions):
+        return self.torch_dist(dist).log_prob(actions)
+    
+    def kl_divergence(self, dist1, dist2):
+        d1 = self.torch_dist(dist1)
+        d2 = self.torch_dist(dist2)
+        return kl_divergence(d1, d2)
+
+class Policy(BasePolicy):
+    
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.nn = nn.Sequential(
+            nn.LazyLinear(50),
+            nn.Tanh(),
+            nn.LazyLinear(50),
+            nn.Tanh(),
+            nn.LazyLinear(2 * self.action_dim),
+        )
+
+    def forward(self, states):
+        return self.nn(states)
+
+class DiscretePolicy(BasePolicy):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.nn = nn.Sequential(
+            nn.LazyLinear(50),
+            nn.Tanh(),
+            nn.LazyLinear(50),
+            nn.Tanh(),
+            nn.LazyLinear(self.action_dim),
+        )
+
+    def forward(self, states):
+        return self.nn(states)
+    
+    def torch_dist(self, dist):
+        return Categorical(logits=dist)
+    
+    def _predict(self, dist, state=None, episode_start=None, deterministic=True):
+        if deterministic:
+            _, actions = dist.max(-1)
+        else:
+            actions = self.sample(dist)
+        return actions, None
+
+class SetPolicy(Policy):
+
+    def forward(self, states):
+        batch_size = states.shape[:-2]
+        states = torch.cat((states[..., :1, [0, 1]], states[..., :, [2, 5]]), axis=-2).reshape(*batch_size, -1)
+        return super().forward(states)
+
+class SetDiscretePolicy(DiscretePolicy):
+
+    def forward(self, states):
+        batch_size = states.shape[:-2]
+        states = torch.cat((states[..., :1, [0, 1]], states[..., :, [2, 5]]), axis=-2).reshape(*batch_size, -1)
+        return super().forward(states)
+
+class DeepSetPolicy(BasePolicy):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.elem = nn.Sequential(
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(10),
+        )
+        self.glob = nn.Sequential(
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(2 * self.action_dim),
+        )
+    
+    def forward(self, states):
+        return self.glob(self.elem(states).sum(-2))

src/core/ppo.py

@@ -0,0 +1,72 @@
+import torch
+from src.core.sampling import rollout
+from src.core.value_estimation import gae
+
+def ppo(env_fn, value, policy, epochs, rollout_episodes, rollout_steps, gamma, gae_lambda, clip_ratio, pi_opt, pi_iters, v_opt, v_iters, target_kl=None, max_grad_norm=None):
+
+    for epoch in range(epochs):
+        policy.eval()
+        states, actions, rewards, dones = rollout(env_fn, policy, rollout_episodes, rollout_steps)
+
+        print('mean', states[~dones].mean(0))
+        print('std', states[~dones].std(0))
+
+        print(f'Iteration {epoch} mean episode length {(~dones).sum() / states.shape[0]}')
+        print(f'Iteration {epoch} mean reward per episode {rewards[~dones].sum() / states.shape[0]}')
+
+        policy.train()
+        value.train()
+        value, policy = ppo_step(value, policy, states, actions, rewards, dones, clip_ratio, gamma, gae_lambda, pi_opt, pi_iters, v_opt, v_iters, target_kl, max_grad_norm)
+
+    return value, policy
+
+def ppo_step(value, policy, states, actions, rewards, dones, clip_ratio, gamma, gae_lambda, pi_opt, pi_iters, v_opt, v_iters, target_kl, max_grad_norm):
+
+    states = states.detach()
+    actions = actions.detach()
+    rewards = rewards.detach()
+    dones = dones.detach()
+
+    advantages, returns, valid = gae(states, rewards, value(states), dones, gamma, gae_lambda)
+    advantages = advantages.detach()
+    returns = returns.detach()
+
+    # update value function
+
+    for _ in range(v_iters):
+        v_opt.zero_grad()
+        value_loss = (value(states) - returns).pow(2)[valid].mean()
+        value_loss.backward()
+        v_opt.step()
+
+    # update policy
+
+    old_dist = policy(states).detach()
+    old_logprob = policy.log_prob(old_dist, actions).detach()
+
+    def g(advantages, clip_ratio):
+        return torch.where(advantages >= 0, (1 + clip_ratio) * advantages, (1 - clip_ratio) * advantages)
+
+    def L(states, actions, advantages, clip_ratio):
+        return torch.minimum(
+            (policy.log_prob(policy(states), actions) - old_logprob).exp() * advantages,
+            g(advantages, clip_ratio)
+        )[valid].mean()
+
+    for _ in range(pi_iters):
+        pi_opt.zero_grad()
+        ppo_loss = -L(states, actions, advantages, clip_ratio)
+        ppo_loss.backward()
+
+        if max_grad_norm:
+            torch.nn.utils.clip_grad_norm(policy.parameters(), max_grad_norm)
+        
+        pi_opt.step()
+
+        kl = policy.kl_divergence(policy(states), old_dist)[valid].mean()
+        if target_kl and kl > target_kl:
+            break
+    
+    print('KL', kl.item())
+
+    return value, policy

src/core/reparam_module.py

@@ -0,0 +1,173 @@
+# Source: https://github.com/SsnL/PyTorch-Reparam-Module
+
+import torch
+import torch.nn as nn
+import warnings
+import types
+from collections import namedtuple
+from contextlib import contextmanager
+
+class ReparamModule(nn.Module):
+    def __init__(self, module):
+        super(ReparamModule, self).__init__()
+        self.module = module
+
+        param_infos = []
+        shared_param_memo = {}
+        shared_param_infos = []
+        params = []
+        param_numels = []
+        param_shapes = []
+        for m in self.modules():
+            for n, p in m.named_parameters(recurse=False):
+                if p is not None:
+                    if p in shared_param_memo:
+                        shared_m, shared_n = shared_param_memo[p]
+                        shared_param_infos.append((m, n, shared_m, shared_n))
+                    else:
+                        shared_param_memo[p] = (m, n)
+                        param_infos.append((m, n))
+                        params.append(p.detach())
+                        param_numels.append(p.numel())
+                        param_shapes.append(p.size())
+
+        assert len(set(p.dtype for p in params)) <= 1, \
+            "expects all parameters in module to have same dtype"
+
+        # store the info for unflatten
+        self._param_infos = tuple(param_infos)
+        self._shared_param_infos = tuple(shared_param_infos)
+        self._param_numels = tuple(param_numels)
+        self._param_shapes = tuple(param_shapes)
+
+        # flatten
+        flat_param = nn.Parameter(torch.cat([p.reshape(-1) for p in params], 0))
+        self.register_parameter('flat_param', flat_param)
+        self.param_numel = flat_param.numel()
+        del params
+        del shared_param_memo
+
+        # deregister the names as parameters
+        for m, n in self._param_infos:
+            delattr(m, n)
+        for m, n, _, _ in self._shared_param_infos:
+            delattr(m, n)
+
+        # register the views as plain attributes
+        self._unflatten_param(self.flat_param)
+
+        # now buffers
+        # they are not reparametrized. just store info as (module, name, buffer)
+        buffer_infos = []
+        for m in self.modules():
+            for n, b in m.named_buffers(recurse=False):
+                if b is not None:
+                    buffer_infos.append((m, n, b))
+
+        self._buffer_infos = tuple(buffer_infos)
+        self._traced_self = None
+
+    def trace(self, example_input, **trace_kwargs):
+        assert self._traced_self is None, 'This ReparamModule is already traced'
+
+        if isinstance(example_input, torch.Tensor):
+            example_input = (example_input,)
+        example_input = tuple(example_input)
+        example_param = (self.flat_param.detach().clone(),)
+        example_buffers = (tuple(b.detach().clone() for _, _, b in self._buffer_infos),)
+
+        self._traced_self = torch.jit.trace_module(
+            self,
+            inputs=dict(
+                _forward_with_param=example_param + example_input,
+                _forward_with_param_and_buffers=example_param + example_buffers + example_input,
+            ),
+            **trace_kwargs,
+        )
+
+        # replace forwards with traced versions
+        self._forward_with_param = self._traced_self._forward_with_param
+        self._forward_with_param_and_buffers = self._traced_self._forward_with_param_and_buffers
+        return self
+
+    def clear_views(self):
+        for m, n in self._param_infos:
+            setattr(m, n, None)  # This will set as plain attr
+
+    def _apply(self, *args, **kwargs):
+        if self._traced_self is not None:
+            self._traced_self._apply(*args, **kwargs)
+            return self
+        return super(ReparamModule, self)._apply(*args, **kwargs)
+
+    def _unflatten_param(self, flat_param):
+        ps = (t.view(s) for (t, s) in zip(flat_param.split(self._param_numels), self._param_shapes))
+        for (m, n), p in zip(self._param_infos, ps):
+            setattr(m, n, p)  # This will set as plain attr
+        for (m, n, shared_m, shared_n) in self._shared_param_infos:
+            setattr(m, n, getattr(shared_m, shared_n))
+
+    @contextmanager
+    def unflattened_param(self, flat_param):
+        saved_views = [getattr(m, n) for m, n in self._param_infos]
+        self._unflatten_param(flat_param)
+        yield
+        # Why not just `self._unflatten_param(self.flat_param)`?
+        # 1. because of https://github.com/pytorch/pytorch/issues/17583
+        # 2. slightly faster since it does not require reconstruct the split+view
+        #    graph
+        for (m, n), p in zip(self._param_infos, saved_views):
+            setattr(m, n, p)
+        for (m, n, shared_m, shared_n) in self._shared_param_infos:
+            setattr(m, n, getattr(shared_m, shared_n))
+
+    @contextmanager
+    def replaced_buffers(self, buffers):
+        for (m, n, _), new_b in zip(self._buffer_infos, buffers):
+            setattr(m, n, new_b)
+        yield
+        for m, n, old_b in self._buffer_infos:
+            setattr(m, n, old_b)
+
+    def _forward_with_param_and_buffers(self, flat_param, buffers, *inputs, **kwinputs):
+        with self.unflattened_param(flat_param):
+            with self.replaced_buffers(buffers):
+                return self.module(*inputs, **kwinputs)
+
+    def _forward_with_param(self, flat_param, *inputs, **kwinputs):
+        with self.unflattened_param(flat_param):
+            return self.module(*inputs, **kwinputs)
+
+    def forward(self, *inputs, flat_param=None, buffers=None, **kwinputs):
+        if flat_param is None:
+            flat_param = self.flat_param
+        if buffers is None:
+            return self._forward_with_param(flat_param, *inputs, **kwinputs)
+        else:
+            return self._forward_with_param_and_buffers(flat_param, tuple(buffers), *inputs, **kwinputs)
+
+
+class ReparamPolicy(ReparamModule):
+
+    def sample(self, *args, **kwargs):
+        return self.module.sample(*args, **kwargs)
+    
+    def log_prob(self, *args, **kwargs):
+        return self.module.log_prob(*args, **kwargs)
+    
+    def kl_divergence(self, *args, **kwargs):
+        return self.module.kl_divergence(*args, **kwargs)
+    
+    def predict(self, obs, *args, **kwargs):
+        obs = torch.tensor(obs)
+        return self.module._predict(self.forward(obs), *args, **kwargs)
+    
+    def unsafe_probability_mass(self, *args, **kwargs):
+        return self.module.unsafe_probability_mass(*args, **kwargs)
+
+class ReparamSafePolicy(ReparamPolicy):
+
+    def predict(self, obs, *args, **kwargs):
+        observation = torch.tensor(obs['observation'])
+        safe_actions = torch.tensor(obs['safe_actions'])
+        return self.module._predict(self.forward(observation, safe_actions), *args, **kwargs)

src/core/sampling.py

@@ -0,0 +1,73 @@
+import torch
+import gym
+from stable_baselines3.common.vec_env import DummyVecEnv as VecEnv
+from tqdm import tqdm
+
+def rollout(env_fn, policy, n_episodes, max_steps_per_episode):
+    env = env_fn(0)
+    states = torch.zeros(n_episodes, max_steps_per_episode + 1, *env.observation_space.shape)
+    actions = torch.zeros(n_episodes, max_steps_per_episode + 1, *env.action_space.shape)
+    rewards = torch.zeros(n_episodes, max_steps_per_episode + 1)
+    dones = torch.ones(n_episodes, max_steps_per_episode + 1, dtype=bool)
+
+    env = VecEnv(list(map(lambda i: (lambda: env_fn(i)), range(n_episodes))))
+
+    states[:, 0] = torch.tensor(env.reset()).clone().detach()
+    dones[:, 0] = False
+
+    for s in range(max_steps_per_episode):
+        actions[:, s] = policy.sample(policy(states[:, s])).clone().detach()
+
+        clipped_actions = actions[:, s]
+        if isinstance(env.action_space, gym.spaces.Box):
+            clipped_actions = torch.clamp(clipped_actions, torch.from_numpy(env.action_space.low), torch.from_numpy(env.action_space.high))
+
+        o, r, d, _ = env.step(clipped_actions)
+        states[:, s + 1] = torch.tensor(o).clone().detach()
+        rewards[:, s] = torch.tensor(r).clone().detach()
+        dones[:, s + 1] = torch.tensor(d).clone().detach()
+
+    dones = dones.cumsum(1) > 0
+
+    states = states[:, :max_steps_per_episode]
+    actions = actions[:, :max_steps_per_episode]
+    rewards = rewards[:, :max_steps_per_episode]
+    dones = dones[:, :max_steps_per_episode]
+
+    return states, actions, rewards, dones
+
+
+def rollout_sb3(env, policy, n_episodes, max_steps_per_episode):
+    states = torch.zeros(n_episodes, max_steps_per_episode + 1, *env.observation_space.shape)
+    actions = torch.zeros(n_episodes, max_steps_per_episode + 1, *env.action_space.shape)
+    rewards = torch.zeros(n_episodes, max_steps_per_episode + 1)
+    dones = torch.ones(n_episodes, max_steps_per_episode + 1, dtype=bool)
+
+    for e in tqdm(range(n_episodes)):
+        states[e, 0] = torch.tensor(env.reset()).clone().detach()
+        dones[e, 0] = False
+
+        for s in range(max_steps_per_episode):
+            action, _ = policy.predict(states[e, s])
+            actions[e, s] = torch.tensor(action).clone().detach()
+
+            clipped_actions = actions[e, s]
+            if isinstance(env.action_space, gym.spaces.Box):
+                clipped_actions = torch.clamp(clipped_actions, torch.from_numpy(env.action_space.low), torch.from_numpy(env.action_space.high))
+
+            o, r, d, _ = env.step(clipped_actions)
+            states[e, s + 1] = torch.tensor(o).clone().detach()
+            rewards[e, s] = torch.tensor(r).clone().detach()
+            dones[e, s + 1] = torch.tensor(d).clone().detach()
+
+            if d:
+                break
+
+    dones = dones.cumsum(1) > 0
+
+    states = states[:, :max_steps_per_episode]
+    actions = actions[:, :max_steps_per_episode]
+    rewards = rewards[:, :max_steps_per_episode]
+    dones = dones[:, :max_steps_per_episode]
+
+    return states, actions, rewards, dones

src/core/test_optimization.py

@@ -0,0 +1,23 @@
+import torch
+from optimization import conjugate_gradient
+
+def test_cg_eye():
+    A = torch.eye(2)
+    b = torch.tensor([1., 2.])
+    x1 = conjugate_gradient(lambda x: A @ x, b, 2)
+    x2 = torch.inverse(A) @ b
+    assert torch.allclose(x1, x2)
+
+def test_cg_eyep1():
+    A = torch.eye(2) + 1
+    b = torch.tensor([1., 2.])
+    x1 = conjugate_gradient(lambda x: A @ x, b, 2)
+    x2 = torch.inverse(A) @ b
+    assert torch.allclose(x1, x2, atol=1e-7)
+
+def test_cg3():
+    A = torch.tensor([[4., 2.], [2., 4.]])
+    b = torch.tensor([2., 1.])
+    x1 = conjugate_gradient(lambda x: A @ x, b, 100)
+    x2 = torch.inverse(A) @ b
+    assert torch.allclose(x1, x2)

src/core/trpo.py

@@ -0,0 +1,79 @@
+import torch
+from src.core.reparam_module import ReparamPolicy
+from src.core.sampling import rollout
+from src.core.value_estimation import gae
+from src.core.optimization import conjugate_gradient, line_search
+
+def trpo(env_fn, value, policy, epochs, rollout_episodes, rollout_steps, gamma, gae_lambda, delta, backtrack_coeff, backtrack_iters, v_opt, v_iters, cg_iters=10, cg_damping=0.1):
+
+    policy(torch.zeros(env_fn(0).observation_space.shape))
+    policy = ReparamPolicy(policy)
+
+    for epoch in range(epochs):
+        policy.eval()
+        states, actions, rewards, dones = rollout(env_fn, policy, rollout_episodes, rollout_steps)
+
+        print('mean', states[~dones].mean(0))
+        print('std', states[~dones].std(0))
+
+        print(f'Iteration {epoch} mean episode length {(~dones).sum() / states.shape[0]}')
+        print(f'Iteration {epoch} mean reward per episode {rewards[~dones].sum() / states.shape[0]}')
+
+        policy.train()
+        value.train()
+        value, policy = trpo_step(value, policy, states, actions, rewards, dones, gamma, gae_lambda, delta, backtrack_coeff, backtrack_iters, v_opt, v_iters, cg_iters, cg_damping)
+
+    return value, policy
+
+def trpo_step(value, policy, states, actions, rewards, dones, gamma, gae_lambda, delta, backtrack_coeff, backtrack_iters, v_opt, v_iters, cg_iters=10, cg_damping=0.1):
+
+    states = states.detach()
+    actions = actions.detach()
+    rewards = rewards.detach()
+    dones = dones.detach()
+
+    advantages, returns, valid = gae(states, rewards, value(states), dones, gamma, gae_lambda)
+    advantages = advantages.detach()
+    returns = returns.detach()
+
+    # update value function
+
+    for _ in range(v_iters):
+        v_opt.zero_grad()
+        value_loss = (value(states) - returns).pow(2)[valid].mean()
+        value_loss.backward()
+        v_opt.step()
+
+    # compute policy gradient
+
+    plogprob = policy.log_prob(policy(states), actions)
+    surrogate_advantage = (plogprob * advantages)[valid].sum() / states.shape[0]
+    g = torch.cat(torch.autograd.grad(surrogate_advantage, policy.flat_param)).detach()
+    
+    def Hx(x):
+        kl = policy.kl_divergence(policy(states), policy(states).detach())[valid].mean()
+        dKL = torch.cat(torch.autograd.grad(kl, policy.flat_param, create_graph=True))
+        H_x = torch.cat(torch.autograd.grad(dKL.T @ x, policy.flat_param)).detach()
+        return H_x + cg_damping * x
+
+    x = conjugate_gradient(Hx, g, cg_iters)
+    npg = torch.sqrt(2 * delta / (x.T @ Hx(x))) * x
+
+    # perform line search
+
+    def L(theta):
+        rplogprob = policy.log_prob(policy(states, flat_param=theta), actions)
+        return ((rplogprob - plogprob.detach()).exp() * advantages)[valid].sum() / advantages.shape[0]
+
+    condition = lambda theta: policy.kl_divergence(policy(states, flat_param=theta), policy(states))[valid].mean() < delta
+
+    x0 = policy.flat_param
+    g0 = torch.cat(torch.autograd.grad(L(x0), x0))
+    theta = line_search(L, x0, npg, g0, backtrack_coeff, condition, max_steps=backtrack_iters)
+
+    # update policy parameters
+    
+    with torch.no_grad():
+        policy.flat_param.copy_(theta)
+
+    return value, policy

src/core/value.py

@@ -0,0 +1,48 @@
+import torch
+import torch.nn as nn
+from torch.distributions import Normal
+from torch.distributions.kl import kl_divergence
+
+class Value(nn.Module):
+    
+    def __init__(self):
+        super().__init__()
+        self.nn = nn.Sequential(
+            nn.LazyLinear(50),
+            nn.Tanh(),
+            nn.LazyLinear(50),
+            nn.Tanh(),
+            nn.LazyLinear(1),
+        )
+    
+    def forward(self, states):
+        return self.nn(states).squeeze(-1)
+
+class SetValue(Value):
+
+    def forward(self, states):
+        batch_size = states.shape[:-2]
+        states = torch.cat((states[..., :1, [0, 1]], states[..., :, [2, 5]]), axis=-2).reshape(*batch_size, -1)
+        return super().forward(states)
+
+class DeepSetValue(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.elem = nn.Sequential(
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(10),
+        )
+        self.glob = nn.Sequential(
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(10),
+            nn.Tanh(),
+            nn.LazyLinear(1),
+        )
+    
+    def forward(self, states):
+        return self.glob(self.elem(states).sum(-2)).squeeze(-1)

src/core/value_estimation.py

@@ -0,0 +1,40 @@
+from operator import index
+import torch
+
+def gae(states, rewards, values, dones, gamma, gae_lambda):
+    assert rewards.shape == values.shape == dones.shape
+    n_episodes, n_steps = rewards.shape
+
+    valid = ~dones
+    valid[..., -1] = False
+
+    td = rewards + gamma * torch.roll(values, shifts=-1, dims=1) - values
+    adv = td.repeat(n_steps, 1, 1).transpose(0, 1)
+    assert adv.shape == (n_episodes, n_steps, n_steps)
+
+    step_start, step = torch.meshgrid(torch.arange(n_steps), torch.arange(n_steps), indexing='ij')
+    past = step < step_start
+
+    # add up discounted temporal differences 
+    discount = torch.minimum(torch.tensor(gamma).log() * (step - step_start), torch.tensor(0.)).exp()
+    discount = discount * ~past
+    discount = discount * valid.unsqueeze(1)
+
+    adv = adv * discount
+    adv = adv.cumsum(2) # eq. (14)
+    assert adv.shape == (n_episodes, n_steps, n_steps)
+
+    # add up discounted k-advantages
+    lambda_discount = torch.minimum(torch.tensor(gae_lambda).log() * (step - step_start), torch.tensor(0.)).exp()
+    lambda_discount = lambda_discount * ~past
+    lambda_discount = lambda_discount * valid.unsqueeze(1)
+
+    adv = adv * lambda_discount
+    adv = adv.sum(2) / (lambda_discount.sum(2) + 1e-10) # eq. (16)
+
+    adv = (adv - adv[valid].mean()) / adv[valid].std()
+    assert adv.shape == rewards.shape == values.shape
+
+    returns = adv + values
+
+    return adv, returns, valid

src/eval_main.py

@@ -0,0 +1,387 @@
+from tqdm import tqdm
+from copy import deepcopy
+import stable_baselines3 as sb3
+import intersim
+from intersim.envs import Intersimple
+from stable_baselines3.common.base_class import BaseAlgorithm
+from src.baselines import IDMRulePolicy
+from src.evaluation import IntersimpleEvaluation
+import src.gail.options as options_envs
+from src.evaluation.metrics import divergence, visualize_distribution, rwse
+from src.evaluation.utils import save_metrics
+from src.core.policy import SetPolicy, SetDiscretePolicy
+from src.core.reparam_module import ReparamPolicy, ReparamSafePolicy
+from src.options import envs as options_envs2
+from src.safe_options.policy import SetMaskedDiscretePolicy
+from src.safe_options import options as options_envs3
+
+from typing import Optional, List, Dict, Tuple
+import torch
+import numpy as np
+
+#(method, policy_file, policy_kwargs, eval_env)
+
+def load_policy(method:str, 
+    policy_file:str, 
+    policy_kwargs:dict, 
+    env: Intersimple) -> BaseAlgorithm:
+    """
+    Load a model given a path and the method
+    
+    Args:
+        method (str): the method for the model
+        policy_file (str): the path to the model
+        policy_kwargs (str): the path to the model
+        env (Intersimple): Intersimple environment for evaluation (necessary for IDM policy)
+    Returns:
+        policy (Optional[BaseAlgorithm]): the policy to evaluate
+    """
+    if method == 'idm':
+        policy = IDMRulePolicy(env, **policy_kwargs)
+    elif method == 'bc':
+        policy = SetPolicy(env.action_space.shape[-1])
+        policy.load_state_dict(torch.load(policy_file))
+        policy.eval()
+    elif method == 'gail':
+        policy = SetPolicy(env.action_space.shape[-1])
+        policy(torch.zeros(env.observation_space.shape))
+        policy = ReparamPolicy(policy)
+        policy.load_state_dict(torch.load(policy_file))
+        policy.eval()
+    elif method == 'gail-ppo':
+        policy = SetPolicy(env.action_space.shape[-1])
+        policy.load_state_dict(torch.load(policy_file))
+        policy.eval()
+    elif method == 'rail':
+        raise NotImplementedError
+    elif method == 'ogail':
+        policy = SetDiscretePolicy(env.action_space.n)
+        policy(torch.zeros(env.observation_space.shape))
+        policy = ReparamPolicy(policy)
+        policy.load_state_dict(torch.load(policy_file))
+        policy.eval()
+    elif method == 'ogail-ppo':
+        policy = SetDiscretePolicy(env.action_space.n)
+        policy.load_state_dict(torch.load(policy_file))
+        policy.eval()
+    elif method == 'sgail':
+        policy = SetMaskedDiscretePolicy(env.action_space.n)
+        policy(
+            torch.zeros(env.observation_space['observation'].shape),
+            torch.zeros(env.observation_space['safe_actions'].shape)
+        )
+        policy = ReparamSafePolicy(policy)
+        policy.load_state_dict(torch.load(policy_file))
+        policy.eval()
+    elif method == 'sgail-ppo':
+        policy = SetMaskedDiscretePolicy(env.action_space.n)
+        policy.load_state_dict(torch.load(policy_file))
+        policy.eval()
+    else:
+        raise NotImplementedError
+    return policy
+
+def form_expert_metrics(states:torch.Tensor, actions:torch.Tensor) ->Dict[str, list]:
+    """
+    Given experts of tensor states and actions, form a dictionary of metrics
+
+    Args:
+        states (torch.tensor): (T+1, nv, 5) expert states for track file
+        actions (torch.tensor): (T, nv, 1) expert actions for track file
+    
+    Returns:
+        metrics (dict): dictionary maping strings to lists
+    """
+
+    T1, nv, _ = states.shape
+    T, nv2, _ = actions.shape
+    assert(nv==nv2)
+    assert(T1==T+1)
+    states = states[:T]
+
+    # make sure metric keys and calculations match that in src.evaluation.IntersimpleEvaluation
+    
+    hard_brake = -3.
+    timestep = 0.1
+    keys = ['col_all','x_all','y_all','v_all', 'a_all','j_all', 'v_avg', 'a_avg', 'col', 'brake', 't']
+    metrics = {key:[None]*nv for key in keys}
+    
+    for i in range(nv):
+        
+        nni = ~torch.isnan(states[:,i,0])
+
+        metrics['col_all'][i] = [False] * sum(nni) 
+        metrics['x_all'][i] = states[nni,i,0].numpy()
+        metrics['y_all'][i] = states[nni,i,1].numpy()
+        metrics['v_all'][i] = states[nni,i,2].numpy()
+        metrics['a_all'][i] = actions[nni,i,0].numpy()
+        
+        # jerk
+        metrics['j_all'][i] = np.diff(metrics['a_all'][i]) / timestep
+
+        # average velocity and acceleration
+        metrics['v_avg'][i] = np.mean(metrics['v_all'][i])
+        metrics['a_avg'][i] = np.mean(metrics['a_all'][i])
+
+        # collision?
+        metrics['col'][i] = any(metrics['col_all'][i]) 
+
+        # brake?
+        metrics['brake'][i] = any(metrics['a_all'][i] < hard_brake)
+
+        # time length 
+        metrics['t'][i] = sum(nni) 
+
+    for key in metrics.keys():
+        assert(len(metrics[key])==nv)
+    return metrics
+
+def generate_expert_metrics(locations: List[Tuple[int,int]]) -> List[Dict[str, list]]:
+    """"
+    Given a list of locations, for and return a list of metrics for each location
+
+    Args:
+        locations (list): list of (roundabout, track) ints
+
+    Returns:
+        expert_metrics (list of dicts): expert_metrics[i][j][k] returns the value of metric 'j' 
+            evaluated on the kth episode (car) of the ith expert roundabout trackfile
+    """
+    expert_metrics = []
+    for (roundabout, track) in locations:
+        states, actions = load_expert_states(roundabout, track)
+        expert_metrics.append(form_expert_metrics(states, actions))
+    return expert_metrics
+
+def load_expert_states(roundabout:int, track:int):
+    """
+    Load expert states from roundabout/track info
+    Args:
+        roundabout (int): roundabout index
+        track (int): 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
+    """
+    rname = intersim.LOCATIONS[roundabout]
+    state_path = 'expert_data/%s/track%04i/joint_expert_states.pt'%(rname, track)
+    action_path = 'expert_data/%s/track%04i/joint_expert_actions.pt'%(rname, track)
+    states = torch.load(state_path)
+    actions = torch.load(action_path)
+    return states, actions
+
+def evaluate_policy(locations:List[Tuple[int,int]], 
+    env_class:str, 
+    env_kwargs:dict,
+    method: str,
+    policy_file: str,
+    policy_kwargs:dict) -> List[Dict[str,list]]:
+    """
+    Evaluate policy on an incrementing agent environment at all locations. 
+    Return metrics for that policy
+
+    Args:
+        policy (BaseAlgorithm): policy to evaluate
+        locations (list of tuples): list of locations to evaluate policy
+        env_class (str): name of environment to evaluate policy with
+        env_kwargs (dict): key word arguments to initialize environment with
+        method (str): policy method
+        policy_file (str): policy file path
+        policy_kwargs (dict): policy kwargs
+    
+    Returns:
+        policy_metrics (list of dicts): policy_metrics[i][j][k] returns the value of metric 'j' 
+            evaluated on the kth episode (car) of the ith roundabout trackfile under a policy
+    """
+    envs_dict = dict(intersim.envs.intersimple.__dict__)
+    envs_dict.update(dict(options_envs.__dict__))
+    envs_dict.update(dict(options_envs2.__dict__))
+    envs_dict.update(dict(options_envs3.__dict__))
+    policy_metrics = [None]* len(locations)
+
+    # iterate through vehicles
+    for i, location in tqdm(enumerate(locations)):
+        
+        # add roundabout and track to environent 
+        iround, track = location
+        rname = intersim.LOCATIONS[iround]
+        it_env_kwargs = deepcopy(env_kwargs)
+        loc_kwargs = {
+            'loc':iround,
+            'track':track
+        }
+        it_env_kwargs.update(loc_kwargs)
+
+        # initialize environment
+        Env = envs_dict[env_class]
+        eval_env = Env(**env_kwargs)
+        evaluator = IntersimpleEvaluation(eval_env)
+
+        # load policy
+        policy = load_policy(method, policy_file, policy_kwargs, eval_env)
+
+        # run policy on environment
+        policy_metrics[i] = evaluator.evaluate(policy)
+    
+    return policy_metrics
+        
+def summary_metrics(metrics:List[Dict[str,list]]) -> Dict[str,float]:
+    """
+    Summarize and print metrics averaged over vehicles and roundabouts
+
+    Args:
+        metrics (list of dicts): policy_metrics[i][j][k] returns the value of metric 'j' 
+            evaluated on the kth episode (car) of the ith roundabout trackfile under a policy
+    
+    Returns:
+        summary_metrics (Dict[str,float]): maps summary metric descriptions to values
+    """
+    # keys = ['col_all','v_all', 'a_all','j_all', 'v_avg', 'a_avg', 'col', 'brake', 't']
+    summary_metrics = {}
+
+    # average average-velocity
+    all_vavgs =  sum([d['v_avg'] for d in metrics],[]) # aggregate to single list
+    summary_metrics['mean average velocity'] = sum(all_vavgs)/len(all_vavgs)
+    
+    # average acceleration
+    all_aalls = np.concatenate([np.concatenate(d['a_all']) for d in metrics])
+    summary_metrics['mean acceleration'] = np.mean(all_aalls)
+
+    # average +acceleration
+    pos_accels = all_aalls[all_aalls>0]
+    summary_metrics['mean positive acceleration'] = np.mean(pos_accels)
+
+    # average deceleration
+    decels = all_aalls[all_aalls<0]
+    summary_metrics['mean deceleration'] = np.mean(decels)
+
+    # average jerk
+    all_jerks = np.concatenate([np.concatenate(d['j_all']) for d in metrics])
+    summary_metrics['mean jerk'] = np.mean(all_jerks)
+
+    # average |jerk|
+    summary_metrics['mean |jerk|'] = np.mean(np.abs(all_jerks))
+
+    # collision rate
+    all_collisions = sum([d['col'] for d in metrics],[]) # aggregate to single list
+    summary_metrics['collision rate'] = sum(all_collisions)/len(all_collisions)
+
+    # hard brake rate
+    all_hard_brakes = sum([d['brake'] for d in metrics],[]) # aggregate to single list
+    summary_metrics['hard brake rate'] = sum(all_hard_brakes)/len(all_hard_brakes)
+
+    # average number of timesteps
+    all_ts = sum([d['t'] for d in metrics],[]) # aggregate to single list
+    summary_metrics['mean episode length'] = sum(all_ts)/len(all_ts)
+
+    for key in summary_metrics.keys():
+        print(f'{key}: {summary_metrics[key]}')
+
+    return summary_metrics
+
+def comparison_metrics(policy_metrics:List[Dict[str,list]], 
+    expert_metrics:List[Dict[str,list]], outbase:str='' ) -> Dict[str,float]:
+    """
+    Provide distributional comparison between different sets of metrics
+
+    Args:
+        policy_metrics (list of dicts): policy_metrics[i][j][k] returns the value of metric 'j' 
+            evaluated on the kth episode (car) of the ith roundabout trackfile under a policy
+        expert_metrics (list of dicts): expert_metrics[i][j][k] returns the value of metric 'j' 
+            evaluated on the kth episode (car) of the ith expert roundabout trackfile
+        outbase (str): path to save output figs to
+
+    Returns:   
+        comparison_metrics (Dict[str,float]): dict mapping comparison metric description to value
+    """
+    comparison_metrics = {}
+
+    # rwse
+    expert_traj, policy_traj = [], []
+    for iR in range(len(policy_metrics)):
+        for iTraj in range(len(policy_metrics[iR]['x_all'])):
+            expert_traj.append(np.vstack((expert_metrics[iR]['x_all'][iTraj], expert_metrics[iR]['y_all'][iTraj]))) 
+            policy_traj.append(np.vstack((policy_metrics[iR]['x_all'][iTraj], policy_metrics[iR]['y_all'][iTraj]))) 
+    assert len(expert_traj)==len(policy_traj)
+    comparison_metrics['rwse'] = rwse(expert_traj, policy_traj)
+
+    # average velocity shortfall
+    expert_vavg = np.array(sum([d['v_avg'] for d in expert_metrics],[]))
+    policy_vavg = np.array(sum([d['v_avg'] for d in policy_metrics],[]))
+    assert len(expert_vavg)==len(policy_vavg)
+    comparison_metrics['mean shortfall velocity'] = np.mean(expert_vavg - policy_vavg)
+
+    # Average |Delta V average|
+    comparison_metrics['average absolute average velocity'] = np.mean(np.abs(expert_vavg - policy_vavg))
+    
+    # velocity JSD
+    expert_vs = torch.tensor(np.concatenate([np.concatenate(d['v_all']) for d in expert_metrics]))
+    policy_vs = torch.tensor(np.concatenate([np.concatenate(d['v_all']) for d in policy_metrics]))
+    comparison_metrics['velocity distribution divergence'] = divergence(expert_vs, policy_vs)
+    visualize_distribution(expert_vs, policy_vs, outbase+'_velocity_jsd')
+    
+    # acceleration JSD
+    expert_as = torch.tensor(np.concatenate([np.concatenate(d['a_all']) for d in expert_metrics]))
+    policy_as = torch.tensor(np.concatenate([np.concatenate(d['a_all']) for d in policy_metrics]))
+    comparison_metrics['acceleration distribution divergence'] = divergence(expert_as, policy_as)
+    visualize_distribution(expert_as, policy_as, outbase+'_accel_jsd')
+    
+    # jerk JSD
+    expert_jerks = torch.tensor(np.concatenate([np.concatenate(d['j_all']) for d in expert_metrics]))
+    policy_jerks = torch.tensor(np.concatenate([np.concatenate(d['j_all']) for d in policy_metrics]))
+    comparison_metrics['jerk distribution divergence'] = divergence(expert_jerks, policy_jerks)
+    visualize_distribution(expert_jerks, policy_jerks, outbase+'_jerk_jsd')
+
+    for key in comparison_metrics.keys():
+        print(f'{key}: {comparison_metrics[key]}')
+        
+    return comparison_metrics
+
+def eval_main(
+    locations: List[Tuple[int,int]]= [(0,0)],
+    method: str='expert', 
+    policy_file: str='', 
+    policy_kwargs: dict={},
+    env: str='NRasterizedRouteIncrementingAgent', 
+    env_kwargs: dict={},
+    seed: int=0):
+    """
+    Test a particular model at different testing locations/tracks and compute average metrics 
+    over all files. 
+
+    Args:
+        locations (list of tuples): list of (roundabout, track) integer pair testing locations
+        method (str): method string
+        policy_file (str): path to saved policy
+        env (str): environment class
+        method (str): method (expert, bc, gail, rail, hgail, hrail)
+    """
+    print(f'Evaluating {method} on {env}')
+
+    # set seed
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    pfilename = policy_file.split('/')[-1].split('.')[0]
+    outbase = f'out/{method}/{pfilename}_seed{seed}'
+
+    # load expert metrics
+    expert_metrics = generate_expert_metrics(locations)
+
+    # no comparison for expert
+    if method=='expert':
+        smetrics = summary_metrics(expert_metrics)
+        save_metrics(smetrics, outbase+'_summary.pkl')
+    
+    # otherwise evaluate policy on roundabouts and generate metrics
+    else:
+
+        # evaluate it on the given roundabouts
+        policy_metrics = evaluate_policy(locations, env, env_kwargs, method, policy_file, policy_kwargs)
+        smetrics = summary_metrics(policy_metrics)
+        save_metrics(smetrics, outbase+'_summary.pkl')
+        cmetrics = comparison_metrics(policy_metrics, expert_metrics, outbase=outbase)
+        save_metrics(cmetrics, outbase+'_comparison.pkl')
+
+if __name__=='__main__':
+    import fire
+    fire.Fire(eval_main)

src/evaluation/__init__.py

@@ -0,0 +1,2 @@
+from src.evaluation.evaluation import IntersimpleEvaluation
+from src.evaluation.metrics import divergence, visualize_distribution