def transform_observations(observations,
target_dim=16,
omap: OccupancyMap = None):
new_obs = observations
new_obs["rgb_filled"] = observations["rgb"]
new_obs["taskonomy"] = observations["rgb"]
new_obs["target"] = np.moveaxis(
np.tile(transform_target(observations["pointgoal"]),
(target_dim, target_dim, 1)), -1, 0)
if omap is not None:
new_obs['map'] = omap.construct_occupancy_map()
new_obs['global_pos'] = omap.get_current_global_pos()
del new_obs['rgb']
return new_obs
if not os.path.exists(TRAJ_DIR):
os.makedirs(TRAJ_DIR)
# main loop: collect data
for episode in tqdm(range(len(env.habitat_env.episodes))):
observations = env.reset()
images = []
traj = []
omap = OccupancyMap(initial_pg=observations['pointgoal'], map_kwargs=DEFAULT_MAP_KWARGS)
while not env.habitat_env.episode_over:
# postprocess and log (state, action) pairs
observations['rgb_filled'] = observations['rgb']
observations['target'] = np.moveaxis(np.tile(transform_target(observations['pointgoal']), (target_dim, target_dim, 1)), -1, 0)
observations['map'] = omap.construct_occupancy_map()
del observations['rgb']
# agent step
best_action = follower.get_next_action(env.habitat_env.current_episode.goals[0].position).value
traj.append([observations, best_action])
# env step
observations, reward, done, info = env.step(best_action)
omap.add_pointgoal(observations['pointgoal']) # s_{t+1}
omap.step(best_action) # a_t
# viz
if SAVE_VIDEO:
im = observations["rgb"]
top_down_map = draw_top_down_map(