Only take targets driving roughly in the same direction as IDM target

2022-02-28 11:37:08 +01:00
43 changed files with 368 additions and 188 deletions
--- a/bc-experiment.py
+++ b/bc-experiment.py
@@ -142,7 +142,7 @@ if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--train', choices=['A', 'B'])
-    parser.add_argument('--epochs', type=int, default=500)
+    parser.add_argument('--epochs', type=int, default=1000)
    parser.add_argument('--test', type=str, help='path to config file to run final training on')
    parser.add_argument('--test_seeds', type=int, default=5)
    parser.add_argument('--test_cpus', type=int, help='number of cpus available to split test seed training over')
@@ -162,10 +162,10 @@ if __name__ == '__main__':
                },
                'policy': {
                    'learning_rate': 3e-4, 
-                    'learning_rate_decay': tune.grid_search([0.999, 1.0]), 
-                    'hidden_layer_size': tune.grid_search([10, 20, 40]),
+                    'learning_rate_decay': 1.0, 
+                    'hidden_layer_size': tune.grid_search([20, 40]),
                    'n_hidden_layers': tune.grid_search([2, 3]), 
-                    'activation':tune.grid_search([0, 1]),
+                    'activation':0,
                },
                'train_epochs': args.epochs,
                'seed': 0,
@@ -209,6 +209,4 @@ if __name__ == '__main__':
            s = analysis._checkpoints[i]['config']['seed']
            check_dir = analysis._checkpoints[i]['logdir']
            shutil.copyfile(os.path.join(check_dir,'policy_final.pt'), 
-                os.path.join(savepath, f'policy_seed{s}.pt'))
-            shutil.copyfile(os.path.join(check_dir,'params.json'),
-                os.path.join(savepath, 'config.json')) # copy config automatically
+                os.path.join(savepath, f'policy_seed{s}.pt'))
--- a/test_policies/gail/expA/config.json
+++ b/test_policies/gail/expA/config.json
@@ -6,7 +6,7 @@
    "policy": {
        "learning_rate": 0.0003,
        "learning_rate_decay": 1.0,
-        "clip_ratio": 0.2,
+        "delta": 0.01,
        "iterations_per_epoch": 100,
        "hidden_layer_size": 40,
        "n_hidden_layers": 2,
--- a/test_policies/gail/expB/config.json
+++ b/test_policies/gail/expB/config.json
@@ -6,7 +6,7 @@
    "policy": {
        "learning_rate": 0.0003,
        "learning_rate_decay": 1.0,
-        "clip_ratio": 0.2,
+        "delta": 0.01,
        "iterations_per_epoch": 100,
        "hidden_layer_size": 40,
        "n_hidden_layers": 2,
--- a/best_configs/hail-etienne2_expA.json
+++ b/best_configs/hail-etienne2_expA.json
@@ -0,0 +1,33 @@
+{
+    "experiment": "A",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": null,
+        "abort_unsafe_collision_method": null
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 10,
+        "n_hidden_layers": 3,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}
--- a/best_configs/hail-etienne2_expB.json
+++ b/best_configs/hail-etienne2_expB.json
@@ -0,0 +1,33 @@
+{
+    "experiment": "B",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": null,
+        "abort_unsafe_collision_method": null
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 10,
+        "n_hidden_layers": 3,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}
--- a/best_configs/hail-etienne3_expA.json
+++ b/best_configs/hail-etienne3_expA.json
@@ -11,7 +11,7 @@
        "clip_ratio": 0.2,
        "iterations_per_epoch": 100,
        "hidden_layer_size": 40,
-        "n_hidden_layers": 2,
+        "n_hidden_layers": 3,
        "activation": 0,
        "option": 0
    },
@@ -23,11 +23,11 @@
        "learning_rate": 0.001,
        "weight_decay": 0.0001,
        "iterations_per_epoch": 100,
-        "n_hidden_layers_element": 4,
-        "n_hidden_layers_global": 1,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
        "hidden_layer_size": 10,
        "activation": 0
    },
-    "train_epochs": 90,
+    "train_epochs": 100,
    "seed": 0
 }
