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"),
)
