Update IDM script

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')
+from intersim.envs import NRasterizedIncrementingAgent, NRasterizedRandomAgent, NRasterized
-# 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.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()
+colliding_agents = []
-env.render(mode='post')
+
-for i in range(1000):
+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)
+        # print('step', i, 'reward', reward)
         if done:
-        obs = env.reset()
+            if reward < -500:
-env.close()
+                collising_agents.append(agent)
+                print("  Collision")
+            break
+    env.close(filestr='idm/agent_{}'.format(agent))
 
+print(len(colliding_agents), "colliding_agents")
+print(colliding_agents)
 # %%

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)