def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model-path", type=str, required=True)
parser.add_argument("--sim-gpu-id", type=int, required=True)
parser.add_argument("--pth-gpu-id", type=int, required=True)
parser.add_argument("--num-processes", type=int, required=True)
parser.add_argument("--hidden-size", type=int, default=512)
parser.add_argument("--count-test-episodes", type=int, default=100)
parser.add_argument(
"--sensors",
type=str,
default="RGB_SENSOR,DEPTH_SENSOR",
help="comma separated string containing different"
"sensors to use, currently 'RGB_SENSOR' and"
"'DEPTH_SENSOR' are supported",
)
parser.add_argument(
"--task-config",
type=str,
default="configs/tasks/pointnav.yaml",
help="path to config yaml containing information about task",
)
args = parser.parse_args()
device = torch.device("cuda:{}".format(args.pth_gpu_id))
env_configs = []
baseline_configs = []
for _ in range(args.num_processes):
config_env = get_config(config_paths=args.task_config)
config_env.defrost()
config_env.DATASET.SPLIT = "val"
agent_sensors = args.sensors.strip().split(",")
for sensor in agent_sensors:
assert sensor in ["RGB_SENSOR", "DEPTH_SENSOR"]
config_env.SIMULATOR.AGENT_0.SENSORS = agent_sensors
config_env.freeze()
env_configs.append(config_env)
config_baseline = cfg_baseline()
baseline_configs.append(config_baseline)
assert len(baseline_configs) > 0, "empty list of datasets"
envs = habitat.VectorEnv(
make_env_fn=make_env_fn,
env_fn_args=tuple(
tuple(zip(env_configs, baseline_configs,
range(args.num_processes)))),
)
ckpt = torch.load(args.model_path, map_location=device)
actor_critic = Policy(
observation_space=envs.observation_spaces[0],
action_space=envs.action_spaces[0],
hidden_size=512,
goal_sensor_uuid=env_configs[0].TASK.GOAL_SENSOR_UUID,
)
actor_critic.to(device)
ppo = PPO(
actor_critic=actor_critic,
clip_param=0.1,
ppo_epoch=4,
num_mini_batch=32,
value_loss_coef=0.5,
entropy_coef=0.01,
lr=2.5e-4,
eps=1e-5,
max_grad_norm=0.5,
)
ppo.load_state_dict(ckpt["state_dict"])
actor_critic = ppo.actor_critic
observations = envs.reset()
batch = batch_obs(observations)
for sensor in batch:
batch[sensor] = batch[sensor].to(device)
episode_rewards = torch.zeros(envs.num_envs, 1, device=device)
episode_spls = torch.zeros(envs.num_envs, 1, device=device)
episode_success = torch.zeros(envs.num_envs, 1, device=device)
episode_counts = torch.zeros(envs.num_envs, 1, device=device)
current_episode_reward = torch.zeros(envs.num_envs, 1, device=device)
test_recurrent_hidden_states = torch.zeros(args.num_processes,
args.hidden_size,
device=device)
not_done_masks = torch.zeros(args.num_processes, 1, device=device)
while episode_counts.sum() < args.count_test_episodes:
with torch.no_grad():
_, actions, _, test_recurrent_hidden_states = actor_critic.act(
batch,
test_recurrent_hidden_states,
not_done_masks,
deterministic=False,
)
outputs = envs.step([a[0].item() for a in actions])
observations, rewards, dones, infos = [list(x) for x in zip(*outputs)]
batch = batch_obs(observations)
for sensor in batch:
batch[sensor] = batch[sensor].to(device)
not_done_masks = torch.tensor(
[[0.0] if done else [1.0] for done in dones],
dtype=torch.float,
device=device,
)
for i in range(not_done_masks.shape[0]):
if not_done_masks[i].item() == 0:
episode_spls[i] += infos[i]["roomnavmetric"]
if infos[i]["roomnavmetric"] > 0:
episode_success[i] += 1
rewards = torch.tensor(rewards, dtype=torch.float,
device=device).unsqueeze(1)
current_episode_reward += rewards
episode_rewards += (1 - not_done_masks) * current_episode_reward
episode_counts += 1 - not_done_masks
current_episode_reward *= not_done_masks
episode_reward_mean = (episode_rewards / episode_counts).mean().item()
episode_spl_mean = (episode_spls / episode_counts).mean().item()
episode_success_mean = (episode_success / episode_counts).mean().item()
print("Average episode reward: {:.6f}".format(episode_reward_mean))
print("Average episode success: {:.6f}".format(episode_success_mean))
print("Average episode spl: {:.6f}".format(episode_spl_mean))
def main():
parser = ppo_args()
args = parser.parse_args()
random.seed(args.seed)
device = torch.device("cuda:{}".format(args.pth_gpu_id))
logger.add_filehandler(args.log_file)
if not os.path.isdir(args.checkpoint_folder):
os.makedirs(args.checkpoint_folder)
for p in sorted(list(vars(args))):
logger.info("{}: {}".format(p, getattr(args, p)))
envs = construct_envs(args)
actor_critic = Policy(
observation_space=envs.observation_spaces[0],
action_space=envs.action_spaces[0],
hidden_size=args.hidden_size,
)
actor_critic.to(device)
agent = PPO(
actor_critic,
args.clip_param,
args.ppo_epoch,
args.num_mini_batch,
args.value_loss_coef,
args.entropy_coef,
lr=args.lr,
eps=args.eps,
max_grad_norm=args.max_grad_norm,
)
logger.info("agent number of parameters: {}".format(
sum(param.numel() for param in agent.parameters())))
observations = envs.reset()
batch = batch_obs(observations)
rollouts = RolloutStorage(
args.num_steps,
envs.num_envs,
envs.observation_spaces[0],
envs.action_spaces[0],
args.hidden_size,
)
for sensor in rollouts.observations:
rollouts.observations[sensor][0].copy_(batch[sensor])
rollouts.to(device)
episode_rewards = torch.zeros(envs.num_envs, 1)
episode_counts = torch.zeros(envs.num_envs, 1)
current_episode_reward = torch.zeros(envs.num_envs, 1)
window_episode_reward = deque()
window_episode_counts = deque()
t_start = time()
env_time = 0
pth_time = 0
count_steps = 0
count_checkpoints = 0
for update in range(args.num_updates):
if args.use_linear_lr_decay:
update_linear_schedule(agent.optimizer, update, args.num_updates,
args.lr)
agent.clip_param = args.clip_param * (1 - update / args.num_updates)
for step in range(args.num_steps):
t_sample_action = time()
# sample actions
with torch.no_grad():
step_observation = {
k: v[step]
for k, v in rollouts.observations.items()
}
(
values,
actions,
actions_log_probs,
recurrent_hidden_states,
) = actor_critic.act(
step_observation,
rollouts.recurrent_hidden_states[step],
rollouts.masks[step],
)
pth_time += time() - t_sample_action
t_step_env = time()
outputs = envs.step([a[0].item() for a in actions])
observations, rewards, dones, infos = [
list(x) for x in zip(*outputs)
]
env_time += time() - t_step_env
t_update_stats = time()
batch = batch_obs(observations)
rewards = torch.tensor(rewards, dtype=torch.float)
rewards = rewards.unsqueeze(1)
masks = torch.tensor([[0.0] if done else [1.0] for done in dones],
dtype=torch.float)
current_episode_reward += rewards
episode_rewards += (1 - masks) * current_episode_reward
episode_counts += 1 - masks
current_episode_reward *= masks
rollouts.insert(
batch,
recurrent_hidden_states,
actions,
actions_log_probs,
values,
rewards,
masks,
)
count_steps += envs.num_envs
pth_time += time() - t_update_stats
if len(window_episode_reward) == args.reward_window_size:
window_episode_reward.popleft()
window_episode_counts.popleft()
window_episode_reward.append(episode_rewards.clone())
window_episode_counts.append(episode_counts.clone())
t_update_model = time()
with torch.no_grad():
last_observation = {
k: v[-1]
for k, v in rollouts.observations.items()
}
next_value = actor_critic.get_value(
last_observation,
rollouts.recurrent_hidden_states[-1],
rollouts.masks[-1],
).detach()
rollouts.compute_returns(next_value, args.use_gae, args.gamma,
args.tau)
value_loss, action_loss, dist_entropy = agent.update(rollouts)
rollouts.after_update()
pth_time += time() - t_update_model
# log stats
if update > 0 and update % args.log_interval == 0:
logger.info("update: {}\tfps: {:.3f}\t".format(
update, count_steps / (time() - t_start)))
logger.info("update: {}\tenv-time: {:.3f}s\tpth-time: {:.3f}s\t"
"frames: {}".format(update, env_time, pth_time,
count_steps))
window_rewards = (window_episode_reward[-1] -
window_episode_reward[0]).sum()
window_counts = (window_episode_counts[-1] -
window_episode_counts[0]).sum()
if window_counts > 0:
logger.info("Average window size {} reward: {:3f}".format(
len(window_episode_reward),
(window_rewards / window_counts).item(),
))
else:
logger.info("No episodes finish in current window")
# checkpoint model
if update % args.checkpoint_interval == 0:
checkpoint = {"state_dict": agent.state_dict()}
torch.save(
checkpoint,
os.path.join(
args.checkpoint_folder,
"ckpt.{}.pth".format(count_checkpoints),
),
)
count_checkpoints += 1
episode_rewards = torch.zeros(envs.num_envs, 1, device=device)
episode_spls = torch.zeros(envs.num_envs, 1, device=device)
episode_success = torch.zeros(envs.num_envs, 1, device=device)
episode_counts = torch.zeros(envs.num_envs, 1, device=device)
current_episode_reward = torch.zeros(envs.num_envs, 1, device=device)
test_recurrent_hidden_states = torch.zeros(
1, 512, device=device
)
not_done_masks = torch.zeros(1, 1, device=device)
while episode_counts.sum() < 100:
with torch.no_grad():
_, actions, _, test_recurrent_hidden_states = actor_critic.act(
batch,
test_recurrent_hidden_states,
not_done_masks,
deterministic=False,
)
outputs = envs.step([a[0].item() for a in actions])
observations, rewards, dones, infos = [list(x) for x in zip(*outputs)]
batch = batch_obs(observations)
for sensor in batch:
batch[sensor] = batch[sensor].to(device)
not_done_masks = torch.tensor(
[[0.0] if done else [1.0] for done in dones],
dtype=torch.float,
device=device,
)
def eval_checkpoint(checkpoint_path, args, writer, cur_ckpt_idx=0):
env_configs = []
baseline_configs = []
device = torch.device("cuda", args.pth_gpu_id)
for _ in range(args.num_processes):
config_env = get_config(config_paths=args.task_config)
config_env.defrost()
config_env.DATASET.SPLIT = "val"
agent_sensors = args.sensors.strip().split(",")
for sensor in agent_sensors:
assert sensor in ["RGB_SENSOR", "DEPTH_SENSOR"]
config_env.SIMULATOR.AGENT_0.SENSORS = agent_sensors
if args.video_option:
config_env.TASK.MEASUREMENTS.append("TOP_DOWN_MAP")
config_env.TASK.MEASUREMENTS.append("COLLISIONS")
config_env.freeze()
env_configs.append(config_env)
config_baseline = cfg_baseline()
baseline_configs.append(config_baseline)
assert len(baseline_configs) > 0, "empty list of datasets"
envs = habitat.VectorEnv(
make_env_fn=make_env_fn,
env_fn_args=tuple(
tuple(
zip(env_configs, baseline_configs, range(args.num_processes))
)
),
)
ckpt = torch.load(checkpoint_path, map_location=device)
actor_critic = Policy(
observation_space=envs.observation_spaces[0],
action_space=envs.action_spaces[0],
hidden_size=512,
goal_sensor_uuid=env_configs[0].TASK.GOAL_SENSOR_UUID,
)
actor_critic.to(device)
ppo = PPO(
actor_critic=actor_critic,
clip_param=0.1,
ppo_epoch=4,
num_mini_batch=32,
value_loss_coef=0.5,
entropy_coef=0.01,
lr=2.5e-4,
eps=1e-5,
max_grad_norm=0.5,
)
ppo.load_state_dict(ckpt["state_dict"])
actor_critic = ppo.actor_critic
observations = envs.reset()
batch = batch_obs(observations)
for sensor in batch:
batch[sensor] = batch[sensor].to(device)
episode_rewards = torch.zeros(envs.num_envs, 1, device=device)
episode_spls = torch.zeros(envs.num_envs, 1, device=device)
episode_success = torch.zeros(envs.num_envs, 1, device=device)
episode_counts = torch.zeros(envs.num_envs, 1, device=device)
current_episode_reward = torch.zeros(envs.num_envs, 1, device=device)
test_recurrent_hidden_states = torch.zeros(
args.num_processes, args.hidden_size, device=device
)
not_done_masks = torch.zeros(args.num_processes, 1, device=device)
stats_episodes = set()
rgb_frames = None
if args.video_option:
rgb_frames = [[]] * args.num_processes
os.makedirs(args.video_dir, exist_ok=True)
while episode_counts.sum() < args.count_test_episodes:
current_episodes = envs.current_episodes()
with torch.no_grad():
_, actions, _, test_recurrent_hidden_states = actor_critic.act(
batch,
test_recurrent_hidden_states,
not_done_masks,
deterministic=False,
)
outputs = envs.step([a[0].item() for a in actions])
observations, rewards, dones, infos = [list(x) for x in zip(*outputs)]
batch = batch_obs(observations)
for sensor in batch:
batch[sensor] = batch[sensor].to(device)
not_done_masks = torch.tensor(
[[0.0] if done else [1.0] for done in dones],
dtype=torch.float,
device=device,
)
for i in range(not_done_masks.shape[0]):
if not_done_masks[i].item() == 0:
episode_spls[i] += infos[i]["spl"]
if infos[i]["spl"] > 0:
episode_success[i] += 1
rewards = torch.tensor(
rewards, dtype=torch.float, device=device
).unsqueeze(1)
current_episode_reward += rewards
episode_rewards += (1 - not_done_masks) * current_episode_reward
episode_counts += 1 - not_done_masks
current_episode_reward *= not_done_masks
next_episodes = envs.current_episodes()
envs_to_pause = []
n_envs = envs.num_envs
for i in range(n_envs):
if next_episodes[i].episode_id in stats_episodes:
envs_to_pause.append(i)
# episode ended
if not_done_masks[i].item() == 0:
stats_episodes.add(current_episodes[i].episode_id)
if args.video_option:
generate_video(
args,
rgb_frames[i],
current_episodes[i].episode_id,
cur_ckpt_idx,
infos[i]["spl"],
writer,
)
rgb_frames[i] = []
# episode continues
elif args.video_option:
frame = observations_to_image(observations[i], infos[i])
rgb_frames[i].append(frame)
# stop tracking ended episodes if they exist
if len(envs_to_pause) > 0:
state_index = list(range(envs.num_envs))
for idx in reversed(envs_to_pause):
state_index.pop(idx)
envs.pause_at(idx)
# indexing along the batch dimensions
test_recurrent_hidden_states = test_recurrent_hidden_states[
:, state_index
]
not_done_masks = not_done_masks[state_index]
current_episode_reward = current_episode_reward[state_index]
for k, v in batch.items():
batch[k] = v[state_index]
if args.video_option:
rgb_frames = [rgb_frames[i] for i in state_index]
episode_reward_mean = (episode_rewards / episode_counts).mean().item()
episode_spl_mean = (episode_spls / episode_counts).mean().item()
episode_success_mean = (episode_success / episode_counts).mean().item()
logger.info("Average episode reward: {:.6f}".format(episode_reward_mean))
logger.info("Average episode success: {:.6f}".format(episode_success_mean))
logger.info("Average episode SPL: {:.6f}".format(episode_spl_mean))
writer.add_scalars(
"eval_reward", {"average reward": episode_reward_mean}, cur_ckpt_idx
)
writer.add_scalars(
"eval_SPL", {"average SPL": episode_spl_mean}, cur_ckpt_idx
)
writer.add_scalars(
"eval_success", {"average success": episode_success_mean}, cur_ckpt_idx
)
def eval_checkpoint(checkpoint_path, args, writer, cur_ckpt_idx=0):
env_configs = []
baseline_configs = []
device = torch.device("cuda", args.pth_gpu_id)
for _ in range(args.num_processes):
config_env = get_config(config_paths=args.task_config)
config_env.defrost()
config_env.DATASET.SPLIT = "val"
agent_sensors = args.sensors.strip().split(",")
for sensor in agent_sensors:
assert sensor in ["RGB_SENSOR", "DEPTH_SENSOR"]
config_env.SIMULATOR.AGENT_0.SENSORS = agent_sensors
if args.video_option:
config_env.TASK.MEASUREMENTS.append("TOP_DOWN_MAP")
config_env.TASK.MEASUREMENTS.append("COLLISIONS")
config_env.freeze()
env_configs.append(config_env)
config_baseline = cfg_baseline()
baseline_configs.append(config_baseline)
assert len(baseline_configs) > 0, "empty list of datasets"
envs = habitat.VectorEnv(
make_env_fn=make_env_fn,
env_fn_args=tuple(
tuple(zip(env_configs, baseline_configs,
range(args.num_processes)))),
)
ckpt = torch.load(checkpoint_path, map_location=device)
actor_critic = Policy(
observation_space=envs.observation_spaces[0],
action_space=envs.action_spaces[0],
hidden_size=512,
goal_sensor_uuid=env_configs[0].TASK.GOAL_SENSOR_UUID,
)
actor_critic.to(device)
ppo = PPO(
actor_critic=actor_critic,
clip_param=0.1,
ppo_epoch=4,
num_mini_batch=32,
value_loss_coef=0.5,
entropy_coef=0.01,
lr=2.5e-4,
eps=1e-5,
max_grad_norm=0.5,
)
ppo.load_state_dict(ckpt["state_dict"])
actor_critic = ppo.actor_critic
observations = envs.reset()
batch = batch_obs(observations)
for sensor in batch:
batch[sensor] = batch[sensor].to(device)
current_episode_reward = torch.zeros(envs.num_envs, 1, device=device)
test_recurrent_hidden_states = torch.zeros(args.num_processes,
args.hidden_size,
device=device)
not_done_masks = torch.zeros(args.num_processes, 1, device=device)
stats_episodes = dict() # dict of dicts that stores stats per episode
while episode_counts.sum() < args.count_test_episodes:
# test_recurrent_hidden_states_list.append(test_recurrent_hidden_states)
# pickle_out = open("hab_recurrent_states.pickle","wb")
# pickle.dump(test_recurrent_hidden_states_list, pickle_out)
# pickle_out.close()
# obs_list.append(observations[0])
# pickle_out = open("hab_obs_list.pickle","wb")
# pickle.dump(obs_list, pickle_out)
# pickle_out.close()
# mask_list.append(not_done_masks)
# pickle_out = open("hab_mask_list.pickle","wb")
# pickle.dump(mask_list, pickle_out)
# pickle_out.close()
with torch.no_grad():
_, actions, _, test_recurrent_hidden_states = actor_critic.act(
batch,
test_recurrent_hidden_states,
not_done_masks,
deterministic=True,
)
print("action_id is " + str(actions.item()))
print('point goal is ' + str(observations[0]['pointgoal']))
outputs = envs.step([a[0].item() for a in actions])
observations, rewards, dones, infos = [list(x) for x in zip(*outputs)]
#for visualizing where robot is going
#cv2.imshow("RGB", transform_rgb_bgr(observations[0]["rgb"]))
cv2.imshow("Depth", observations[0]["depth"])
cv2.waitKey(100)
time.sleep(0.2)
batch = batch_obs(observations)
for sensor in batch:
batch[sensor] = batch[sensor].to(device)
not_done_masks = torch.tensor(
[[0.0] if done else [1.0] for done in dones],
dtype=torch.float,
device=device,
)
rewards = torch.tensor(rewards, dtype=torch.float,
device=device).unsqueeze(1)
current_episode_reward += rewards
next_episodes = envs.current_episodes()
envs_to_pause = []
n_envs = envs.num_envs
for i in range(n_envs):
if (
next_episodes[i].scene_id,
next_episodes[i].episode_id,
) in stats_episodes:
envs_to_pause.append(i)
# episode ended
if not_done_masks[i].item() == 0:
episode_stats = dict()
episode_stats["spl"] = infos[i]["spl"]
episode_stats["success"] = int(infos[i]["spl"] > 0)
episode_stats["reward"] = current_episode_reward[i].item()
current_episode_reward[i] = 0
# use scene_id + episode_id as unique id for storing stats
stats_episodes[(
current_episodes[i].scene_id,
current_episodes[i].episode_id,
)] = episode_stats
if args.video_option:
generate_video(
args,
rgb_frames[i],
current_episodes[i].episode_id,
cur_ckpt_idx,
infos[i]["spl"],
writer,
)
rgb_frames[i] = []
# episode continues
elif args.video_option:
frame = observations_to_image(observations[i], infos[i])
rgb_frames[i].append(frame)
# pausing envs with no new episode
if len(envs_to_pause) > 0:
state_index = list(range(envs.num_envs))
for idx in reversed(envs_to_pause):
state_index.pop(idx)
envs.pause_at(idx)
# indexing along the batch dimensions
test_recurrent_hidden_states = test_recurrent_hidden_states[
state_index]
not_done_masks = not_done_masks[state_index]
current_episode_reward = current_episode_reward[state_index]
for k, v in batch.items():
batch[k] = v[state_index]
if args.video_option:
rgb_frames = [rgb_frames[i] for i in state_index]
aggregated_stats = dict()
for stat_key in next(iter(stats_episodes.values())).keys():
aggregated_stats[stat_key] = sum(
[v[stat_key] for v in stats_episodes.values()])
num_episodes = len(stats_episodes)
episode_reward_mean = aggregated_stats["reward"] / num_episodes
episode_spl_mean = aggregated_stats["spl"] / num_episodes
episode_success_mean = aggregated_stats["success"] / num_episodes
logger.info("Average episode reward: {:.6f}".format(episode_reward_mean))
logger.info("Average episode success: {:.6f}".format(episode_success_mean))
logger.info("Average episode SPL: {:.6f}".format(episode_spl_mean))
writer.add_scalars("eval_reward", {"average reward": episode_reward_mean},
cur_ckpt_idx)
writer.add_scalars("eval_SPL", {"average SPL": episode_spl_mean},
cur_ckpt_idx)
writer.add_scalars("eval_success",
{"average success": episode_success_mean}, cur_ckpt_idx)