From a3280893afb5a4affe7f419305c8a13eaf64ba0a Mon Sep 17 00:00:00 2001
From: Johannes Fischer <johannes.fischer@kit.edu>
Date: Wed, 23 Feb 2022 18:16:01 +0100
Subject: [PATCH] Update IDM script

---
 scratch/johannes/intersimple/idm.py | 63 +++++++++++------------------
 src/baselines/rule_policies.py      |  1 +
 2 files changed, 25 insertions(+), 39 deletions(-)

diff --git a/scratch/johannes/intersimple/idm.py b/scratch/johannes/intersimple/idm.py
index 1d1c223..8853255 100644
--- a/scratch/johannes/intersimple/idm.py
+++ b/scratch/johannes/intersimple/idm.py
@@ -3,55 +3,40 @@ import torch
 from src.baselines.rule_policies import IDMRulePolicy
 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 NRasterizedIncrementingAgent, NRasterizedRandomAgent
+from intersim.envs import NRasterizedIncrementingAgent, NRasterizedRandomAgent, NRasterized
 from intersim.envs.intersimple import speed_reward
 import functools
 
 
-env = NRasterizedRandomAgent(
+env = NRasterizedIncrementingAgent(
     # agent = 4,
     reward=functools.partial(
         speed_reward,
-        collision_penalty=0
+        collision_penalty=1000
     ),
-    stop_on_collision=False,
+    stop_on_collision=True,
 )
 policy = IDMRulePolicy(env)
 
-obs = env.reset()
-env.render(mode='post')
-for i in range(1000):
-    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:
-        obs = env.reset()
-env.close()
+colliding_agents = []
 
+for agent in range(151):
+    print("Start agent", agent)
+    obs = env.reset()
+    env.render(mode='post')
+    for i in range(300):
+        action, _ = policy.predict(torch.tensor(obs))
+        # action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
+        obs, reward, done, _ = env.step(action)
+        env.render(mode='post')
+        # print('step', i, 'reward', reward)
+        if done:
+            if reward < -500:
+                collising_agents.append(agent)
+                print("  Collision")
+            break
+    env.close(filestr='idm/agent_{}'.format(agent))
+
+print(len(colliding_agents), "colliding_agents")
+print(colliding_agents)
 # %%
diff --git a/src/baselines/rule_policies.py b/src/baselines/rule_policies.py
index ecc807c..9888b5d 100644
--- a/src/baselines/rule_policies.py
+++ b/src/baselines/rule_policies.py
@@ -204,6 +204,7 @@ class IDMRulePolicy(BaseAlgorithm):
         else:
             d = np.Inf
             d_des = self.d_min
+            self._env._env._graph._neighbor_dict={}
 
         assert (d_des>= self.d_min)
         action = self.a_max*(1 - (v_ego/self.v_max)**4 - (d_des/d)**2)