Use ray tune

.gitignore

@@ -1,8 +1,6 @@
 *.pkl
 *.pt
 *.zip
-**/ray/*
-**/runs/*
 
 # Byte-compiled / optimized / DLL files
 __pycache__/
@@ -151,4 +149,3 @@ output/
 
 # Dependencies
 InteractionSimulator/
-imitation/

evaluate_models.sh

@@ -1,42 +0,0 @@
-# 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}'

generate_demos.sh

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

requirements.txt

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

scratch/arec/intersimple/commands.txt

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

scratch/arec/intersimple/data/expert.py

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

scratch/arec/intersimple/data/generate.sh

@@ -1,9 +0,0 @@
-#python -m expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51.pkl'
-#python -m expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --policy_args='{mu:0.005}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.005.pkl'
-#python -m expert --env=IntersimpleReward --min_timesteps=200 --env_args='{agent:51}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.001.pkl'
-#python -m expert --env=NRasterized --min_timesteps=200 --env_args='{agent:51,width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl'
-# python -m expert --env=NRasterizedRandomAgent --min_timesteps=10000 --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001N10000_NRasterizedRandomAgentw36h36mppx2.pkl'
-#python -m expert --env=NRasterized --min_timesteps=200 --env_args='{agent:51,width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl'
-#python -m expert --env=NRasterized --min_timesteps=3000 --video --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedRandomAgentw36h36mppx2.pkl'
-#python -m expert --env=NRasterizedIncrementingAgent --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}' --path='NormalizedIntersimpleExpertMu.001_NRasterizedIncrementingAgentw36h36mppx2.pkl'
-python -m process_all_experts --env_args='{width:36,height:36,m_per_px:2}' --policy_args='{mu:0.001}'

scratch/arec/intersimple/data/load_experts.py

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

scratch/arec/intersimple/data/process_all_experts.py

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

scratch/arec/intersimple/gail/discriminator.py

@@ -1,101 +0,0 @@
-import torch
-
-# imitation.rewards.discrim_nets.DiscrimNetGAIL is composed of self.discriminator (nn.Module),
-# which gets called with inputs (state, action) when needed.
-
-class CnnDiscriminator(torch.nn.Module):
-    """ConvNet similar to stable_baselines3.common.policies.ActorCriticCnnPolicy."""
-
-    def __init__(self, env):
-        super().__init__()
-
-        obs_channels, _, _ = env.observation_space.shape
-        (action_size,) = env.action_space.shape
-        in_channels = obs_channels + action_size
-
-        self.cnn = torch.nn.Sequential(
-            torch.nn.Conv2d(in_channels, 32, kernel_size=(8, 8), stride=(4, 4)), # 5+1 -> 32
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2)), # 32 -> 64
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1)), # 64 -> 64
-            torch.nn.ReLU(),
-            torch.nn.Flatten(start_dim=1, end_dim=-1),
-            torch.nn.LazyLinear(512), # 28224 -> 512
-            torch.nn.ReLU(),
-            torch.nn.LazyLinear(1), # 512 -> 1
-        )
-    
-    @staticmethod
-    def _concatenate(state, action):
-        b, _, h, w = state.shape
-        _, a = action.shape
-        act = action.unsqueeze(-1).unsqueeze(-1).expand((b, a, h, w))
-        sa = torch.cat((state, act), -3)
-        return sa
-
-    def forward(self, state, action):
-        sa = self._concatenate(state, action)
-        assert sa.ndim == 4
-        return self.cnn(sa).squeeze(1)
-
-class CnnDiscriminatorFlatAction(torch.nn.Module):
-    """ConvNet similar to stable_baselines3.common.policies.ActorCriticCnnPolicy."""
-
-    def __init__(self, env):
-        super().__init__()
-
-        obs_channels, _, _ = env.observation_space.shape
-        (action_size,) = env.action_space.shape
-        in_channels = obs_channels
-
-        self.cnn = torch.nn.Sequential(
-            torch.nn.Conv2d(in_channels, 32, kernel_size=(8, 8), stride=(4, 4)), # in_channels -> 32
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2)), # 32 -> 64
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1)), # 64 -> 64
-            torch.nn.ReLU(),
-            torch.nn.Flatten(start_dim=1, end_dim=-1),
-            torch.nn.LazyLinear(128), # 28224 -> 128
-        )
-        self.decoder = torch.nn.Sequential(
-            torch.nn.LazyLinear(64), #128 + 2 -> 64
-            torch.nn.ReLU(),
-            torch.nn.LazyLinear(64), #64 -> 64
-            torch.nn.ReLU(),
-            torch.nn.LazyLinear(1) #64 -> 1
-        )
-    
-    @staticmethod
-    def _concatenate(state, action):
-        b, s= state.shape
-        b, a = action.shape
-        sa = torch.cat((state, action), -1)
-        return sa
-
-    def forward(self, state, action):
-        s = self.cnn(state.float())
-        sa = self._concatenate(s, action)
-        assert sa.ndim == 2
-        return self.decoder(sa).squeeze(1)
-
-class MlpDiscriminator(torch.nn.Module):
-    """MLP similar to stable_baselines3.common.policies.ActorCriticPolicy."""
-
-    def __init__(self, env=None):
-        super().__init__()
-        self.flatten = torch.nn.Flatten(start_dim=1, end_dim=-1)
-        self.mlp = torch.nn.Sequential(
-            torch.nn.LazyLinear(64), # 42 -> 64
-            torch.nn.Tanh(),
-            torch.nn.LazyLinear(64), # 64 -> 64
-            torch.nn.Tanh(),
-            torch.nn.LazyLinear(1),  # 64 -> 1
-        )
-    
-    def forward(self, state, action):
-        flat = self.flatten(state)
-        sa = torch.cat((action, flat), -1)
-        assert sa.ndim == 2
-        return self.mlp(sa).squeeze(1)

scratch/arec/intersimple/gail_image_multiagent_nocollision.py

@@ -1,70 +0,0 @@
-# %%
-import pathlib
-import pickle
-import tempfile
-
-import stable_baselines3 as sb3
-from stable_baselines3.common.env_util import make_vec_env
-
-from imitation.algorithms import adversarial, bc
-from imitation.data import rollout
-from imitation.util import logger
-
-from intersim.envs.intersimple import NRasterized
-
-from gail.discriminator import CnnDiscriminatorFlatAction
-
-model_name = 'gail_image_multiagent_nocollision'
-
-# %%
-# Load pickled test demonstrations.
-with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedRandomAgentw36h36mppx2.pkl", "rb") as f:
-    # This is a list of `imitation.data.types.Trajectory`, where
-    # every instance contains observations and actions for a single expert
-    # demonstration.
-    trajectories = pickle.load(f)
-
-# %%
-# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
-# This is a more general dataclass containing unordered
-# (observation, actions, next_observation) transitions.
-transitions = rollout.flatten_trajectories(trajectories)
-
-venv = make_vec_env(NRasterized, n_envs=2, env_kwargs={'stop_on_collision':False, 'width': 36, 'height': 36, 'm_per_px': 2})
-
-tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
-tempdir_path = pathlib.Path(tempdir.name)
-print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
-
-# Train GAIL on expert data.
-# GAIL, and AIRL also accept as `expert_data` any Pytorch-style DataLoader that
-# iterates over dictionaries containing observations, actions, and next_observations.
-logger.configure(tempdir_path / "GAIL/")
-gail_trainer = adversarial.GAIL(
-    venv,
-    expert_data=transitions,
-    expert_batch_size=32,
-    #n_disc_updates_per_round=2048,
-    discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
-    gen_algo=sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=1024),
-    allow_variable_horizon=True,
-)
-gail_trainer.train(total_timesteps=100000)
-gail_trainer.gen_algo.save(model_name)
-
-#del gail_trainer
-
-# %%
-model = sb3.PPO.load(model_name)
-
-env = NRasterized(stop_on_collision=False, width=36, height=36, m_per_px=2)
-
-obs = env.reset()
-while True:
-    action, _states = model.predict(obs)
-    obs, rewards, done, info = env.step(action)
-    env.render(mode='post')
-    if done:
-        break
-
-env.close(filestr='render/'+model_name)

scratch/arec/intersimple/gail_image_singleagent_nocollision.py

@@ -1,70 +0,0 @@
-# %%
-import pathlib
-import pickle
-import tempfile
-
-import stable_baselines3 as sb3
-from stable_baselines3.common.env_util import make_vec_env
-
-from imitation.algorithms import adversarial, bc
-from imitation.data import rollout
-from imitation.util import logger
-
-from intersim.envs.intersimple import NRasterized
-
-from gail.discriminator import CnnDiscriminator
-
-model_name = 'gail_image_singleagent_nocollision'
-
-# %%
-# Load pickled test demonstrations.
-with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
-    # This is a list of `imitation.data.types.Trajectory`, where
-    # every instance contains observations and actions for a single expert
-    # demonstration.
-    trajectories = pickle.load(f)
-
-# %%
-# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
-# This is a more general dataclass containing unordered
-# (observation, actions, next_observation) transitions.
-transitions = rollout.flatten_trajectories(trajectories)
-
-venv = make_vec_env(NRasterized, n_envs=2, env_kwargs={'agent':51, 'stop_on_collision':False, 'width': 36, 'height': 36, 'm_per_px': 2})
-
-tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
-tempdir_path = pathlib.Path(tempdir.name)
-print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
-
-# Train GAIL on expert data.
-# GAIL, and AIRL also accept as `expert_data` any Pytorch-style DataLoader that
-# iterates over dictionaries containing observations, actions, and next_observations.
-logger.configure(tempdir_path / "GAIL/")
-gail_trainer = adversarial.GAIL(
-    venv,
-    expert_data=transitions,
-    expert_batch_size=32,
-    #n_disc_updates_per_round=2048,
-    discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
-    gen_algo=sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=1024),
-    allow_variable_horizon=True,
-)
-gail_trainer.train(total_timesteps=100000)
-gail_trainer.gen_algo.save(model_name)
-
-#del gail_trainer
-
-# %%
-model = sb3.PPO.load(model_name)
-
-env = NRasterized(agent=51, width=36, height=36, m_per_px=2, stop_on_collision=False)
-
-obs = env.reset()
-while True:
-    action, _states = model.predict(obs)
-    obs, rewards, done, info = env.step(action)
-    env.render(mode='post')
-    if done:
-        break
-
-env.close(filestr='render/'+model_name)

scratch/arec/intersimple/gail_options_image.py

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

scratch/arec/intersimple/gail_options_scratch.py

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

scratch/arec/intersimple/options_gail.py

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

scratch/arec/intersimple/plan.txt

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

scratch/arec/intersimple/render_env_from_model.py

@@ -1,58 +0,0 @@
-
-import stable_baselines3 as sb3
-from intersim.envs.intersimple import NRasterized
-
-
-def render_env(model_name='gail_image_multiagent_nocollision', agent=51, environment=NRasterized):
-    """
-    Render a video from an model, agent, and environment
-    Args:
-        model_name (str): name of the model
-        agent (int): agent to start the video from
-        environment (gym.Env): gym environment class to render environment on
-    """
-
-    model = sb3.PPO.load(model_name)
-
-    env = environment(stop_on_collision=False, width=36, height=36, m_per_px=2, agent=agent)
-
-    obs = env.reset()
-    i=0
-    while True and i < 600:
-        i+=1
-        action, _states = model.predict(obs)
-        obs, rewards, done, info = env.step(action)
-        env.render(mode='post')
-        if done:
-            break
-
-    env.close(filestr='render/'+model_name+'_agent%i'%(agent)) 
-
-def render_options_env(model_name='gail_image_multiagent_nocollision', agent=51, environment=NRasterized):
-    """
-    Render a video from an model, agent, and environment
-    Args:
-        model_name (str): name of the model
-        agent (int): agent to start the video from
-        environment (gym.Env): gym environment class to render environment on
-    """
-
-    model = sb3.PPO.load(model_name)
-
-    env = environment(stop_on_collision=False, width=36, height=36, m_per_px=2, agent=agent)
-
-    obs = env.reset()
-    i=0
-    while True and i < 600:
-        i+=1
-        action, _states = model.predict(obs)
-        obs, rewards, done, info = env.step(action)
-        env.render(mode='post')
-        if done:
-            break
-
-    env.close(filestr='render/'+model_name+'_agent%i'%(agent)) 
-
-if __name__ == '__main__':
-    import fire
-    fire.Fire(render_env)

scratch/arec/intersimple/render_options.py

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

scratch/etienne/intersimple/data/expert.py

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

scratch/etienne/intersimple/data/generate.sh

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

scratch/etienne/intersimple/gail/discriminator.py

@@ -36,49 +36,7 @@ class CnnDiscriminator(torch.nn.Module):
 
     def forward(self, state, action):
         sa = self._concatenate(state, action)
-        assert sa.ndim == 4
+        return self.cnn(sa).squeeze()
-        return self.cnn(sa).squeeze(1)
-
-class CnnDiscriminatorFlatAction(torch.nn.Module):
-    """ConvNet similar to stable_baselines3.common.policies.ActorCriticCnnPolicy."""
-
-    def __init__(self, env):
-        super().__init__()
-
-        obs_channels, _, _ = env.observation_space.shape
-        (action_size,) = env.action_space.shape
-        in_channels = obs_channels
-
-        self.cnn = torch.nn.Sequential(
-            torch.nn.Conv2d(in_channels, 32, kernel_size=(8, 8), stride=(4, 4)), # in_channels -> 32
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2)), # 32 -> 64
-            torch.nn.ReLU(),
-            torch.nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1)), # 64 -> 64
-            torch.nn.ReLU(),
-            torch.nn.Flatten(start_dim=1, end_dim=-1),
-            torch.nn.LazyLinear(128), # 28224 -> 128
-        )
-        self.decoder = torch.nn.Sequential(
-            torch.nn.LazyLinear(64), #128 + 2 -> 64
-            torch.nn.ReLU(),
-            torch.nn.LazyLinear(64), #64 -> 64
-            torch.nn.ReLU(),
-            torch.nn.LazyLinear(1) #64 -> 1
-        )
-    
-    @staticmethod
-    def _concatenate(state, action):
-        b, s= state.shape
-        b, a = action.shape
-        sa = torch.cat((state, action), -1)
-        return sa
-
-    def forward(self, state, action):
-        s = self.cnn(state.float())
-        sa = self._concatenate(s, action)
-        assert sa.ndim == 2
-        return self.decoder(sa).squeeze(1)
 
 class MlpDiscriminator(torch.nn.Module):
     """MLP similar to stable_baselines3.common.policies.ActorCriticPolicy."""
@@ -97,5 +55,4 @@ class MlpDiscriminator(torch.nn.Module):
     def forward(self, state, action):
         flat = self.flatten(state)
         sa = torch.cat((action, flat), -1)
-        assert sa.ndim == 2
+        return self.mlp(sa).squeeze()
-        return self.mlp(sa).squeeze(1)

scratch/etienne/intersimple/gail/envs.py

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

scratch/etienne/intersimple/gail/options2.py

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

scratch/etienne/intersimple/gail_flat.py

@@ -10,15 +10,22 @@ from imitation.algorithms import adversarial, bc
 from imitation.data import rollout
 from imitation.util import logger
 
-from intersim.envs.intersimple import IntersimpleReward
+from intersimple.intersimple import IntersimpleReward, speed_reward
 
 from gail.discriminator import MlpDiscriminator
+import numpy as np
+import functools
+from stable_baselines3.common.evaluation import evaluate_policy
+from ray import tune
+import os
+import torch
 
 model_name = 'gail_flat'
 
 # %%
 # Load pickled test demonstrations.
-with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51.pkl", "rb") as f:
+#with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51.pkl", "rb") as f:
+with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.001.pkl", "rb") as f:
     # This is a list of `imitation.data.types.Trajectory`, where
     # every instance contains observations and actions for a single expert
     # demonstration.
@@ -36,21 +43,50 @@ tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
 tempdir_path = pathlib.Path(tempdir.name)
 print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
 
-# Train GAIL on expert data.
+def training_function(config, checkpoint_dir=None):
-# GAIL, and AIRL also accept as `expert_data` any Pytorch-style DataLoader that
+    logger.configure(tempdir_path / "GAIL/")
-# iterates over dictionaries containing observations, actions, and next_observations.
+
-logger.configure(tempdir_path / "GAIL/")
+    discriminator = MlpDiscriminator()
-gail_trainer = adversarial.GAIL(
+    if checkpoint_dir:
+        discriminator.load_state_dict(torch.load(os.path.join(checkpoint_dir, 'disc_checkpoint')))
+        generator = sb3.PPO.load(os.path.join(checkpoint_dir, 'gen_checkpoint'))
+    else:
+        generator = sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=config['n_steps'])
+
+    gail_trainer = adversarial.GAIL(
         venv,
         expert_data=transitions,
-    expert_batch_size=150,
+        expert_batch_size=config['expert_batch_size'],
-    n_disc_updates_per_round=32,
+        n_disc_updates_per_round=config['n_disc_updates_per_round'],
         discrim_kwargs={'discrim_net': MlpDiscriminator()},
-    gen_algo=sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=4530),
+        gen_algo=generator,
-    allow_variable_horizon=True,
+    )
+
+    def callback(epoch):
+        eval_env = IntersimpleReward(agent=51, reward=functools.partial(speed_reward, collision_penalty=0.))
+        #sync_envs_normalization(self.training_env, self.eval_env)
+        episode_rewards, episode_lengths = evaluate_policy(generator, eval_env)
+        tune.report(progress=np.mean(episode_rewards))
+
+        with tune.checkpoint_dir(step=epoch) as checkpoint_dir:
+            gail_trainer.gen_algo.save(os.path.join(checkpoint_dir, 'gen_checkpoint'))
+            torch.save(discriminator.state_dict(), os.path.join(checkpoint_dir, 'disc_checkpoint'))
+
+    gail_trainer.train(total_timesteps=400000, callback=callback)
+
+analysis = tune.run(
+    training_function,
+    config = {
+        'expert_batch_size': tune.randint(1, 220), #220,
+        'n_disc_updates_per_round': tune.randint(2, 100), #16,
+        'n_steps': tune.randint(1, 10000), #4096,
+    },
+    resources_per_trial={
+        'gpu': 1,
+    }
 )
-gail_trainer.train(total_timesteps=400000)
+
-gail_trainer.gen_algo.save(model_name)
+print('Best config', analysis.get_best_config(metric='progress', mode='max'))
 
 #del gail_trainer
 

scratch/etienne/intersimple/gail_flat_ray.py

@@ -1,115 +0,0 @@
-# %%
-import pathlib
-import pickle
-import tempfile
-
-import stable_baselines3 as sb3
-from stable_baselines3.common.env_util import make_vec_env
-
-from imitation.algorithms import adversarial, bc
-from imitation.data import rollout
-from imitation.util import logger
-
-from intersim.envs.intersimple import IntersimpleReward, speed_reward
-
-from gail.discriminator import MlpDiscriminator
-import numpy as np
-import functools
-from stable_baselines3.common.evaluation import evaluate_policy
-from ray import tune
-import os
-import torch
-
-model_name = 'gail_flat'
-
-# %%
-# Load pickled test demonstrations.
-#with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51.pkl", "rb") as f:
-with open("data/NormalizedIntersimpleExpert_IntersimpleRewardAgent51Mu.001.pkl", "rb") as f:
-    # This is a list of `imitation.data.types.Trajectory`, where
-    # every instance contains observations and actions for a single expert
-    # demonstration.
-    trajectories = pickle.load(f)
-
-# %%
-# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
-# This is a more general dataclass containing unordered
-# (observation, actions, next_observation) transitions.
-transitions = rollout.flatten_trajectories(trajectories)
-
-venv = make_vec_env(IntersimpleReward, n_envs=2, env_kwargs={'agent': 51})
-
-tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
-tempdir_path = pathlib.Path(tempdir.name)
-print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
-
-def training_function(config, checkpoint_dir=None):
-    logger.configure(tempdir_path / "GAIL/")
-
-    discriminator = MlpDiscriminator()
-    if checkpoint_dir:
-        discriminator.load_state_dict(torch.load(os.path.join(checkpoint_dir, 'disc_checkpoint')))
-        generator = sb3.PPO.load(os.path.join(checkpoint_dir, 'gen_checkpoint'))
-    else:
-        generator = sb3.PPO("MlpPolicy", venv, verbose=1, n_steps=config['n_steps'])
-
-    gail_trainer = adversarial.GAIL(
-        venv,
-        expert_data=transitions,
-        expert_batch_size=config['expert_batch_size'],
-        n_disc_updates_per_round=config['n_disc_updates_per_round'],
-        discrim_kwargs={'discrim_net': MlpDiscriminator()},
-        gen_algo=generator,
-        allow_variable_horizon=True,
-    )
-
-    def callback(epoch):
-        print("callback")
-        eval_env = IntersimpleReward(agent=51, reward=functools.partial(speed_reward, collision_penalty=0.))
-        #sync_envs_normalization(self.training_env, self.eval_env)
-        episode_rewards, episode_lengths = evaluate_policy(generator, eval_env, return_episode_rewards=True)
-        tune.report(
-            reward=np.mean(episode_rewards),
-            length=np.mean(episode_lengths),
-            training_iteration=epoch,
-        )
-
-        with tune.checkpoint_dir(step=epoch) as checkpoint_dir:
-            gail_trainer.gen_algo.save(os.path.join(checkpoint_dir, 'gen_checkpoint'))
-            torch.save(discriminator.state_dict(), os.path.join(checkpoint_dir, 'disc_checkpoint'))
-
-    gail_trainer.train(total_timesteps=40000, callback=callback)
-
-analysis = tune.run(
-    training_function,
-    config = {
-        'expert_batch_size': tune.randint(1, 22), #220,
-        'n_disc_updates_per_round': tune.randint(2, 100), #16,
-        'n_steps': tune.randint(1, 10000), #4096,
-    },
-    resources_per_trial={
-        'cpu': 1,
-        # 'gpu': 1,
-    },
-    local_dir='ray',
-    num_samples=10,
-)
-
-print('Best config', analysis.get_best_config(metric='progress', mode='max'))
-
-#del gail_trainer
-
-# %%
-model = sb3.PPO.load(model_name)
-
-env = IntersimpleReward(agent=51)
-
-obs = env.reset()
-while True:
-    action, _states = model.predict(obs)
-    obs, rewards, done, info = env.step(action)
-    env.render(mode='post')
-    if done:
-        break
-
-env.close(filestr='render/'+model_name)

scratch/etienne/intersimple/gail_image_multiagent_nocollision.py

@@ -1,70 +0,0 @@
-# %%
-import pathlib
-import pickle
-import tempfile
-
-import stable_baselines3 as sb3
-from stable_baselines3.common.env_util import make_vec_env
-
-from imitation.algorithms import adversarial, bc
-from imitation.data import rollout
-from imitation.util import logger
-
-from intersim.envs.intersimple import NRasterized
-
-from gail.discriminator import CnnDiscriminatorFlatAction
-
-model_name = 'gail_image_multiagent_nocollision'
-
-# %%
-# Load pickled test demonstrations.
-with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedRandomAgentw36h36mppx2.pkl", "rb") as f:
-    # This is a list of `imitation.data.types.Trajectory`, where
-    # every instance contains observations and actions for a single expert
-    # demonstration.
-    trajectories = pickle.load(f)
-
-# %%
-# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
-# This is a more general dataclass containing unordered
-# (observation, actions, next_observation) transitions.
-transitions = rollout.flatten_trajectories(trajectories)
-
-venv = make_vec_env(NRasterized, n_envs=2, env_kwargs={'stop_on_collision':False, 'width': 36, 'height': 36, 'm_per_px': 2})
-
-tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
-tempdir_path = pathlib.Path(tempdir.name)
-print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
-
-# Train GAIL on expert data.
-# GAIL, and AIRL also accept as `expert_data` any Pytorch-style DataLoader that
-# iterates over dictionaries containing observations, actions, and next_observations.
-logger.configure(tempdir_path / "GAIL/")
-gail_trainer = adversarial.GAIL(
-    venv,
-    expert_data=transitions,
-    expert_batch_size=32,
-    #n_disc_updates_per_round=2048,
-    discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
-    gen_algo=sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=1024),
-    allow_variable_horizon=True,
-)
-gail_trainer.train(total_timesteps=100000)
-gail_trainer.gen_algo.save(model_name)
-
-#del gail_trainer
-
-# %%
-model = sb3.PPO.load(model_name)
-
-env = NRasterized(stop_on_collision=False, width=36, height=36, m_per_px=2)
-
-obs = env.reset()
-while True:
-    action, _states = model.predict(obs)
-    obs, rewards, done, info = env.step(action)
-    env.render(mode='post')
-    if done:
-        break
-
-env.close(filestr='render/'+model_name)

scratch/etienne/intersimple/gail_image_random.py

@@ -1,79 +0,0 @@
-# %%
-import pathlib
-import pickle
-import tempfile
-
-import stable_baselines3 as sb3
-from stable_baselines3.common.env_util import make_vec_env
-
-from imitation.algorithms import adversarial, bc
-from imitation.data import rollout
-from imitation.util import logger
-
-from intersim.envs.intersimple import NRasterizedRandomAgent, IntersimpleReward, speed_reward
-import functools
-from stable_baselines3.common.evaluation import evaluate_policy
-
-from gail.discriminator import CnnDiscriminator
-
-model_name = 'gail_image_random'
-
-# %%
-# Load pickled test demonstrations.
-with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedRandomAgentw36h36mppx2.pkl", "rb") as f:
-    # This is a list of `imitation.data.types.Trajectory`, where
-    # every instance contains observations and actions for a single expert
-    # demonstration.
-    trajectories = pickle.load(f)
-
-# %%
-# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
-# This is a more general dataclass containing unordered
-# (observation, actions, next_observation) transitions.
-transitions = rollout.flatten_trajectories(trajectories)
-
-env_kwargs = {'width': 36, 'height': 36, 'm_per_px': 2}
-venv = make_vec_env(NRasterizedRandomAgent, n_envs=2, env_kwargs=env_kwargs)
-
-tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
-tempdir_path = pathlib.Path(tempdir.name)
-print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
-
-# Train GAIL on expert data.
-# GAIL, and AIRL also accept as `expert_data` any Pytorch-style DataLoader that
-# iterates over dictionaries containing observations, actions, and next_observations.
-logger.configure(tempdir_path / "GAIL/")
-generator = sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=1024)
-gail_trainer = adversarial.GAIL(
-    venv,
-    expert_data=transitions,
-    expert_batch_size=32,
-    #n_disc_updates_per_round=2048,
-    discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
-    gen_algo=generator,
-    allow_variable_horizon=True,
-)
-def callback(round):
-    eval_env = NRasterizedRandomAgent(reward=functools.partial(speed_reward, collision_penalty=0.), **env_kwargs)
-    #sync_envs_normalization(self.training_env, self.eval_env)
-    episode_rewards, episode_lengths = evaluate_policy(generator, eval_env, return_episode_rewards=True)
-
-gail_trainer.train(total_timesteps=100000, callback=callback)
-gail_trainer.gen_algo.save(model_name)
-
-#del gail_trainer
-
-# %%
-model = sb3.PPO.load(model_name)
-
-env = NRasterizedRandomAgent(width=36, height=36, m_per_px=2)
-
-obs = env.reset()
-while True:
-    action, _states = model.predict(obs)
-    obs, rewards, done, info = env.step(action)
-    env.render(mode='post')
-    if done:
-        break
-
-env.close(filestr='render/'+model_name)

scratch/etienne/intersimple/gail_image_random_ray.py

@@ -1,171 +0,0 @@
-# %%
-import pathlib
-import pickle
-import tempfile
-import os
-import random
-import numpy as np
-import torch
-
-# set up ray tune
-import ray
-from ray import tune
-from ray.tune import Analysis, ExperimentAnalysis
-from ray.tune.schedulers import ASHAScheduler
-from ray.tune.suggest.hyperopt import HyperOptSearch
-from ray.tune.suggest import ConcurrencyLimiter
-
-import stable_baselines3 as sb3
-from stable_baselines3.common.env_util import make_vec_env
-
-
-from imitation.algorithms import adversarial, bc
-from imitation.data import rollout
-from imitation.util import logger
-
-from intersim.envs.intersimple import NRasterizedRandomAgent, IntersimpleReward, speed_reward, NRasterized, NRasterizedRandomAgentVerbose
-import functools
-from stable_baselines3.common.evaluation import evaluate_policy
-from gym.wrappers import TimeLimit
-
-from gail.discriminator import CnnDiscriminator
-
-model_name = 'gail_image_random_ray'
-env_kwargs={'width': 36, 'height': 36, 'm_per_px': 2}
-
-# %%
-
-import argparse
-parser = argparse.ArgumentParser()
-parser.add_argument("--outdir", help="result directory", default='ray')
-parser.add_argument("--test", help="test run", default=False, action="store_true")
-args = parser.parse_args()
-outdir = args.outdir
-
-# %%
-# Load pickled test demonstrations.
-with open("data/NormalizedIntersimpleExpertMu.001N10000_NRasterizedRandomAgentw36h36mppx2.pkl", "rb") as f:
-    # This is a list of `imitation.data.types.Trajectory`, where
-    # every instance contains observations and actions for a single expert
-    # demonstration.
-    trajectories = pickle.load(f)
-
-# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
-# This is a more general dataclass containing unordered
-# (observation, actions, next_observation) transitions.
-transitions = rollout.flatten_trajectories(trajectories)
-# Store transitions in shared ray memory
-ray_transitions = ray.put(transitions)
-
-# %%
-venv = make_vec_env(NRasterizedRandomAgent, n_envs=2, env_kwargs=env_kwargs)
-
-tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
-tempdir_path = pathlib.Path(tempdir.name)
-print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
-logger.configure(tempdir_path / "GAIL/")
-
-def get_ray_config(test=False):
-    if test:
-        return {
-            'expert_batch_size': 2,
-            'ppo_n_steps': 2,
-            'ppo_batch_size': 2,
-            'ppo_n_epochs': 1,
-            'total_timesteps': 10,
-        }
-    else:
-        return {
-            'expert_batch_size': tune.choice([2**x for x in range(6,10)]),
-            'ppo_n_steps': tune.choice([2048, 3072, 4096]),
-            'ppo_batch_size': tune.choice([2**x for x in range(9,13)]),
-            'ppo_n_epochs': tune.choice([6,10]),
-            'total_timesteps': 400_000,
-        }
-
-
-def ray_train(config, checkpoint_dir=None):
-    # Train GAIL on expert data.
-    # GAIL, and AIRL also accept as `expert_data` any Pytorch-style DataLoader that
-    # iterates over dictionaries containing observations, actions, and next_observations.
-
-    discriminator = CnnDiscriminator(venv)
-    if checkpoint_dir:
-        discriminator.load_state_dict(torch.load(os.path.join(checkpoint_dir, 'disc_checkpoint')))
-        generator = sb3.PPO.load(os.path.join(checkpoint_dir, 'gen_checkpoint'))
-    else:
-        generator = sb3.PPO(
-            "CnnPolicy", venv, verbose=0, 
-            n_steps=config["ppo_n_steps"], 
-            batch_size=config["ppo_batch_size"], 
-            n_epochs=config["ppo_n_epochs"]
-        )
-    gail_trainer = adversarial.GAIL(
-        venv,
-        expert_data=ray.get(ray_transitions),
-        expert_batch_size=config["expert_batch_size"],
-        #n_disc_updates_per_round=2048,
-        discrim_kwargs={'discrim_net': discriminator},
-        gen_algo=generator,
-        allow_variable_horizon=True,
-    )
-    def callback(round):
-        # eval_env = NRasterized(agent=51, reward=functools.partial(speed_reward, collision_penalty=0.), **env_kwargs)
-        eval_env = TimeLimit(NRasterizedRandomAgent(reward=functools.partial(speed_reward, collision_penalty=0.), **env_kwargs), max_episode_steps=1000)
-        episode_rewards, episode_lengths = evaluate_policy(generator, eval_env, return_episode_rewards=True)
-        tune.report(
-            reward=np.mean(episode_rewards),
-            length=np.mean(episode_lengths),
-            training_iteration=round,
-        )
-        with tune.checkpoint_dir(step=round) as checkpoint_dir:
-            gail_trainer.gen_algo.save(os.path.join(checkpoint_dir, 'gen_checkpoint'))
-            torch.save(discriminator.state_dict(), os.path.join(checkpoint_dir, 'disc_checkpoint'))
-
-    gail_trainer.train(total_timesteps=config['total_timesteps'], callback=callback)
-
-
-ray_config = get_ray_config(args.test)
-search = HyperOptSearch(ray_config, metric='length', mode="max",)
-search = ConcurrencyLimiter(search, max_concurrent=10)
-custom_scheduler = ASHAScheduler(time_attr='training_iteration', metric='length', mode="max", grace_period=15)
-
-analysis = tune.run(
-    ray_train, 
-    # config=ray_config,
-    search_alg=search,
-    scheduler=custom_scheduler,
-    local_dir=outdir,
-    resources_per_trial={"cpu":10, "gpu": 0.2},
-    num_samples=1 if args.test else 100,
-)
-
-del analysis
-
-# %%
-# outdir = "ray/ray_train_2021-09-20_13-33-50/ray_train_f06785b0_33_expert_batch_size=128,ppo_batch_size=1024,ppo_n_epochs=6,ppo_n_steps=2048,total_timesteps=400000_2021-09-20_15-52-05"
-
-# %%
-analysis = Analysis(outdir, default_metric="length", default_mode="max")
-filepath = analysis.get_best_logdir()
-print("Best ray experiment:", filepath)
-config = analysis.get_best_config()
-print("Best config:", config)
-
-# %%
-
-model = sb3.PPO.load(os.path.join(analysis.get_last_checkpoint(), 'gen_checkpoint'))
-
-# env = NRasterized(agent=51, **env_kwargs)
-env = TimeLimit(NRasterizedRandomAgent(reward=functools.partial(speed_reward, collision_penalty=0.), **env_kwargs), max_episode_steps=1000)
-obs = env.reset()
-while True:
-    action, _states = model.predict(obs)
-    obs, rewards, done, info = env.step(action)
-    env.render(mode='post')
-    if done:
-        break
-
-env.env.close(filestr='render/'+model_name)
-# %%
-

scratch/etienne/intersimple/gail_image_singleagent_nocollision.py

@@ -1,70 +0,0 @@
-# %%
-import pathlib
-import pickle
-import tempfile
-
-import stable_baselines3 as sb3
-from stable_baselines3.common.env_util import make_vec_env
-
-from imitation.algorithms import adversarial, bc
-from imitation.data import rollout
-from imitation.util import logger
-
-from intersim.envs.intersimple import NRasterized
-
-from gail.discriminator import CnnDiscriminator
-
-model_name = 'gail_image_singleagent_nocollision'
-
-# %%
-# Load pickled test demonstrations.
-with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
-    # This is a list of `imitation.data.types.Trajectory`, where
-    # every instance contains observations and actions for a single expert
-    # demonstration.
-    trajectories = pickle.load(f)
-
-# %%
-# Convert List[types.Trajectory] to an instance of `imitation.data.types.Transitions`.
-# This is a more general dataclass containing unordered
-# (observation, actions, next_observation) transitions.
-transitions = rollout.flatten_trajectories(trajectories)
-
-venv = make_vec_env(NRasterized, n_envs=2, env_kwargs={'agent':51, 'stop_on_collision':False, 'width': 36, 'height': 36, 'm_per_px': 2})
-
-tempdir = tempfile.TemporaryDirectory(prefix="quickstart")
-tempdir_path = pathlib.Path(tempdir.name)
-print(f"All Tensorboards and logging are being written inside {tempdir_path}/.")
-
-# Train GAIL on expert data.
-# GAIL, and AIRL also accept as `expert_data` any Pytorch-style DataLoader that
-# iterates over dictionaries containing observations, actions, and next_observations.
-logger.configure(tempdir_path / "GAIL/")
-gail_trainer = adversarial.GAIL(
-    venv,
-    expert_data=transitions,
-    expert_batch_size=32,
-    #n_disc_updates_per_round=2048,
-    discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
-    gen_algo=sb3.PPO("CnnPolicy", venv, verbose=1, n_steps=1024),
-    allow_variable_horizon=True,
-)
-gail_trainer.train(total_timesteps=100000)
-gail_trainer.gen_algo.save(model_name)
-
-#del gail_trainer
-
-# %%
-model = sb3.PPO.load(model_name)
-
-env = NRasterized(agent=51, width=36, height=36, m_per_px=2, stop_on_collision=False)
-
-obs = env.reset()
-while True:
-    action, _states = model.predict(obs)
-    obs, rewards, done, info = env.step(action)
-    env.render(mode='post')
-    if done:
-        break
-
-env.close(filestr='render/'+model_name)

scratch/etienne/intersimple/gail_options_image.py

@@ -1,8 +1,5 @@
 # %%
-import sys
+from gail.discriminator import CnnDiscriminator
-sys.path.append('../../../')
-
-from src.discriminator import CnnDiscriminator, CnnDiscriminatorFlatAction
 from imitation.algorithms import adversarial
 import stable_baselines3
 import torch.utils.data
@@ -18,17 +15,198 @@ import tempfile
 import pathlib
 from imitation.util import logger
 from stable_baselines3.common.env_util import make_vec_env
-from tqdm import tqdm
-from src.policies.options import OptionsCnnPolicy
-from src.gail.options import OptionsEnv, LLOptions, HLOptions, RenderOptions
-from src.gail.train import train_discriminator, train_generator
 
 model_name = 'gail_options_image'
 env_settings = {'agent': 51, 'width': 36, 'height': 36, 'm_per_px': 2}
 
-ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10]] # option 0 is safe fallback
+ALL_OPTIONS = [(v,t) for v in [0,2,4,6,8] for t in [5, 10, 20]] # option 0 is safe fallback
 
-def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, discount=0.99):
+class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy):
+
+    def __init__(self, observation_space, *args, **kwargs):
+        super().__init__(observation_space['obs'], *args, **kwargs)
+
+    def _prior_distribution(self, s):
+        latent_pi, latent_vf, latent_sde = self._get_latent(s)
+        distribution = self._get_action_dist_from_latent(latent_pi, latent_sde)
+        values = self.value_net(latent_vf)
+        return values, distribution.distribution
+
+    def predict(self, obs):
+        s, m = obs['obs'], obs['mask']
+        values, prior = self._prior_distribution(s)
+        posterior = Categorical(prior.probs * m)
+        ch = posterior.sample()
+        return ch, values, posterior.log_prob(ch)
+
+    def evaluate_actions(self, obs, ch):
+        s, m = obs['obs'], obs['mask']
+        values, prior = self._prior_distribution(s)
+        posterior = Categorical(prior.probs * m)
+        return values, posterior.log_prob(ch), posterior.entropy() # additional values used by PPO.train
+
+def available_actions(env):
+    """Return mask of available actions given current `env` state."""
+    valid = np.array([feasible(env, generate_plan(env, i), i) for i in range(len(ALL_OPTIONS))])
+    return valid
+
+def target_velocity_plan(current_v: float, target_v: float, t: int, dt: float):
+    """Smoothly target a velocity in a given number of steps"""
+    # for now, constant acceleration
+    a = (target_v - current_v)  / (t * dt)
+    return a*np.ones((t,))
+
+def generate_plan(env, i):
+    """Generate input profile for high-level action `i`."""
+    assert i < len(ALL_OPTIONS), "Invalid option index {i}"
+    target_v, t = ALL_OPTIONS[i]
+    current_v = env._env.state[env._agent, 1].item() # extract from env
+    plan = target_velocity_plan(current_v, target_v, t, env._env._dt)
+    assert len(plan) == t, "incorrect plan length"
+    return plan
+
+def check_future_collisions_fast(env, actions):
+    """Checks whether `env._agent` would collide with other agents assuming `actions` as input.
+    
+    Vehicles are (over-)approximated by single circles.
+
+    Args:
+        env (gym.Env): current environment state
+        actions (list of torch.Tensor): list of B (T, nv, adims) T-length action profiles
+    Returns:
+        feasible (torch.Tensor): tensor of shape (B,) indicating whether the respective action profiles are collision-free
+    """
+    B, (T, nv, _) = len(actions), actions[0].shape
+
+    states = torch.stack(env._env.propagate_action_profile(actions), axis=0)
+    assert states.shape == (B, T, nv, 5)
+
+    distance = ((states[:, :, :, :2] - states[:, :, env._agent:env._agent+1, :2])**2).sum(-1).sqrt()
+    distance = torch.where(distance.isnan(), np.inf*torch.ones_like(distance), distance) # only collide with spawned agents
+    distance[:, :, env._agent] = np.inf # cannot collide with itself
+    assert distance.shape == (B, T, nv)
+
+    radius = (env._env._lengths**2 + env._env._widths**2).sqrt() / 2
+    min_distance = radius[env._agent] + radius
+    min_distance = min_distance.unsqueeze(0).unsqueeze(0)
+    assert min_distance.shape == (1, 1, nv)
+
+    return (distance > min_distance).all(-1).all(-1)
+
+def feasible(env, plan, ch):
+    """Check if input profile is feasible given current `env` state. Action `ch=0` is safe fallback."""
+    
+    # zero pad plan - Take (T,) np plan and convert it to (T, nv, 1) torch.Tensor
+    full_plan = torch.zeros(len(plan), env._env._nv, 1)
+    full_plan[:, env._agent, 0] = torch.tensor(plan)
+    valid = check_future_collisions_fast(env, [full_plan]) # check_future_collisions_fast takes in B-list and outputs (B,) bool tensor
+    return ch == 0 or valid.item()
+
+def sample(env, generator, discriminator, level: str):
+    """
+        Sample low-level (state, action, next_state) tuples for discriminator training or 
+        high-level (state, action, reward) tuples for generator training.
+    """
+    done = True
+    while True:
+        episode_start = False
+        if done:
+            s = env.reset()
+            m = available_actions(env)
+            done = False
+            episode_start = True
+
+        obs = {'obs': s, 'mask': m}
+        ch, value, log_prob = generator.policy.predict({
+            'obs': torch.tensor(s).unsqueeze(0).to(generator.policy.device),
+            'mask': torch.tensor(m).unsqueeze(0).to(generator.policy.device),
+        })
+        plan = list(map(float, generate_plan(env, ch)))
+
+        assert not done
+        assert plan
+        assert feasible(env, plan, ch), f'Infeasible hl action {ch}'
+
+        r = 0
+        discount = 1
+        while not done and plan and feasible(env, plan, ch):
+            a, plan = env._normalize(plan[0]), plan[1:]
+            if level == 'high':
+                r += discount * discriminator.discrim_net.discriminator(
+                    torch.tensor(s).unsqueeze(0).to(discriminator.discrim_net.device()),
+                    torch.tensor([[a]]).to(discriminator.discrim_net.device()),
+                )
+                discount *= env.discount
+
+            nexts, _, done, _ = env.step(a)
+            m = available_actions(env)
+
+            if level == 'low':
+                yield {
+                    'obs': s,
+                    'next_obs': nexts,
+                    'acts': np.array((a,)),
+                    'dones': np.array(done),
+                }
+            s = nexts
+
+        if level == 'high':
+            yield {
+                'obs': obs,
+                'option': ch,
+                'reward': r.detach(),
+                'episode_start': episode_start,
+                'value': value.detach(),
+                'log_prob': log_prob.detach(),
+                'done': done,
+            }
+
+def flatten_transitions(transitions):
+    return {
+        'obs': np.stack(list(t['obs'] for t in transitions), axis=0),
+        'next_obs': np.stack(list(t['next_obs'] for t in transitions), axis=0),
+        'acts': np.stack(list(t['acts'] for t in transitions), axis=0),
+        'dones': np.stack(list(t['dones'] for t in transitions), axis=0),
+    }
+
+def train_discriminator(env, generator, discriminator, num_samples):
+    transitions = list(itertools.islice(sample(env, generator, None, 'low'), num_samples))
+    generator_samples = flatten_transitions(transitions)
+    discriminator.train_disc(gen_samples=generator_samples)
+
+def train_generator(env, generator, discriminator, num_samples):
+    generator_samples = list(itertools.islice(sample(env, generator, discriminator, 'high'), num_samples+1))
+    
+    generator.rollout_buffer.reset()
+    for s in generator_samples[:-1]:
+        generator.rollout_buffer.add(
+            obs=s['obs'],
+            action=s['option'].cpu(),
+            reward=s['reward'].cpu(),
+            episode_start=s['episode_start'],
+            value=s['value'],
+            log_prob=s['log_prob'],
+        )
+    
+    generator.rollout_buffer.compute_returns_and_advantage(
+        last_values=generator_samples[-1]['value'],
+        dones=generator_samples[-1]['done'],
+    )
+    
+    generator.train()
+
+class OptionsEnv(gym.Wrapper):
+    
+    def __init__(self, env):
+        super().__init__(env)
+        num_hl_options = len(ALL_OPTIONS)
+        self.action_space = gym.spaces.Discrete(num_hl_options)
+        self.observation_space = gym.spaces.Dict({
+            'obs': env.observation_space,
+            'mask': gym.spaces.Box(low=0, high=1, shape=(num_hl_options,)),
+        })
+
+def train(expert_data, epochs=10, expert_batch_size=32, generator_steps=2048, discount=0.99):
     env = NRasterized(**env_settings)
     env.discount = discount
 
@@ -41,15 +219,14 @@ def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, di
     discriminator = adversarial.GAIL(
         expert_data=expert_data,
         expert_batch_size=expert_batch_size,
-        discrim_kwargs={'discrim_net': CnnDiscriminatorFlatAction(venv)},
+        discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
-        #discrim_kwargs={'discrim_net': CnnDiscriminator(venv)},
         venv=venv, # unused
         gen_algo=stable_baselines3.PPO("CnnPolicy", venv), # unused
     )
 
     generator = stable_baselines3.PPO(
         OptionsCnnPolicy,
-        OptionsEnv(env, options=ALL_OPTIONS),
+        OptionsEnv(env),
         verbose=1,
         n_steps=generator_steps,
     )
@@ -61,30 +238,60 @@ def train(expert_data, epochs=20, expert_batch_size=32, generator_steps=1024, di
         generator.tensorboard_log,
     )
 
-    for _ in tqdm(range(epochs)):
+    for _ in range(epochs):
-        train_discriminator(LLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=expert_batch_size)
+        train_discriminator(env, generator, discriminator, num_samples=expert_batch_size)
-        train_generator(HLOptions(env, options=ALL_OPTIONS), generator, discriminator, num_samples=generator_steps)
+        train_generator(env, generator, discriminator, num_samples=generator_steps)
     
     return generator
 
 # %%
 if __name__ == '__main__':
     # %%
- 
     with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
         trajectories = pickle.load(f)
     transitions = rollout.flatten_trajectories(trajectories)
-    generator = train(transitions)
+    generator = train(transitions, generator_steps=200)
 
     generator.save(model_name)
 
     # %%
     model = stable_baselines3.PPO.load(model_name)
 
-    env = RenderOptions(NRasterized(**env_settings), options=ALL_OPTIONS)
+    env = NRasterized(**env_settings)
 
-    for s in env.sample_ll(model):
+    for transition in sample(env, generator, None, 'low'):
-        if s['dones']:
+        env.render()
+        if transition['dones']:
             break
 
     env.close(filestr='render/'+model_name)
+
+# %% Tests
+
+def test_ll_transitions_vs_expert_data():
+    with open("data/NormalizedIntersimpleExpertMu.001_NRasterizedAgent51w36h36mppx2.pkl", "rb") as f:
+        expert_trajectories = pickle.load(f)
+    expert_transitions = rollout.flatten_trajectories(expert_trajectories)
+
+    env = NRasterized(agent=51, width=36, height=36, m_per_px=2)
+
+    gen_transitions = list(itertools.islice(sample(
+        env=NRasterized(**env_settings),
+        generator=stable_baselines3.PPO(
+            OptionsCnnPolicy,
+            OptionsEnv(env),
+            verbose=1,
+        ),
+        discriminator=None,
+        level='low'
+    ), 10))
+    gen_transitions = flatten_transitions(gen_transitions)
+
+    assert expert_transitions[:10].obs.shape == gen_transitions['obs'].shape
+    assert expert_transitions[:10].next_obs.shape == gen_transitions['next_obs'].shape
+    assert expert_transitions[:10].acts.shape == gen_transitions['acts'].shape
+    assert expert_transitions[:10].dones.shape == gen_transitions['dones'].shape
+
+
+def test_hl_transitions():
+    pass

scratch/etienne/intersimple/gail_options_image_alltracks.py

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

scratch/etienne/intersimple/gail_options_image_random.py

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

scratch/etienne/intersimple/gail_options_image_random_location.py

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

scratch/etienne/intersimple/ppo_speed_lidar.py

@@ -1,50 +0,0 @@
-# %%
-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

@@ -1,49 +0,0 @@
-# %%
-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/render_env_from_model.py

@@ -1,33 +0,0 @@
-
-import stable_baselines3 as sb3
-from intersim.envs.intersimple import NRasterized
-
-
-def render_env(model_name='gail_image_multiagent_nocollision', agent=51, environment=NRasterized):
-    """
-    Render a video from an model, agent, and environment
-    Args:
-        model_name (str): name of the model
-        agent (int): agent to start the video from
-        environment (gym.Env): gym environment class to render environment on
-    """
-
-    model = sb3.PPO.load(model_name)
-
-    env = environment(stop_on_collision=False, width=36, height=36, m_per_px=2, agent=agent)
-
-    obs = env.reset()
-    i=0
-    while True and i < 600:
-        i+=1
-        action, _states = model.predict(obs)
-        obs, rewards, done, info = env.step(action)
-        env.render(mode='post')
-        if done:
-            break
-
-    env.close(filestr='render/'+model_name+'_agent%i'%(agent)) 
-
-if __name__ == '__main__':
-    import fire
-    fire.Fire(render_env)

scratch/etienne/intersimple/train_discrim.py

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

scratch/etienne/intersimple/trpo_speed_lidar.py

@@ -1,52 +0,0 @@
-# %%
-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

@@ -1,78 +0,0 @@
-# %%
-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

@@ -1,74 +0,0 @@
-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

@@ -1,100 +0,0 @@
-# %%
-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

@@ -1,74 +0,0 @@
-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

@@ -1,74 +0,0 @@
-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

@@ -1,97 +0,0 @@
-# %%
-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

@@ -1,101 +0,0 @@
-# %%
-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

@@ -1,54 +0,0 @@
-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

@@ -1,97 +0,0 @@
-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

@@ -1,97 +0,0 @@
-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

@@ -1,104 +0,0 @@
-# %%
-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

@@ -1,39 +0,0 @@
-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

@@ -1,75 +0,0 @@
-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

@@ -1,75 +0,0 @@
-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

@@ -1,102 +0,0 @@
-# %%
-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

@@ -1,55 +0,0 @@
-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

@@ -1,96 +0,0 @@
-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

@@ -1,96 +0,0 @@
-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

@@ -1,103 +0,0 @@
-# %%
-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

@@ -1,54 +0,0 @@
-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

@@ -1,53 +0,0 @@
-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

@@ -1,54 +0,0 @@
-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

@@ -1,54 +0,0 @@
-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

@@ -1,53 +0,0 @@
-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

@@ -1,25 +0,0 @@
-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

@@ -1,63 +0,0 @@
-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

@@ -1,62 +0,0 @@
-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

@@ -1,61 +0,0 @@
-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

@@ -1,41 +0,0 @@
-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

@@ -1,89 +0,0 @@
-# %%
-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

@@ -1,27 +0,0 @@
-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

@@ -1,7 +0,0 @@
-|                     | 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

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

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

@@ -1,111 +0,0 @@
-# %%
-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

@@ -1,110 +0,0 @@
-# %%
-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

@@ -1,88 +0,0 @@
-# %%
-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

@@ -1,87 +0,0 @@
-# %%
-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

@@ -1,62 +0,0 @@
-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

@@ -1,90 +0,0 @@
-# %%
-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

@@ -1,87 +0,0 @@
-# %%
-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

@@ -1,42 +0,0 @@
-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

@@ -1,91 +0,0 @@
-# %%
-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

@@ -1,17 +0,0 @@
-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

@@ -1,30 +0,0 @@
-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

@@ -1,26 +0,0 @@
-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

@@ -1,346 +0,0 @@
-{
- "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

@@ -1,76 +0,0 @@
-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')