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')
-# 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)
 # %%