--- a/best_configs/hail-etienne3_expB.json
+++ b/best_configs/hail-etienne3_expB.json
@@ -0,0 +1,33 @@
+{
+    "experiment": "B",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": null,
+        "abort_unsafe_collision_method": null
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 40,
+        "n_hidden_layers": 3,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}
--- a/best_configs/hail-etienne_expA.json
+++ b/best_configs/hail-etienne_expA.json
@@ -0,0 +1,33 @@
+{
+    "experiment": "A",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": null,
+        "abort_unsafe_collision_method": null
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 20,
+        "n_hidden_layers": 4,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}
--- a/best_configs/hail-etienne_expB.json
+++ b/best_configs/hail-etienne_expB.json
@@ -11,7 +11,7 @@
        "clip_ratio": 0.2,
        "iterations_per_epoch": 100,
        "hidden_layer_size": 20,
-        "n_hidden_layers": 2,
+        "n_hidden_layers": 4,
        "activation": 0,
        "option": 0
    },
@@ -23,11 +23,11 @@
        "learning_rate": 0.001,
        "weight_decay": 0.0001,
        "iterations_per_epoch": 100,
-        "n_hidden_layers_element": 4,
+        "n_hidden_layers_element": 3,
        "n_hidden_layers_global": 2,
        "hidden_layer_size": 10,
        "activation": 0
    },
-    "train_epochs": 85,
+    "train_epochs": 100,
    "seed": 0
 }
--- a/best_configs/shail-etienne2_expA.json
+++ b/best_configs/shail-etienne2_expA.json
@@ -0,0 +1,33 @@
+{
+    "experiment": "A",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": "circle",
+        "abort_unsafe_collision_method": "circle"
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 10,
+        "n_hidden_layers": 3,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}
--- a/best_configs/shail-etienne2_expB.json
+++ b/best_configs/shail-etienne2_expB.json
@@ -0,0 +1,33 @@
+{
+    "experiment": "B",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": "circle",
+        "abort_unsafe_collision_method": "circle"
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 10,
+        "n_hidden_layers": 3,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}
--- a/best_configs/shail-etienne3_expA.json
+++ b/best_configs/shail-etienne3_expA.json
@@ -11,7 +11,7 @@
        "clip_ratio": 0.2,
        "iterations_per_epoch": 100,
        "hidden_layer_size": 40,
-        "n_hidden_layers": 2,
+        "n_hidden_layers": 3,
        "activation": 0,
        "option": 0
    },
@@ -23,11 +23,11 @@
        "learning_rate": 0.001,
        "weight_decay": 0.0001,
        "iterations_per_epoch": 100,
-        "n_hidden_layers_element": 4,
-        "n_hidden_layers_global": 1,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
        "hidden_layer_size": 10,
        "activation": 0
    },
-    "train_epochs": 90,
+    "train_epochs": 100,
    "seed": 0
 }
--- a/best_configs/shail-etienne3_expB.json
+++ b/best_configs/shail-etienne3_expB.json
@@ -0,0 +1,33 @@
+{
+    "experiment": "B",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": "circle",
+        "abort_unsafe_collision_method": "circle"
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 40,
+        "n_hidden_layers": 3,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}
--- a/best_configs/shail-etienne_expA.json
+++ b/best_configs/shail-etienne_expA.json
@@ -0,0 +1,33 @@
+{
+    "experiment": "A",
+    "trainenv": {
+        "stop_on_collision": false,
+        "safe_actions_collision_method": "circle",
+        "abort_unsafe_collision_method": "circle"
+    },
+    "policy": {
+        "learning_rate": 0.0003,
+        "learning_rate_decay": 1.0,
+        "clip_ratio": 0.2,
+        "iterations_per_epoch": 100,
+        "hidden_layer_size": 20,
+        "n_hidden_layers": 4,
+        "activation": 0,
+        "option": 0
+    },
+    "value": {
+        "learning_rate": 0.001,
+        "iterations_per_epoch": 1000
+    },
+    "discriminator": {
+        "learning_rate": 0.001,
+        "weight_decay": 0.0001,
+        "iterations_per_epoch": 100,
+        "n_hidden_layers_element": 3,
+        "n_hidden_layers_global": 2,
+        "hidden_layer_size": 10,
+        "activation": 0
+    },
+    "train_epochs": 100,
+    "seed": 0
+}
--- a/best_configs/shail-etienne_expB.json
+++ b/best_configs/shail-etienne_expB.json
@@ -11,7 +11,7 @@
        "clip_ratio": 0.2,
        "iterations_per_epoch": 100,
        "hidden_layer_size": 20,
-        "n_hidden_layers": 2,
+        "n_hidden_layers": 4,
        "activation": 0,
        "option": 0
    },
