Update scratch for plotting

scratch/johannes/intersimple/idm.py

@@ -3,13 +3,14 @@ import torch
 from src.baselines.rule_policies import IDMRulePolicy
 from tqdm import tqdm
 
-from intersim.envs import NRasterizedIncrementingAgent, NRasterizedRandomAgent, NRasterized
+from intersim.envs import NRasterizedIncrementingAgent, NRasterizedRandomAgent, NRasterized,IntersimpleLidarFlat
 from intersim.envs.intersimple import speed_reward
 import functools
 
 
-env = NRasterizedIncrementingAgent(
-    # agent = 4,
+env = IntersimpleLidarFlat(
+    agent = 51,
+    n_rays=5,
     reward=functools.partial(
         speed_reward,
         collision_penalty=1000
@@ -20,11 +21,12 @@ policy = IDMRulePolicy(env)
 
 colliding_agents = []
 
-for agent in range(151):
-    print("Start agent", agent)
-    obs = env.reset()
-    env.render(mode='post')
-    for i in range(300):
+# 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)
@@ -35,7 +37,7 @@ for agent in range(151):
             colliding_agents.append(agent)
             print("  Collision")
         break
-    env.close(filestr='idm/agent_{}'.format(agent))
+env.close(filestr='idm3/agent_{}'.format(agent))
 
 print(len(colliding_agents), "colliding_agents")
 print(colliding_agents)