plt.title(f"Box for class {class_}")
plt.show()
classes = np.array(classes)
return boxes, classes
if __name__ == "__main__":
DEBUG = False
NPZ_INDEX = 0
DATA_TO_COLLECT = 1000 # we want to get 1000 data points
seed(123)
environment = launch_env()
policy = PurePursuitPolicy(environment)
MAX_STEPS = 500
SAMPLE_FREQ = 10
while NPZ_INDEX < DATA_TO_COLLECT:
obs = environment.reset()
environment.render(segment=True)
rewards = []
nb_of_steps = 0
while NPZ_INDEX < DATA_TO_COLLECT:
nb_of_steps = 0
# we interate over several maps to get more diverse data
possible_maps = [
"loop_pedestrians", "udem1", "loop_dyn_duckiebots", "zigzag_dists"
]
env_id = 0
env = None
while True:
if env is not None:
env.window.close()
env.close()
if env_id >= len(possible_maps):
env_id = env_id % len(possible_maps)
env = launch_env(possible_maps[env_id])
policy = PurePursuitPolicy(env)
obs = env.reset()
inner_steps = 0
if nb_of_steps >= MAX_STEPS:
break
while True:
if nb_of_steps >= MAX_STEPS or inner_steps > 100:
break
action = policy.predict(np.array(obs))
obs, rew, done, misc = env.step(action)
seg = env.render_obs(True)