@@ -23,11 +23,11 @@
        "learning_rate": 0.001,
        "weight_decay": 0.0001,
        "iterations_per_epoch": 100,
-        "n_hidden_layers_element": 4,
+        "n_hidden_layers_element": 3,
        "n_hidden_layers_global": 2,
        "hidden_layer_size": 10,
        "activation": 0
    },
-    "train_epochs": 85,
+    "train_epochs": 100,
    "seed": 0
 }
--- a/eval_experiments.py
+++ b/eval_experiments.py
@@ -1,23 +1,17 @@
 import os
 from src.eval_main import eval_main
 from src.evaluation.utils import load_and_average
-import torch
-import json
-
-activations = [torch.nn.Tanh, torch.nn.LeakyReLU]

 def main(method:str='expert', folder:str=None, locations=[(0,0)], skip_running=False):
    
-    exclude_keys_from_policy_kwargs = {'learning_rate', 'learning_rate_decay', 'clip_ratio', 'iterations_per_epoch', 'option'}
    policy_kwargs = {}
-
    if method in ['expert', 'idm']:
        env, env_kwargs ='NRasterizedRouteIncrementingAgent', {}
    elif method in ['bc','gail']:
        env='NormalizedContinuousEvalEnv' 
        env_kwargs={'stop_on_collision':True, 'max_episode_steps':1000}
    elif method in ['hail']:
-        env = 'NormalizedSafeOptionsEvalEnv'
+        env = 'NormalizedOptionsEvalEnv'
        env_kwargs={'stop_on_collision':True, 'max_episode_steps':1000, 'safe_actions_collision_method': None, 'abort_unsafe_collision_method': None}
    elif method in ['shail']:
        env = 'NormalizedSafeOptionsEvalEnv'
