def test_pyrobot(mocker): if "pyrobot" not in sys.modules: # Mock pyrobot package if it is not installed mock_pyrobot = mocker.MagicMock() mock_pyrobot.Robot = RobotMock sys.modules["pyrobot"] = mock_pyrobot # Re-register pyrobot with mock from habitat.sims.registration import _try_register_pyrobot _try_register_pyrobot() config = get_config() with make_sim("PyRobot-v0", config=config.PYROBOT) as reality: _ = reality.reset() _ = reality.step( "go_to_relative", { "xyt_position": [0, 0, (10 / 180) * np.pi], "use_map": False, "close_loop": True, "smooth": False, }, )
def __init__(self, config: Config) -> None: """Constructor :param config: config for the environment. Should contain id for simulator and ``task_name`` which are passed into ``make_sim`` and ``make_task``. :param dataset: reference to dataset for task instance level information. Can be defined as :py:`None` in which case ``_episodes`` should be populated from outside. """ assert config.is_frozen(), ("Freeze the config before creating the " "environment, use config.freeze().") self._config = config self._current_episode_index = None self._current_episode = None self._scenes = config.DATASET.CONTENT_SCENES self._swap_building_every = config.ENVIRONMENT.ITERATOR_OPTIONS.MAX_SCENE_REPEAT_EPISODES self._current_scene_episode_idx = 0 self._current_scene_idx = 0 self._config.defrost() self._config.SIMULATOR.SCENE = os.path.join( config.DATASET.SCENES_DIR, 'mp3d/{}/{}.glb'.format(self._scenes, self._scenes)) self._config.freeze() self._sim = make_sim(id_sim=self._config.SIMULATOR.TYPE, config=self._config.SIMULATOR)
def __init__(self, config: Config, dataset: Optional[Dataset] = None) -> None: assert config.is_frozen(), ("Freeze the config before creating the " "environment, use config.freeze()") self._config = config self._dataset = dataset self._current_episode_index = None if self._dataset is None and config.DATASET.TYPE: self._dataset = make_dataset(id_dataset=config.DATASET.TYPE, config=config.DATASET) self._episodes = self._dataset.episodes if self._dataset else [] self._sim = make_sim(id_sim=self._config.SIMULATOR.TYPE, config=self._config.SIMULATOR) self._task = make_task( self._config.TASK.TYPE, task_config=self._config.TASK, sim=self._sim, dataset=dataset, ) self.observation_space = SpaceDict({ **self._sim.sensor_suite.observation_spaces.spaces, **self._task.sensor_suite.observation_spaces.spaces, }) self.action_space = self._sim.action_space self._max_episode_seconds = ( self._config.ENVIRONMENT.MAX_EPISODE_SECONDS) self._max_episode_steps = self._config.ENVIRONMENT.MAX_EPISODE_STEPS self._elapsed_steps = 0 self._episode_start_time: Optional[float] = None self._episode_over = False
def test_sim_no_sensors(): config = get_config() config.defrost() config.SIMULATOR.AGENT_0.SENSORS = [] if not os.path.exists(config.SIMULATOR.SCENE): pytest.skip("Please download Habitat test data to data folder.") sim = make_sim(config.SIMULATOR.TYPE, config=config.SIMULATOR) sim.reset() sim.close()
def __init__(self, config: Config, dataset: Optional[Dataset] = None) -> None: """Constructor :param config: config for the environment. Should contain id for simulator and ``task_name`` which are passed into ``make_sim`` and ``make_task``. :param dataset: reference to dataset for task instance level information. Can be defined as :py:`None` in which case ``_episodes`` should be populated from outside. """ assert config.is_frozen(), ("Freeze the config before creating the " "environment, use config.freeze().") self._config = config self._dataset = dataset self._current_episode_index = None if self._dataset is None and config.DATASET.TYPE: self._dataset = make_dataset(id_dataset=config.DATASET.TYPE, config=config.DATASET) self._episodes = self._dataset.episodes if self._dataset else [] self._current_episode = None iter_option_dict = { k.lower(): v for k, v in config.ENVIRONMENT.ITERATOR_OPTIONS.items() } self._episode_iterator = self._dataset.get_episode_iterator( **iter_option_dict) # load the first scene if dataset is present if self._dataset: assert (len(self._dataset.episodes) > 0), "dataset should have non-empty episodes list" self._config.defrost() self._config.SIMULATOR.SCENE = self._dataset.episodes[0].scene_id self._config.freeze() self._sim = make_sim(id_sim=self._config.SIMULATOR.TYPE, config=self._config.SIMULATOR) self._task = make_task( self._config.TASK.TYPE, config=self._config.TASK, sim=self._sim, dataset=self._dataset, ) self.observation_space = SpaceDict({ **self._sim.sensor_suite.observation_spaces.spaces, **self._task.sensor_suite.observation_spaces.spaces, }) self.action_space = self._task.action_space self._max_episode_seconds = ( self._config.ENVIRONMENT.MAX_EPISODE_SECONDS) self._max_episode_steps = self._config.ENVIRONMENT.MAX_EPISODE_STEPS self._elapsed_steps = 0 self._episode_start_time: Optional[float] = None self._episode_over = False
def __init__(self, forward_step, angle_step, is_blind=False, sensors=["RGB_SENSOR"]): config = habitat.get_config() log_mesg("env: forward_step: {}, angle_step: {}".format( forward_step, angle_step)) config.defrost() config.PYROBOT.SENSORS = sensors config.PYROBOT.RGB_SENSOR.WIDTH = 256 config.PYROBOT.RGB_SENSOR.HEIGHT = 256 config.PYROBOT.DEPTH_SENSOR.WIDTH = 256 config.PYROBOT.DEPTH_SENSOR.HEIGHT = 256 config.PYROBOT.DEPTH_SENSOR.MAX_DEPTH = 10 config.PYROBOT.DEPTH_SENSOR.MIN_DEPTH = 0.3 config.freeze() self._reality = make_sim(id_sim="PyRobot-v0", config=config.PYROBOT) self._angle = (angle_step / 180) * np.pi self._pointgoal_key = "pointgoal_with_gps_compass" self.is_blind = is_blind if not is_blind: sensors_dict = { **self._reality.sensor_suite.observation_spaces.spaces } else: sensors_dict = {} sensors_dict[self._pointgoal_key] = spaces.Box( low=np.finfo(np.float32).min, high=np.finfo(np.float32).max, shape=(2, ), dtype=np.float32, ) self.observation_space = SpaceDict(sensors_dict) self.action_space = spaces.Discrete(4) self._actions = { "forward": [forward_step, 0, 0], "left": [0, 0, self._angle], "right": [0, 0, -self._angle], "stop": [0, 0, 0], } self._process = None self._last_time = -1
def __init__(self, config: Config, dataset: Optional[Dataset] = None) -> None: assert config.is_frozen(), ("Freeze the config before creating the " "environment, use config.freeze().") self._config = config self._dataset = dataset self._current_episode_index = None if self._dataset is None and config.DATASET.TYPE: self._dataset = make_dataset(id_dataset=config.DATASET.TYPE, config=config.DATASET) self._episodes = self._dataset.episodes if self._dataset else [] self._current_episode = None iter_option_dict = { k.lower(): v for k, v in config.ENVIRONMENT.ITERATOR_OPTIONS.items() } self._episode_iterator = self._dataset.get_episode_iterator( **iter_option_dict) # load the first scene if dataset is present if self._dataset: assert (len(self._dataset.episodes) > 0), "dataset should have non-empty episodes list" self._config.defrost() self._config.SIMULATOR.SCENE = self._dataset.episodes[0].scene_id self._config.freeze() self._sim = make_sim(id_sim=self._config.SIMULATOR.TYPE, config=self._config.SIMULATOR) self._task = make_task( self._config.TASK.TYPE, task_config=self._config.TASK, sim=self._sim, dataset=self._dataset, ) self.observation_space = SpaceDict({ **self._sim.sensor_suite.observation_spaces.spaces, **self._task.sensor_suite.observation_spaces.spaces, }) self.action_space = self._sim.action_space self._max_episode_seconds = ( self._config.ENVIRONMENT.MAX_EPISODE_SECONDS) self._max_episode_steps = self._config.ENVIRONMENT.MAX_EPISODE_STEPS self._elapsed_steps = 0 self._episode_start_time: Optional[float] = None self._episode_over = False
def generate_pointnav_dataset(config, num_episodes): sim = make_sim(id_sim=config.SIMULATOR.TYPE, config=config.SIMULATOR) sim.seed(config.SEED) random.seed(config.SEED) generator = pointnav_generator.generate_pointnav_episode( sim=sim, shortest_path_success_distance=config.TASK.SUCCESS_DISTANCE, shortest_path_max_steps=config.ENVIRONMENT.MAX_EPISODE_STEPS, ) episodes = [] for i in range(num_episodes): print(f"Generating episode {i+1}/{num_episodes}") episode = next(generator) episodes.append(episode) dataset = habitat.Dataset() dataset.episodes = episodes return dataset
def test_sim_geodesic_distance(): config = get_config() if not os.path.exists(config.SIMULATOR.SCENE): pytest.skip("Please download Habitat test data to data folder.") sim = make_sim(config.SIMULATOR.TYPE, config=config.SIMULATOR) sim.seed(0) sim.reset() start_point = sim.sample_navigable_point() navigable_points = [sim.sample_navigable_point() for _ in range(10)] assert np.isclose( sim.geodesic_distance(start_point, navigable_points[0]), 1.3849650 ), "Geodesic distance or sample navigable points mechanism has been changed." assert np.isclose( sim.geodesic_distance(start_point, navigable_points), 0.6194838 ), "Geodesic distance or sample navigable points mechanism has been changed." assert sim.geodesic_distance(start_point, navigable_points) == np.min( [ sim.geodesic_distance(start_point, position) for position in navigable_points ] ), "Geodesic distance for multi target setup isn't equal to separate single target calls." sim.close()
def test_sim_geodesic_distance(): config = get_config() if not os.path.exists(config.SIMULATOR.SCENE): pytest.skip("Please download Habitat test data to data folder.") sim = make_sim(config.SIMULATOR.TYPE, config=config.SIMULATOR) sim.reset() with open( os.path.join( os.path.dirname(__file__), "data", "test-sim-geodesic-distance-test-golden.json", ), "r", ) as f: test_data = json.load(f) for test_case in test_data["single_end"]: assert np.isclose( sim.geodesic_distance(test_case["start"], test_case["end"]), test_case["expected"], ), "Geodesic distance mechanism has been changed" for test_case in test_data["multi_end"]: assert np.isclose( sim.geodesic_distance(test_case["start"], test_case["ends"]), test_case["expected"], ), "Geodesic distance mechanism has been changed" assert np.isclose( sim.geodesic_distance(test_case["start"], test_case["ends"]), np.min([ sim.geodesic_distance(test_case["start"], end) for end in test_case["ends"] ]), ), "Geodesic distance for multi target setup isn't equal to separate single target calls." sim.close()
def init_sim(): config = get_config() if not os.path.exists(config.SIMULATOR.SCENE): pytest.skip("Please download Habitat test data to data folder.") return make_sim(config.SIMULATOR.TYPE, config=config.SIMULATOR)
def __init__(self, config: Config) -> None: """Constructor :param config: config for the environment. Should contain id for simulator and ``task_name`` which are passed into ``make_sim`` and ``make_task``. :param dataset: reference to dataset for task instance level information. Can be defined as :py:`None` in which case ``_episodes`` should be populated from outside. """ assert config.is_frozen(), ("Freeze the config before creating the " "environment, use config.freeze().") self._config = config self._current_episode_index = None self._current_episode = None iter_option_dict = { k.lower(): v for k, v in config.ENVIRONMENT.ITERATOR_OPTIONS.items() } self._scenes = config.DATASET.CONTENT_SCENES self._swap_building_every = config.ENVIRONMENT.ITERATOR_OPTIONS.MAX_SCENE_REPEAT_EPISODES print('[HabitatEnv] Total {} scenes : '.format(len(self._scenes)), self._scenes) print('[HabitatEnv] swap building every', self._swap_building_every) self._current_scene_episode_idx = 0 self._current_scene_idx = 0 self._config.defrost() if 'mp3d' in config.DATASET.DATA_PATH: self._config.SIMULATOR.SCENE = os.path.join( config.DATASET.SCENES_DIR, 'mp3d/{}/{}.glb'.format(self._scenes[0], self._scenes[0])) else: self._config.SIMULATOR.SCENE = os.path.join( config.DATASET.SCENES_DIR, 'gibson_habitat/{}.glb'.format(self._scenes[0])) if not os.path.exists(self._config.SIMULATOR.SCENE): self._config.SIMULATOR.SCENE = os.path.join( self._config.DATASET.SCENES_DIR, 'gibson_more/{}.glb'.format(self._scenes[0])) self._config.freeze() self._sim = make_sim(id_sim=self._config.SIMULATOR.TYPE, config=self._config.SIMULATOR) self._task = make_task(self._config.TASK.TYPE, config=self._config.TASK, sim=self._sim) self.observation_space = SpaceDict({ **self._sim.sensor_suite.observation_spaces.spaces, **self._task.sensor_suite.observation_spaces.spaces, }) self.action_space = self._task.action_space self._max_episode_seconds = ( self._config.ENVIRONMENT.MAX_EPISODE_SECONDS) self._max_episode_steps = self._config.ENVIRONMENT.MAX_EPISODE_STEPS self._elapsed_steps = 0 self._episode_start_time: Optional[float] = None self._episode_over = False #TODO listup demonstration data self._episode_dataset = {}
def __init__(self, scene, level, housetype='mp3d'): # -- scene = data/mp3d/house/house.glb self.scene = scene self.level = level # -- int self.house = scene.split('/')[-2] self.housetype = housetype #-- setup config self.config = get_config() self.config.defrost() self.config.SIMULATOR.SCENE = scene self.config.SIMULATOR.AGENT_0.SENSORS = [ "RGB_SENSOR", "DEPTH_SENSOR", "SEMANTIC_SENSOR" ] # -- Original resolution self.config.SIMULATOR.FORWARD_STEP_SIZE = 0.1 self.config.SIMULATOR.TURN_ANGLE = 9 # -- fine resolution setps #self.config.SIMULATOR.FORWARD_STEP_SIZE = 0.05 #self.config.SIMULATOR.TURN_ANGLE = 3 # -- render High Rez images #self.config.SIMULATOR.RGB_SENSOR.HEIGHT = 720 #self.config.SIMULATOR.RGB_SENSOR.WIDTH = 1280 # -- LOOK DOWN #theta = 30 * np.pi / 180 #self.config.SIMULATOR.RGB_SENSOR.ORIENTATION = [-theta, 0.0, 0.0] #self.config.SIMULATOR.DEPTH_SENSOR.ORIENTATION = [-theta, 0.0, 0.0] #self.config.SIMULATOR.SEMANTIC_SENSOR.ORIENTATION = [-theta, 0.0, 0.0] # -- OUTDATED (might be able to re-instantiate those in future commits) #self.config.TASK.POSSIBLE_ACTIONS = ["STOP", "MOVE_FORWARD", # "TURN_LEFT", "TURN_RIGHT", # "LOOK_UP", "LOOK_DOWN"] # -- ObjNav settings #self.config.SIMULATOR.FORWARD_STEP_SIZE = 0.25 #self.config.SIMULATOR.TURN_ANGLE = 30 self.config.freeze() self.agent_height = self.config.SIMULATOR.AGENT_0.HEIGHT self.sim = make_sim(id_sim=self.config.SIMULATOR.TYPE, config=self.config.SIMULATOR) self.semantic_annotations = self.sim.semantic_annotations() self.sim.reset() agent_state = self.get_agent_state() self.position = agent_state.position self.rotation = agent_state.rotation # -- get level dimensions # -- read it directly from the saved data from the .house files # Tries to set the agent on the given floor. It's actually quite hard.. if housetype == 'mp3d': env = '_'.join([self.house, str(self.level)]) houses_dim = json.load(open('data/houses_dim.json', 'r')) self.center = np.array(houses_dim[env]['center']) self.sizes = np.array(houses_dim[env]['sizes']) self.start_height = self.center[1] - self.sizes[1] / 2 self.set_agent_on_level() else: pass self.all_objects = self.get_objects_in_house()
def _eval(self): start_time = time.time() if self.config.MANUAL_COMMANDS: init_time = None manual_step_start_time = None total_manual_time = 0.0 checkpoint_index = int( (re.findall('\d+', self.config.EVAL_CKPT_PATH_DIR))[-1]) ckpt_dict = torch.load(self.config.EVAL_CKPT_PATH_DIR, map_location="cpu") print( f'Number of steps of the ckpt: {ckpt_dict["extra_state"]["step"]}') config = self._setup_config(ckpt_dict) ppo_cfg = config.RL.PPO ans_cfg = config.RL.ANS self.mapper_rollouts = None self._setup_actor_critic_agent(ppo_cfg, ans_cfg) self.mapper_agent.load_state_dict(ckpt_dict["mapper_state_dict"]) if self.local_agent is not None: self.local_agent.load_state_dict(ckpt_dict["local_state_dict"]) self.local_actor_critic = self.local_agent.actor_critic else: self.local_actor_critic = self.ans_net.local_policy self.global_agent.load_state_dict(ckpt_dict["global_state_dict"]) self.mapper = self.mapper_agent.mapper self.global_actor_critic = self.global_agent.actor_critic # Set models to evaluation self.mapper.eval() self.local_actor_critic.eval() self.global_actor_critic.eval() M = ans_cfg.overall_map_size V = ans_cfg.MAPPER.map_size s = ans_cfg.MAPPER.map_scale imH, imW = ans_cfg.image_scale_hw num_steps = self.config.T_EXP prev_action = torch.zeros(1, 1, device=self.device, dtype=torch.long) masks = torch.zeros(1, 1, device=self.device) try: self.sim = make_sim('PyRobot-v1', config=self.config.TASK_CONFIG.PYROBOT) except (KeyboardInterrupt, SystemExit): sys.exit() pose = defaultdict() self.sim._robot.camera.set_tilt(math.radians(self.config.CAMERA_TILT), wait=True) print( f"\nStarting Camera State: {self.sim.get_agent_state()['camera']}") print(f"Starting Agent State: {self.sim.get_agent_state()['base']}") obs = [self.sim.reset()] if self.config.SAVE_OBS_IMGS: cv2.imwrite(f'obs/depth_dirty_s.jpg', obs[0]['depth'] * 255.0) obs[0]['depth'][..., 0] = self._correct_depth(obs, -1) if self.config.SAVE_OBS_IMGS: cv2.imwrite(f'obs/rgb_s.jpg', obs[0]['rgb'][:, :, ::-1]) cv2.imwrite(f'depth_s.jpg', obs[0]['depth'] * 255.0) starting_agent_state = self.sim.get_agent_state() locobot2relative = CoordProjection(starting_agent_state['base']) pose['base'] = locobot2relative(starting_agent_state['base']) print(f"Starting Agent Pose: {pose['base']}\n") batch = self._prepare_batch(obs, -1, device=self.device) if ans_cfg.MAPPER.use_sensor_positioning: batch['pose'] = pose['base'].to(self.device) batch['pose'][0][1:] = -batch['pose'][0][1:] prev_batch = batch num_envs = self.config.NUM_PROCESSES agent_poses_over_time = [] for i in range(num_envs): agent_poses_over_time.append( torch.tensor([(M - 1) / 2, (M - 1) / 2, 0])) state_estimates = { "pose_estimates": torch.zeros(num_envs, 3).to(self.device), "map_states": torch.zeros(num_envs, 2, M, M).to(self.device), "recurrent_hidden_states": torch.zeros(1, num_envs, ans_cfg.LOCAL_POLICY.hidden_size).to(self.device), "visited_states": torch.zeros(num_envs, 1, M, M).to(self.device), } ground_truth_states = { "visible_occupancy": torch.zeros(num_envs, 2, M, M).to(self.device), "pose": torch.zeros(num_envs, 3).to(self.device), "environment_layout": torch.zeros(num_envs, 2, M, M).to(self.device) } # Reset ANS states self.ans_net.reset() # Frames for video creation rgb_frames = [] if len(self.config.VIDEO_OPTION) > 0: os.makedirs(self.config.VIDEO_DIR, exist_ok=True) step_start_time = time.time() for i in range(num_steps): print( f"\n\n---------------------------------------------------<<< STEP {i} >>>---------------------------------------------------" ) ep_time = torch.zeros(num_envs, 1, device=self.device).fill_(i) ( mapper_inputs, local_policy_inputs, global_policy_inputs, mapper_outputs, local_policy_outputs, global_policy_outputs, state_estimates, intrinsic_rewards, ) = self.ans_net.act( batch, prev_batch, state_estimates, ep_time, masks, deterministic=True, ) if self.config.SAVE_MAP_IMGS: cv2.imwrite( f'maps/test_map_{i - 1}.jpg', self._round_map(state_estimates['map_states']) * 255) action = local_policy_outputs["actions"][0][0] distance2ggoal = torch.norm( mapper_outputs['curr_map_position'] - self.ans_net.states["curr_global_goals"], dim=1) * s print(f"Distance to Global Goal: {distance2ggoal}") reached_flag = distance2ggoal < ans_cfg.goal_success_radius if self.config.MANUAL_COMMANDS: if init_time is None: init_time = time.time() - start_time total_manual_time = total_manual_time + init_time if manual_step_start_time is not None: manual_step_time = time.time() - manual_step_start_time total_manual_time = total_manual_time + manual_step_time action = torch.tensor( int(input('Waiting input to start new action: '))) manual_step_start_time = time.time() if action.item() == 3: reached_flag = True prev_action.copy_(action) if not reached_flag and action.item() != 3: print(f'Doing Env Step [{self.ACT_2_NAME[action.item()]}]...') action_command = self.ACT_2_COMMAND[action.item()] obs = self._do_action(action_command) if self.config.SAVE_OBS_IMGS: cv2.imwrite(f'obs/depth_dirty_{i}.jpg', obs[0]['depth'] * 255.0) # Correcting invalid depth pixels obs[0]['depth'][..., 0] = self._correct_depth(obs, i) if self.config.SAVE_OBS_IMGS: cv2.imwrite(f'obs/rgb_{i}.jpg', obs[0]['rgb'][:, :, ::-1]) cv2.imwrite(f'obs/depth_{i}.jpg', obs[0]['depth'] * 255.0) agent_state = self.sim.get_agent_state() prev_batch = batch batch = self._prepare_batch(obs, i, device=self.device) pose = defaultdict() pose['base'] = locobot2relative(agent_state['base']) if ans_cfg.MAPPER.use_sensor_positioning: batch['pose'] = pose['base'].to(self.device) batch['pose'][0][1:] = -batch['pose'][0][1:] map_coords = convert_world2map( batch['pose'], (M, M), ans_cfg.OCCUPANCY_ANTICIPATOR. EGO_PROJECTION.map_scale).squeeze() map_coords = torch.cat( (map_coords, batch['pose'][0][-1].reshape(1))) if self.config.COORD_DEBUG: print('COORDINATES CHECK') print( f'Starting Agent State: {starting_agent_state["base"]}' ) print(f'Current Agent State: {agent_state["base"]}') print( f'Current Sim Agent State: {self.sim.get_agent_state()["base"]}' ) print(f'Current Global Coords: {batch["pose"]}') print(f'Current Map Coords: {map_coords}') agent_poses_over_time.append(map_coords) step_time = time.time() - step_start_time print(f"\nStep Time: {step_time}") step_start_time = time.time() # Create new frame of the video if (len(self.config.VIDEO_OPTION) > 0): frame = observations_to_image( obs[0], observation_size=300, collision_flag=self.config.DRAW_COLLISIONS) # Add ego_map_gt to frame ego_map_gt_i = asnumpy(batch["ego_map_gt"][0]) # (2, H, W) ego_map_gt_i = convert_gt2channel_to_gtrgb(ego_map_gt_i) ego_map_gt_i = cv2.resize(ego_map_gt_i, (300, 300)) # frame = np.concatenate([frame], axis=1) # Generate ANS specific visualizations environment_layout = asnumpy( ground_truth_states["environment_layout"][0]) # (2, H, W) visible_occupancy = mapper_outputs["gt_mt"][0].cpu().numpy( ) # (2, H, W) anticipated_occupancy = mapper_outputs["hat_mt"][0].cpu( ).numpy() # (2, H, W) H = frame.shape[0] visible_occupancy_vis = generate_topdown_allocentric_map( environment_layout, visible_occupancy, agent_poses_over_time, thresh_explored=ans_cfg.thresh_explored, thresh_obstacle=ans_cfg.thresh_obstacle, zoom=False) visible_occupancy_vis = cv2.resize(visible_occupancy_vis, (H, H)) anticipated_occupancy_vis = generate_topdown_allocentric_map( environment_layout, anticipated_occupancy, agent_poses_over_time, thresh_explored=ans_cfg.thresh_explored, thresh_obstacle=ans_cfg.thresh_obstacle, zoom=False) anticipated_occupancy_vis = cv2.resize( anticipated_occupancy_vis, (H, H)) anticipated_action_map = generate_topdown_allocentric_map( environment_layout, anticipated_occupancy, agent_poses_over_time, zoom=False, thresh_explored=ans_cfg.thresh_explored, thresh_obstacle=ans_cfg.thresh_obstacle, ) global_goals = self.ans_net.states["curr_global_goals"] local_goals = self.ans_net.states["curr_local_goals"] if global_goals is not None: cX = int(global_goals[0, 0].item()) cY = int(global_goals[0, 1].item()) anticipated_action_map = cv2.circle( anticipated_action_map, (cX, cY), 10, (255, 0, 0), -1, ) if local_goals is not None: cX = int(local_goals[0, 0].item()) cY = int(local_goals[0, 1].item()) anticipated_action_map = cv2.circle( anticipated_action_map, (cX, cY), 10, (0, 255, 255), -1, ) anticipated_action_map = cv2.resize(anticipated_action_map, (H, H)) maps_vis = np.concatenate( [ visible_occupancy_vis, anticipated_occupancy_vis, anticipated_action_map, ego_map_gt_i ], axis=1, ) if self.config.RL.ANS.overall_map_size == 2001 or self.config.RL.ANS.overall_map_size == 961: if frame.shape[1] < maps_vis.shape[1]: diff = maps_vis.shape[1] - frame.shape[1] npad = ((0, 0), (diff // 2, diff // 2), (0, 0)) frame = np.pad(frame, pad_width=npad, mode='constant', constant_values=0) elif frame.shape[1] > maps_vis.shape[1]: diff = frame.shape[1] - maps_vis.shape[1] npad = ((0, 0), (diff // 2, diff // 2), (0, 0)) maps_vis = np.pad(maps_vis, pad_width=npad, mode='constant', constant_values=0) frame = np.concatenate([frame, maps_vis], axis=0) rgb_frames.append(frame) if self.config.SAVE_VIDEO_IMGS: try: os.mkdir("fig1") except: pass print("Saved imgs for Fig. 1!") cv2.imwrite(f'fig1/rgb_{step_start_time}.jpg', obs[0]['rgb'][:, :, ::-1]) cv2.imwrite(f'fig1/depth_{step_start_time}.jpg', obs[0]['depth'] * 255.0) cv2.imwrite(f'fig1/aap_{step_start_time}.jpg', anticipated_action_map[..., ::-1]) cv2.imwrite(f'fig1/em_{step_start_time}.jpg', ego_map_gt_i[..., ::-1]) if self.config.DEBUG_VIDEO_FRAME: cv2.imwrite('last_frame.jpg', frame) if reached_flag: for f in range(20): rgb_frames.append(frame) # Video creation video_dict = {"t": start_time} if (i + 1) % 10 == 0 or reached_flag: generate_video( video_option=self.config.VIDEO_OPTION, video_dir=self.config.VIDEO_DIR, images=rgb_frames, episode_id=0, checkpoint_idx=checkpoint_index, metrics=video_dict, tb_writer=TensorboardWriter('tb/locobot'), ) if reached_flag: if self.config.MANUAL_COMMANDS: manual_step_time = time.time() - manual_step_start_time total_manual_time = total_manual_time + manual_step_time print(f"Manual elapsed time: {total_manual_time}") print(f"Number of steps: {i + 1}") print(f"Elapsed time: {time.time() - start_time}") print(f"Final Distance to Goal: {distance2ggoal}") if "bump" in obs[0]: print(f"Collision: {obs[0]['bump']}") print("Exiting...") break return
from habitat import get_config from habitat.sims import make_sim from scipy.spatial.transform import Rotation as R from core import _transform3D house = '17DRP5sb8fy' scene = '../data/mp3d/{}/{}.glb'.format(house, house) config = get_config() config.defrost() config.SIMULATOR.SCENE = scene config.SIMULATOR.AGENT_0.SENSORS = ["DEPTH_SENSOR"] config.freeze() sim = make_sim(id_sim=config.SIMULATOR.TYPE, config=config.SIMULATOR) sim.reset() vfov = 67.5 world_shift = torch.FloatTensor([0, 0, 0]) projector = PointCloud( vfov, 1, 480, 640, world_shift, 0.5, device=torch.device("cpu"), )