@@ -29,18 +23,7 @@ def main(method:str='expert', folder:str=None, locations=[(0,0)], skip_running=F

    if folder is not None:
        files = [os.path.join(folder, f) for f in os.listdir(folder) if os.path.isfile(os.path.join(folder, f))]
-        files = [f for f in files if f.endswith('.pt')]
-        with open(os.path.join(folder, 'config.json'), 'rb') as f:
-            config = json.load(f)
-        print('%i policy files found in %s folder' %(len(files), folder))
-        print('found policy config', config['policy'])
-
-        policy_config = {k: v for k, v in config['policy'].items() if k not in exclude_keys_from_policy_kwargs}
-        policy_config['activation'] = activations[policy_config['activation']]
-        print('final policy config', policy_config)
-
-        policy_kwargs.update(policy_config)
-        print('final policy kwargs', policy_kwargs)
+        print('%i folders found in %s folder' %(len(files), folder))

    if not skip_running:
        for policy_file in files:
@@ -77,7 +60,7 @@ def latex_print(am, light=False):
    print('success rate, distance travelled, RWSE_10, |DeltaV|, AccelJSD')
    if light:
        if 'rwse_10s' in am.keys():
-            print("%2.1f& %2.1f & %2.1f & %1.2f& "
+            print("%2.1f& %2.1f & %1.2f & %2.1f& "
            "%0.3f \\\\" %( 100*am['success rate'][0], am['mean travel distance'][0], am['rwse_10s'][0],
                            am['average absolute average velocity'][0],am['acceleration distribution divergence'][0] ))
            return
@@ -88,7 +71,7 @@ def latex_print(am, light=False):
        return
    
    print("%2.1f \\scriptstyle\\pm %2.1f & %2.1f \\scriptstyle\\pm %2.1f & "
-        "%2.1f \\scriptstyle\\pm %1.1f & %1.2f \\scriptstyle\\pm %1.2f & "
+        "%1.2f \\scriptstyle\\pm %1.2f & %2.1f \\scriptstyle\\pm %1.1f & "
        "%0.3f \\scriptstyle\\pm %0.3f \\\\" %( 100*am['success rate'][0], 100*am['success rate'][1],
               am['mean travel distance'][0] , am['mean travel distance'][1]  ,
               am['rwse_10s'][0] , am['rwse_10s'][1]  ,
--- a/gail-experiment.py
+++ b/gail-experiment.py
@@ -235,6 +235,4 @@ if __name__ == '__main__':
            s = analysis._checkpoints[i]['config']['seed']
            check_dir = analysis._checkpoints[i]['logdir']
            shutil.copyfile(os.path.join(check_dir,'policy_final.pt'), 
-                os.path.join(savepath, f'policy_seed{s}.pt'))
-            shutil.copyfile(os.path.join(check_dir,'params.json'),
-                os.path.join(savepath, 'config.json')) # copy config automatically
+                os.path.join(savepath, f'policy_seed{s}.pt'))
--- a/out/hail/expA/loc_r0t0/policy_seed1_tseed0_comparison.pkl
+++ b/out/hail/expA/loc_r0t0/policy_seed1_tseed0_comparison.pkl
--- a/out/hail/expA/loc_r0t0/policy_seed1_tseed0_summary.pkl
+++ b/out/hail/expA/loc_r0t0/policy_seed1_tseed0_summary.pkl
--- a/out/hail/expA/loc_r0t0/policy_seed2_tseed0_comparison.pkl
+++ b/out/hail/expA/loc_r0t0/policy_seed2_tseed0_comparison.pkl
--- a/out/hail/expA/loc_r0t0/policy_seed2_tseed0_summary.pkl
+++ b/out/hail/expA/loc_r0t0/policy_seed2_tseed0_summary.pkl
--- a/out/hail/expA/loc_r0t0/policy_seed3_tseed0_comparison.pkl
+++ b/out/hail/expA/loc_r0t0/policy_seed3_tseed0_comparison.pkl
--- a/out/hail/expA/loc_r0t0/policy_seed3_tseed0_summary.pkl
+++ b/out/hail/expA/loc_r0t0/policy_seed3_tseed0_summary.pkl
--- a/out/hail/expA/loc_r0t0/policy_seed4_tseed0_comparison.pkl
+++ b/out/hail/expA/loc_r0t0/policy_seed4_tseed0_comparison.pkl
--- a/out/hail/expA/loc_r0t0/policy_seed4_tseed0_summary.pkl
+++ b/out/hail/expA/loc_r0t0/policy_seed4_tseed0_summary.pkl
--- a/out/hail/expA/loc_r0t0/policy_seed5_tseed0_comparison.pkl
+++ b/out/hail/expA/loc_r0t0/policy_seed5_tseed0_comparison.pkl
--- a/out/hail/expA/loc_r0t0/policy_seed5_tseed0_summary.pkl
+++ b/out/hail/expA/loc_r0t0/policy_seed5_tseed0_summary.pkl
--- a/out/hail/expB/loc_r0t4/policy_seed1_tseed0_comparison.pkl
+++ b/out/hail/expB/loc_r0t4/policy_seed1_tseed0_comparison.pkl
--- a/out/hail/expB/loc_r0t4/policy_seed1_tseed0_summary.pkl
+++ b/out/hail/expB/loc_r0t4/policy_seed1_tseed0_summary.pkl
--- a/out/hail/expB/loc_r0t4/policy_seed2_tseed0_comparison.pkl
+++ b/out/hail/expB/loc_r0t4/policy_seed2_tseed0_comparison.pkl
--- a/out/hail/expB/loc_r0t4/policy_seed2_tseed0_summary.pkl
+++ b/out/hail/expB/loc_r0t4/policy_seed2_tseed0_summary.pkl
--- a/out/hail/expB/loc_r0t4/policy_seed3_tseed0_comparison.pkl
+++ b/out/hail/expB/loc_r0t4/policy_seed3_tseed0_comparison.pkl
--- a/out/hail/expB/loc_r0t4/policy_seed3_tseed0_summary.pkl
+++ b/out/hail/expB/loc_r0t4/policy_seed3_tseed0_summary.pkl
--- a/out/hail/expB/loc_r0t4/policy_seed4_tseed0_comparison.pkl
+++ b/out/hail/expB/loc_r0t4/policy_seed4_tseed0_comparison.pkl
--- a/out/hail/expB/loc_r0t4/policy_seed4_tseed0_summary.pkl
+++ b/out/hail/expB/loc_r0t4/policy_seed4_tseed0_summary.pkl
--- a/out/hail/expB/loc_r0t4/policy_seed5_tseed0_comparison.pkl
+++ b/out/hail/expB/loc_r0t4/policy_seed5_tseed0_comparison.pkl
--- a/out/hail/expB/loc_r0t4/policy_seed5_tseed0_summary.pkl
+++ b/out/hail/expB/loc_r0t4/policy_seed5_tseed0_summary.pkl
--- a/scratch/johannes/intersimple/idm.py
+++ b/scratch/johannes/intersimple/idm.py
@@ -1,44 +0,0 @@
-# %%
-import torch
-from src.baselines.rule_policies import IDMRulePolicy
-from tqdm import tqdm
-
-from intersim.envs import NRasterizedIncrementingAgent, NRasterizedRandomAgent, NRasterized,IntersimpleLidarFlat
-from intersim.envs.intersimple import speed_reward
-import functools
-
-
-env = IntersimpleLidarFlat(
-    agent = 51,
-    n_rays=5,
-    reward=functools.partial(
-        speed_reward,
-        collision_penalty=1000
-    ),
-    stop_on_collision=True,
-)
-policy = IDMRulePolicy(env)
-
-colliding_agents = []
-
-# for agent in range(151):
-agent = env._agent
-print("Start agent", agent)
-obs = env.reset()
-env.render(mode='post')
-for i in range(300):
-    action, _ = policy.predict(torch.tensor(obs))
-    # action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
-    obs, reward, done, _ = env.step(action)
-    env.render(mode='post')
-    # print('step', i, 'reward', reward)
-    if done:
-        if reward < -500:
-            colliding_agents.append(agent)
-            print("  Collision")
-        break
-env.close(filestr='idm3/agent_{}'.format(agent))
-
-print(len(colliding_agents), "colliding_agents")
-print(colliding_agents)
-# %%
--- a/shail-experiment.py
+++ b/shail-experiment.py
@@ -247,17 +247,8 @@ if __name__ == '__main__':
            os.makedirs(savepath)
        
        import shutil
-
-        # save config
-        shutil.copyfile(
-            args.test,
-            os.path.join(savepath, 'config.json')
-        )
-
        for i in range(args.test_seeds):
            s = analysis._checkpoints[i]['config']['seed']
            check_dir = analysis._checkpoints[i]['logdir']
            shutil.copyfile(os.path.join(check_dir,'policy_final.pt'), 
-                os.path.join(savepath, f'policy_seed{s}.pt'))
-            shutil.copyfile(os.path.join(check_dir,'params.json'),
-                os.path.join(savepath, 'config.json')) # copy config automatically
+                os.path.join(savepath, f'policy_seed{s}.pt'))
--- a/src/baselines/rule_policies.py
+++ b/src/baselines/rule_policies.py
@@ -82,16 +82,15 @@ class IDMRulePolicy(BaseAlgorithm):
        self._env = env
        self.t_future = t_future
        self.half_angle = half_angle
+        self.max_heading_diff = 120

        # Default IDM parameters
        assert target_speed>0, 'negative target speed'
-        self.v_max = target_speed
+        self.s_max = target_speed
        self.a_max = np.array([3.]) # nominal acceleration
        self.tau = 0.5 # desired time headway
        self.b_pref = 2.5 # preferred deceleration
-        self.d_min = 3 #minimum spacing
-        self.max_pos_error = 2 # m, for matching vehicles to ego path
-        self.max_deg_error = 30 # degree, for matching vehicles to ego path
+        self.d_min = 1 #minimum spacing

        # for np.remainder nan warnings
        np.seterr(invalid='ignore')
@@ -126,63 +125,37 @@ class IDMRulePolicy(BaseAlgorithm):
            action (np.ndarray): action for controlled agent to take
        """
        agent = self._env._agent
-        state = self._env._env.state.numpy()
        full_state = self._env._env.projected_state.numpy() #(nv, 5)
        ego_state = full_state[agent] # (5,)
-        v_ego = ego_state[2]
+        s = ego_state[2]
+        xy = full_state[:,0:2] # (nv, 2)
        v = full_state[:,2:3] # (nv, 1)
+        psi = full_state[:,3:4] # (nv, 1)

-        length = 20
-        step = 0.1
-        x, y = self._env._env._generate_paths(delta=step, n=length/step, is_distance=True)
-        heading = to_circle(np.arctan2(np.diff(y), np.diff(x)))
+        d, r, i = self.get_ego_dr(agent, xy, v, psi)

-        paths = np.stack([x[:,:-1],y[:,:-1], heading], axis=1) # (nv, 3, (path_length-1))
-        ego_path = paths[agent:agent+1] # (1, 3, path_length-1)
+        # propagate environment forward at constant velocity
+        for t in self.t_future:
+            if t > 0:
+                xy2 = xy + t * v * np.vstack((np.cos(psi[:,0]), np.sin(psi[:,0]))).T
+                d2, r2, i2 = self.get_ego_dr(agent, xy2, v, psi)

-        # (x,y,phi) of all vehicles
-        poses = np.expand_dims(full_state[:, [0,1,3]], 2) # (nv, 3, 1)
+                # choose closer vehicle (now vs imagined)
+                if d2 < d:
+                    d, r, i = d2, r2, i2

-        diff = ego_path - poses
-        diff[:, 2, :] = to_circle(diff[:, 2, :])
+        # Update environment interaction graph with i
+        if i:
+            self._env._env._graph._neighbor_dict={agent:[i]}

-        # Test if position and heading angle are close for some point on the future vehicle track
-        pos_close = np.sum(diff[:, 0:2, :]**2, 1) <= self.max_pos_error**2 # (nv, path_length-1)
-        heading_close = np.abs(diff[:, 2, :]) <= self.max_deg_error * np.pi / 180 # (nv, path_length-1)
-        # For all vehicles get the path points where they are close to the ego path
-        close = np.logical_and(pos_close, heading_close) # (nv, path_length-1)
-        close[agent, :] = False # exclude ego agent
-
-        leader = agent
-        min_idx = np.Inf
-        # Determine vehicle that is closest to ego in terms of path coordinate
-        for veh_id in range(len(close)):
-            path_idx = np.nonzero(close[veh_id])[0]
-            # veh_id is never close to agent
-            if len(path_idx) == 0:
-                continue
-            # first path index where veh_id is close to agent
-            elif path_idx[0] < min_idx:
-                leader = veh_id
-                min_idx = path_idx[0]
-
-        if leader != agent:
-            # distance along ego path to point with closest distance
-            d = step * min_idx
-
-            # Update environment interaction graph with leader
-            self._env._env._graph._neighbor_dict={agent:[leader]}
-
-            delta_v = v_ego - v[leader, 0]
-            d_des = self.d_min + self.tau * v_ego + v_ego * delta_v / (2* (self.a_max*self.b_pref)**0.5 )
-            d_des = max(d_des, self.d_min)
-        else:
-            d = np.Inf
+        if d == np.inf:
            d_des = self.d_min
-            self._env._env._graph._neighbor_dict={}
+        else:
+            d_des = self.d_min + self.tau * s + s * r / (2* (self.a_max*self.b_pref)**0.5 )
+            d_des = max(d_des, self.d_min)

        assert (d_des>= self.d_min)
-        action = self.a_max*(1 - (v_ego/self.v_max)**4 - (d_des/d)**2)
+        action = self.a_max*(1 - (s/self.s_max)**4 - (d_des/d)**2)

        # normalize action to range if env is a NormalizedActionSpace
        if isinstance(self._env, NormalizedActionSpace):
@@ -191,6 +164,53 @@ class IDMRulePolicy(BaseAlgorithm):
        assert action.shape==(1,)
        return action

+    def get_ego_dr(self, agent:int, xy: np.ndarray,
+        v: np.ndarray, psi: np.ndarray) -> Tuple[float, float, Optional[int]]:
+        """
+        Return distance and relative speed of closest car within half angle from heading
+
+        Args:
+            agent (int): agent index
+            xy (np.ndarray): (nv, 2) x and y positions
+            v (np.ndarray): (nv, 1) velocity
+            psi (np.ndarray): (nv, 1) heading angle
+
+        Returns:
+            d (float): distance to closest vehicle in cone
+            r (float): relative speed between the two vehicles
+            i (Optional[int]): index of closest vehicle, or None
+        """
+        nv, nxy = xy.shape
+        nv2, nvel = v.shape
+        nv3, npsi = psi.shape
+        assert nv==nv2==nv3
+        assert nxy==2
+        assert nvel==npsi==1
+
+        dxys = xy - xy[agent] # (nv, 2)
+        ds = np.linalg.norm(dxys,axis=1) # (nv,)
+        df = (dxys*np.hstack((np.cos(psi),np.sin(psi)))).sum(-1) # (nv, )
+        dl = (dxys*np.hstack((-np.sin(psi), np.cos(psi)))).sum(-1) # (nv, )
+        alpha = to_circle(np.arctan2(dl, df))
+
+        heading_diff = to_circle(psi - psi[agent]).flatten()
+
+        val_idx = np.arange(nv)[
+            (np.abs(alpha) < self.half_angle*np.pi/180) & (np.arange(nv) != agent) & (np.abs(heading_diff) < self.max_heading_diff*np.pi/180)
+            ]
+
+        if len(val_idx)==0:
+            i = None
+            d = float('inf')
+            r = float('inf')
+        else:
+            idx = np.argmin(ds[val_idx]) # closest car which meets requirements
+            i = int(val_idx[idx])
+            d = ds[i]
+            r = v[i,0]-v[agent,0]
+
+        return d, r, i
+
 def to_circle(x: np.ndarray) -> np.ndarray:
    """
    Casts x (in rad) to [-pi, pi)
--- a/src/eval_main.py
+++ b/src/eval_main.py
@@ -41,33 +41,33 @@ def load_policy(method:str,
    if method == 'idm':
        policy = IDMRulePolicy(env, **policy_kwargs)
    elif method == 'bc':
-        policy = SetPolicy(env.action_space.shape[-1], **policy_kwargs)
+        policy = SetPolicy(env.action_space.shape[-1])
        policy.load_state_dict(torch.load(policy_file, map_location=ml))
        policy.eval()
    elif method == 'gail-trpo':
-        policy = SetPolicy(env.action_space.shape[-1], **policy_kwargs)
+        policy = SetPolicy(env.action_space.shape[-1])
        policy(torch.zeros(env.observation_space.shape))
        policy = ReparamPolicy(policy)
        policy.load_state_dict(torch.load(policy_file, map_location=ml))
        policy.eval()
    elif method == 'gail':
-        policy = SetPolicy(env.action_space.shape[-1], **policy_kwargs)
+        policy = SetPolicy(env.action_space.shape[-1])
        policy.load_state_dict(torch.load(policy_file, map_location=ml))
        policy.eval()
    elif method == 'rail':
        raise NotImplementedError
    elif method == 'hail-trpo':
-        policy = SetMaskedDiscretePolicy(env.action_space.n, **policy_kwargs)
+        policy = SetDiscretePolicy(env.action_space.n)
        policy(torch.zeros(env.observation_space.shape))
        policy = ReparamPolicy(policy)
        policy.load_state_dict(torch.load(policy_file, map_location=ml))
        policy.eval()
    elif method == 'hail':
-        policy = SetMaskedDiscretePolicy(env.action_space.n, **policy_kwargs)
+        policy = SetDiscretePolicy(env.action_space.n)
        policy.load_state_dict(torch.load(policy_file, map_location=ml))
        policy.eval()
    elif method == 'shail-trpo':
-        policy = SetMaskedDiscretePolicy(env.action_space.n, **policy_kwargs)
+        policy = SetMaskedDiscretePolicy(env.action_space.n)
        policy(
            torch.zeros(env.observation_space['observation'].shape),
            torch.zeros(env.observation_space['safe_actions'].shape)
@@ -76,7 +76,7 @@ def load_policy(method:str,
        policy.load_state_dict(torch.load(policy_file, map_location=ml))
        policy.eval()
    elif method == 'shail':
-        policy = SetMaskedDiscretePolicy(env.action_space.n, **policy_kwargs)
+        policy = SetMaskedDiscretePolicy(env.action_space.n)
        policy.load_state_dict(torch.load(policy_file, map_location=ml))
        policy.eval()
    else:
--- a/test_policies/bc/expA/config.json
+++ b/test_policies/bc/expA/config.json
@@ -1,15 +0,0 @@
-{
-    "experiment": "A",
-    "trainenv": {
-        "stop_on_collision": false
-    },
-    "policy": {
-        "learning_rate": 0.0003,
-        "learning_rate_decay": 1.0,
-        "hidden_layer_size": 40,
-        "n_hidden_layers": 2,
-        "activation": 0
-    },
-    "train_epochs": 300,
-    "seed": 0
-}
--- a/test_policies/bc/expB/config.json
+++ b/test_policies/bc/expB/config.json
@@ -1,15 +0,0 @@
-{
-    "experiment": "B",
-    "trainenv": {
-        "stop_on_collision": false
-    },
-    "policy": {
-        "learning_rate": 0.0003,
-        "learning_rate_decay": 1.0,
-        "hidden_layer_size": 40,
-        "n_hidden_layers": 2,
-        "activation": 0
-    },
-    "train_epochs": 300,
-    "seed": 0
-}