def _eval_checkpoint(
self,
checkpoint_path: str,
writer: TensorboardWriter,
checkpoint_index: int = 0,
) -> None:
r"""Evaluates a single checkpoint.
Args:
checkpoint_path: path of checkpoint
writer: tensorboard writer object for logging to tensorboard
checkpoint_index: index of cur checkpoint for logging
Returns:
None
"""
config = self.config
config.defrost()
config.TASK_CONFIG.DATASET.SPLIT = self.config.EVAL.SPLIT
config.freeze()
eqa_cnn_pretrain_dataset = EQACNNPretrainDataset(config, mode="val")
eval_loader = DataLoader(
eqa_cnn_pretrain_dataset,
batch_size=config.IL.EQACNNPretrain.batch_size,
shuffle=False,
)
logger.info("[ eval_loader has {} samples ]".format(
len(eqa_cnn_pretrain_dataset)))
model = MultitaskCNN()
state_dict = torch.load(checkpoint_path)
model.load_state_dict(state_dict)
model.to(self.device).eval()
depth_loss = torch.nn.SmoothL1Loss()
ae_loss = torch.nn.SmoothL1Loss()
seg_loss = torch.nn.CrossEntropyLoss()
t = 0
avg_loss = 0.0
avg_l1 = 0.0
avg_l2 = 0.0
avg_l3 = 0.0
with torch.no_grad():
for batch in eval_loader:
t += 1
idx, gt_rgb, gt_depth, gt_seg = batch
gt_rgb = gt_rgb.to(self.device)
gt_depth = gt_depth.to(self.device)
gt_seg = gt_seg.to(self.device)
pred_seg, pred_depth, pred_rgb = model(gt_rgb)
l1 = seg_loss(pred_seg, gt_seg.long())
l2 = ae_loss(pred_rgb, gt_rgb)
l3 = depth_loss(pred_depth, gt_depth)
loss = l1 + (10 * l2) + (10 * l3)
avg_loss += loss.item()
avg_l1 += l1.item()
avg_l2 += l2.item()
avg_l3 += l3.item()
if t % config.LOG_INTERVAL == 0:
logger.info(
"[ Iter: {}; loss: {:.3f} ]".format(t, loss.item()), )
if (config.EVAL_SAVE_RESULTS
and t % config.EVAL_SAVE_RESULTS_INTERVAL == 0):
result_id = "ckpt_{}_{}".format(checkpoint_index,
idx[0].item())
result_path = os.path.join(self.config.RESULTS_DIR,
result_id)
self._save_results(
gt_rgb,
pred_rgb,
gt_seg,
pred_seg,
gt_depth,
pred_depth,
result_path,
)
avg_loss /= len(eval_loader)
avg_l1 /= len(eval_loader)
avg_l2 /= len(eval_loader)
avg_l3 /= len(eval_loader)
writer.add_scalar("avg val total loss", avg_loss, checkpoint_index)
writer.add_scalars(
"avg val individual_losses",
{
"seg_loss": avg_l1,
"ae_loss": avg_l2,
"depth_loss": avg_l3
},
checkpoint_index,
)
logger.info("[ Average loss: {:.3f} ]".format(avg_loss))
logger.info("[ Average seg loss: {:.3f} ]".format(avg_l1))
logger.info("[ Average autoencoder loss: {:.4f} ]".format(avg_l2))
logger.info("[ Average depthloss: {:.4f} ]".format(avg_l3))
def _eval_checkpoint(
self,
checkpoint_path: str,
writer: TensorboardWriter,
checkpoint_index: int = 0,
) -> None:
r"""Evaluates a single checkpoint.
Args:
checkpoint_path: path of checkpoint
writer: tensorboard writer object for logging to tensorboard
checkpoint_index: index of cur checkpoint for logging
Returns:
None
"""
# Map location CPU is almost always better than mapping to a CUDA device.
ckpt_dict = self.load_checkpoint(checkpoint_path, map_location="cpu")
if self.config.EVAL.USE_CKPT_CONFIG:
config = self._setup_eval_config(ckpt_dict["config"])
else:
config = self.config.clone()
ppo_cfg = config.RL.PPO
config.defrost()
config.TASK_CONFIG.DATASET.SPLIT = config.EVAL.SPLIT
config.freeze()
if len(self.config.VIDEO_OPTION) > 0:
config.defrost()
config.TASK_CONFIG.TASK.MEASUREMENTS.append("TOP_DOWN_MAP")
config.TASK_CONFIG.TASK.MEASUREMENTS.append("COLLISIONS")
config.freeze()
logger.info(f"env config: {config}")
self.envs = construct_envs(config, get_env_class(config.ENV_NAME))
self._setup_actor_critic_agent(ppo_cfg)
self.actor_critic.eval()
if self._static_encoder:
self._encoder = self.agent.actor_critic.net.visual_encoder
self.agent.load_state_dict(ckpt_dict["state_dict"])
self.actor_critic = self.agent.actor_critic
observations = self.envs.reset()
batch = batch_obs(observations, device=self.device)
if self._static_encoder:
batch["visual_features"] = self._encoder(batch)
batch["prev_visual_features"] = torch.zeros_like(
batch["visual_features"])
current_episode_reward = torch.zeros(self.envs.num_envs,
1,
device=self.device)
test_recurrent_hidden_states = torch.zeros(
self.actor_critic.net.num_recurrent_layers,
self.config.NUM_PROCESSES,
ppo_cfg.hidden_size,
device=self.device,
)
prev_actions = torch.zeros(self.config.NUM_PROCESSES,
1,
device=self.device,
dtype=torch.long)
not_done_masks = torch.zeros(self.config.NUM_PROCESSES,
1,
device=self.device)
stats_episodes = dict() # dict of dicts that stores stats per episode
rgb_frames = [[] for _ in range(self.config.NUM_PROCESSES)
] # type: List[List[np.ndarray]]
if len(self.config.VIDEO_OPTION) > 0:
os.makedirs(self.config.VIDEO_DIR, exist_ok=True)
number_of_eval_episodes = self.config.TEST_EPISODE_COUNT
if number_of_eval_episodes == -1:
number_of_eval_episodes = sum(self.envs.number_of_episodes)
else:
total_num_eps = sum(self.envs.number_of_episodes)
if total_num_eps < number_of_eval_episodes:
logger.warn(
f"Config specified {number_of_eval_episodes} eval episodes"
", dataset only has {total_num_eps}.")
logger.warn(f"Evaluating with {total_num_eps} instead.")
number_of_eval_episodes = total_num_eps
pbar = tqdm.tqdm(total=number_of_eval_episodes)
self.actor_critic.eval()
while (len(stats_episodes) < number_of_eval_episodes
and self.envs.num_envs > 0):
current_episodes = self.envs.current_episodes()
with torch.no_grad():
step_batch = batch
(
_,
actions,
_,
test_recurrent_hidden_states,
) = self.actor_critic.act(
batch,
test_recurrent_hidden_states,
prev_actions,
not_done_masks,
deterministic=False,
)
prev_actions.copy_(actions)
outputs = self.envs.step([a[0].item() for a in actions])
observations, rewards, dones, infos = [
list(x) for x in zip(*outputs)
]
batch = batch_obs(observations, device=self.device)
if self._static_encoder:
batch["prev_visual_features"] = step_batch["visual_features"]
batch["visual_features"] = self._encoder(batch)
not_done_masks = torch.tensor(
[[0.0] if done else [1.0] for done in dones],
dtype=torch.float,
device=self.device,
)
rewards = torch.tensor(rewards,
dtype=torch.float,
device=self.device).unsqueeze(1)
current_episode_reward += rewards
next_episodes = self.envs.current_episodes()
envs_to_pause = []
n_envs = self.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:
pbar.update()
episode_stats = dict()
episode_stats["reward"] = current_episode_reward[i].item()
episode_stats.update(
self._extract_scalars_from_info(infos[i]))
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 len(self.config.VIDEO_OPTION) > 0:
generate_video(
video_option=self.config.VIDEO_OPTION,
video_dir=self.config.VIDEO_DIR,
images=rgb_frames[i],
episode_id=current_episodes[i].episode_id,
checkpoint_idx=checkpoint_index,
metrics=self._extract_scalars_from_info(infos[i]),
tb_writer=writer,
)
rgb_frames[i] = []
# episode continues
elif len(self.config.VIDEO_OPTION) > 0:
frame = observations_to_image(observations[i], infos[i])
rgb_frames[i].append(frame)
(
self.envs,
test_recurrent_hidden_states,
not_done_masks,
current_episode_reward,
prev_actions,
batch,
rgb_frames,
) = self._pause_envs(
envs_to_pause,
self.envs,
test_recurrent_hidden_states,
not_done_masks,
current_episode_reward,
prev_actions,
batch,
rgb_frames,
)
num_episodes = len(stats_episodes)
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)
for k, v in aggregated_stats.items():
logger.info(f"Average episode {k}: {v:.4f}")
step_id = checkpoint_index
if "extra_state" in ckpt_dict and "step" in ckpt_dict["extra_state"]:
step_id = ckpt_dict["extra_state"]["step"]
writer.add_scalars(
"eval_reward",
{"average reward": aggregated_stats["reward"]},
step_id,
)
metrics = {k: v for k, v in aggregated_stats.items() if k != "reward"}
if len(metrics) > 0:
writer.add_scalars("eval_metrics", metrics, step_id)
self.envs.close()
def _eval_checkpoint(
self,
checkpoint_path: str,
writer: TensorboardWriter,
checkpoint_index: int = 0,
) -> None:
r"""Evaluates a single checkpoint.
Args:
checkpoint_path: path of checkpoint
writer: tensorboard writer object for logging to tensorboard
checkpoint_index: index of cur checkpoint for logging
Returns:
None
"""
if self._is_distributed:
raise RuntimeError("Evaluation does not support distributed mode")
# Map location CPU is almost always better than mapping to a CUDA device.
ckpt_dict = self.load_checkpoint(checkpoint_path, map_location="cpu")
if self.config.EVAL.USE_CKPT_CONFIG:
config = self._setup_eval_config(ckpt_dict["config"])
else:
config = self.config.clone()
ppo_cfg = config.RL.PPO
config.defrost()
config.TASK_CONFIG.DATASET.SPLIT = config.EVAL.SPLIT
config.freeze()
if len(self.config.VIDEO_OPTION) > 0:
config.defrost()
config.TASK_CONFIG.TASK.MEASUREMENTS.append("TOP_DOWN_MAP")
config.TASK_CONFIG.TASK.MEASUREMENTS.append("COLLISIONS")
config.freeze()
if config.VERBOSE:
logger.info(f"env config: {config}")
self._init_envs(config)
self._setup_actor_critic_agent(ppo_cfg)
self.agent.load_state_dict(ckpt_dict["state_dict"])
self.actor_critic = self.agent.actor_critic
observations = self.envs.reset()
batch = batch_obs(
observations, device=self.device, cache=self._obs_batching_cache
)
batch = apply_obs_transforms_batch(batch, self.obs_transforms)
current_episode_reward = torch.zeros(
self.envs.num_envs, 1, device="cpu"
)
test_recurrent_hidden_states = torch.zeros(
self.config.NUM_ENVIRONMENTS,
self.actor_critic.net.num_recurrent_layers,
ppo_cfg.hidden_size,
device=self.device,
)
prev_actions = torch.zeros(
self.config.NUM_ENVIRONMENTS,
1,
device=self.device,
dtype=torch.long,
)
not_done_masks = torch.zeros(
self.config.NUM_ENVIRONMENTS,
1,
device=self.device,
dtype=torch.bool,
)
stats_episodes: Dict[
Any, Any
] = {} # dict of dicts that stores stats per episode
rgb_frames = [
[] for _ in range(self.config.NUM_ENVIRONMENTS)
] # type: List[List[np.ndarray]]
if len(self.config.VIDEO_OPTION) > 0:
os.makedirs(self.config.VIDEO_DIR, exist_ok=True)
number_of_eval_episodes = self.config.TEST_EPISODE_COUNT
if number_of_eval_episodes == -1:
number_of_eval_episodes = sum(self.envs.number_of_episodes)
else:
total_num_eps = sum(self.envs.number_of_episodes)
if total_num_eps < number_of_eval_episodes:
logger.warn(
f"Config specified {number_of_eval_episodes} eval episodes"
", dataset only has {total_num_eps}."
)
logger.warn(f"Evaluating with {total_num_eps} instead.")
number_of_eval_episodes = total_num_eps
pbar = tqdm.tqdm(total=number_of_eval_episodes)
self.actor_critic.eval()
while (
len(stats_episodes) < number_of_eval_episodes
and self.envs.num_envs > 0
):
current_episodes = self.envs.current_episodes()
with torch.no_grad():
(
_,
actions,
_,
test_recurrent_hidden_states,
) = self.actor_critic.act(
batch,
test_recurrent_hidden_states,
prev_actions,
not_done_masks,
deterministic=False,
)
prev_actions.copy_(actions) # type: ignore
# NB: Move actions to CPU. If CUDA tensors are
# sent in to env.step(), that will create CUDA contexts
# in the subprocesses.
# For backwards compatibility, we also call .item() to convert to
# an int
step_data = [a.item() for a in actions.to(device="cpu")]
outputs = self.envs.step(step_data)
observations, rewards_l, dones, infos = [
list(x) for x in zip(*outputs)
]
batch = batch_obs(
observations,
device=self.device,
cache=self._obs_batching_cache,
)
batch = apply_obs_transforms_batch(batch, self.obs_transforms)
not_done_masks = torch.tensor(
[[not done] for done in dones],
dtype=torch.bool,
device="cpu",
)
rewards = torch.tensor(
rewards_l, dtype=torch.float, device="cpu"
).unsqueeze(1)
current_episode_reward += rewards
next_episodes = self.envs.current_episodes()
envs_to_pause = []
n_envs = self.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 not_done_masks[i].item():
pbar.update()
episode_stats = {}
episode_stats["reward"] = current_episode_reward[i].item()
episode_stats.update(
self._extract_scalars_from_info(infos[i])
)
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 len(self.config.VIDEO_OPTION) > 0:
generate_video(
video_option=self.config.VIDEO_OPTION,
video_dir=self.config.VIDEO_DIR,
images=rgb_frames[i],
episode_id=current_episodes[i].episode_id,
checkpoint_idx=checkpoint_index,
metrics=self._extract_scalars_from_info(infos[i]),
tb_writer=writer,
)
rgb_frames[i] = []
# episode continues
elif len(self.config.VIDEO_OPTION) > 0:
# TODO move normalization / channel changing out of the policy and undo it here
frame = observations_to_image(
{k: v[i] for k, v in batch.items()}, infos[i]
)
rgb_frames[i].append(frame)
not_done_masks = not_done_masks.to(device=self.device)
(
self.envs,
test_recurrent_hidden_states,
not_done_masks,
current_episode_reward,
prev_actions,
batch,
rgb_frames,
) = self._pause_envs(
envs_to_pause,
self.envs,
test_recurrent_hidden_states,
not_done_masks,
current_episode_reward,
prev_actions,
batch,
rgb_frames,
)
num_episodes = len(stats_episodes)
aggregated_stats = {}
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
)
for k, v in aggregated_stats.items():
logger.info(f"Average episode {k}: {v:.4f}")
step_id = checkpoint_index
if "extra_state" in ckpt_dict and "step" in ckpt_dict["extra_state"]:
step_id = ckpt_dict["extra_state"]["step"]
writer.add_scalars(
"eval_reward",
{"average reward": aggregated_stats["reward"]},
step_id,
)
metrics = {k: v for k, v in aggregated_stats.items() if k != "reward"}
if len(metrics) > 0:
writer.add_scalars("eval_metrics", metrics, step_id)
self.envs.close()
def _eval_checkpoint(
self,
checkpoint_path: str,
writer: TensorboardWriter,
checkpoint_index: int = 0,
) -> None:
r"""Evaluates a single checkpoint.
Args:
checkpoint_path: path of checkpoint
writer: tensorboard writer object for logging to tensorboard
checkpoint_index: index of cur checkpoint for logging
Returns:
None
"""
ckpt_dict = self.load_checkpoint(checkpoint_path,
map_location=self.device)
config = self._setup_eval_config(ckpt_dict["config"])
ppo_cfg = config.RL.PPO
if len(self.config.VIDEO_OPTION) > 0:
config.defrost()
config.TASK_CONFIG.TASK.MEASUREMENTS.append("TOP_DOWN_MAP")
config.TASK_CONFIG.TASK.MEASUREMENTS.append("COLLISIONS")
config.freeze()
logger.info(f"env config: {config}")
self.envs = construct_envs(self.config,
get_env_class(self.config.ENV_NAME))
self._setup_actor_critic_agent(ppo_cfg)
self.agent.load_state_dict(ckpt_dict["state_dict"])
self.actor_critic = self.agent.actor_critic
# get name of performance metric, e.g. "spl"
metric_name = self.config.TASK_CONFIG.TASK.MEASUREMENTS[0]
metric_cfg = getattr(self.config.TASK_CONFIG.TASK, metric_name)
measure_type = baseline_registry.get_measure(metric_cfg.TYPE)
assert measure_type is not None, "invalid measurement type {}".format(
metric_cfg.TYPE)
self.metric_uuid = measure_type(None, None)._get_uuid()
observations = self.envs.reset()
batch = batch_obs(observations)
for sensor in batch:
batch[sensor] = batch[sensor].to(self.device)
current_episode_reward = torch.zeros(self.envs.num_envs,
1,
device=self.device)
test_recurrent_hidden_states = torch.zeros(
self.actor_critic.net.num_recurrent_layers,
self.config.NUM_PROCESSES,
ppo_cfg.hidden_size,
device=self.device,
)
prev_actions = torch.zeros(self.config.NUM_PROCESSES,
1,
device=self.device,
dtype=torch.long)
not_done_masks = torch.zeros(self.config.NUM_PROCESSES,
1,
device=self.device)
stats_episodes = dict() # dict of dicts that stores stats per episode
rgb_frames = [
[]
] * self.config.NUM_PROCESSES # type: List[List[np.ndarray]]
if len(self.config.VIDEO_OPTION) > 0:
os.makedirs(self.config.VIDEO_DIR, exist_ok=True)
while (len(stats_episodes) < self.config.TEST_EPISODE_COUNT
and self.envs.num_envs > 0):
current_episodes = self.envs.current_episodes()
with torch.no_grad():
_, actions, _, test_recurrent_hidden_states = self.actor_critic.act(
batch,
test_recurrent_hidden_states,
prev_actions,
not_done_masks,
deterministic=False,
)
prev_actions.copy_(actions)
outputs = self.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(self.device)
not_done_masks = torch.tensor(
[[0.0] if done else [1.0] for done in dones],
dtype=torch.float,
device=self.device,
)
rewards = torch.tensor(rewards,
dtype=torch.float,
device=self.device).unsqueeze(1)
current_episode_reward += rewards
next_episodes = self.envs.current_episodes()
envs_to_pause = []
n_envs = self.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[self.metric_uuid] = infos[i][
self.metric_uuid]
episode_stats["success"] = int(
infos[i][self.metric_uuid] > 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 len(self.config.VIDEO_OPTION) > 0:
generate_video(
video_option=self.config.VIDEO_OPTION,
video_dir=self.config.VIDEO_DIR,
images=rgb_frames[i],
episode_id=current_episodes[i].episode_id,
checkpoint_idx=checkpoint_index,
metric_name=self.metric_uuid,
metric_value=infos[i][self.metric_uuid],
tb_writer=writer,
)
rgb_frames[i] = []
# episode continues
elif len(self.config.VIDEO_OPTION) > 0:
frame = observations_to_image(observations[i], infos[i])
rgb_frames[i].append(frame)
# pausing self.envs with no new episode
if len(envs_to_pause) > 0:
state_index = list(range(self.envs.num_envs))
for idx in reversed(envs_to_pause):
state_index.pop(idx)
self.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]
prev_actions = prev_actions[state_index]
for k, v in batch.items():
batch[k] = v[state_index]
if len(self.config.VIDEO_OPTION) > 0:
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_metric_mean = aggregated_stats[self.metric_uuid] / num_episodes
episode_success_mean = aggregated_stats["success"] / num_episodes
logger.info(f"Average episode reward: {episode_reward_mean:.6f}")
logger.info(f"Average episode success: {episode_success_mean:.6f}")
logger.info(
f"Average episode {self.metric_uuid}: {episode_metric_mean:.6f}")
writer.add_scalars(
"eval_reward",
{"average reward": episode_reward_mean},
checkpoint_index,
)
writer.add_scalars(
f"eval_{self.metric_uuid}",
{f"average {self.metric_uuid}": episode_metric_mean},
checkpoint_index,
)
writer.add_scalars(
"eval_success",
{"average success": episode_success_mean},
checkpoint_index,
)
self.envs.close()
def _eval_checkpoint(self,
checkpoint_path: str,
writer: TensorboardWriter,
checkpoint_index: int = 0,
log_diagnostics=[],
output_dir='.',
label='.',
num_eval_runs=1) -> None:
r"""Evaluates a single checkpoint.
Args:
checkpoint_path: path of checkpoint
writer: tensorboard writer object for logging to tensorboard
checkpoint_index: index of cur checkpoint for logging
Returns:
None
"""
if checkpoint_index == -1:
ckpt_file = checkpoint_path.split('/')[-1]
split_info = ckpt_file.split('.')
checkpoint_index = split_info[1]
# Map location CPU is almost always better than mapping to a CUDA device.
ckpt_dict = self.load_checkpoint(checkpoint_path, map_location="cpu")
if self.config.EVAL.USE_CKPT_CONFIG:
config = self._setup_eval_config(ckpt_dict["config"])
else:
config = self.config.clone()
ppo_cfg = config.RL.PPO
task_cfg = config.TASK_CONFIG.TASK
config.defrost()
config.TASK_CONFIG.DATASET.SPLIT = config.EVAL.SPLIT
config.freeze()
if len(self.config.VIDEO_OPTION) > 0:
config.defrost()
config.TASK_CONFIG.TASK.MEASUREMENTS.append("TOP_DOWN_MAP")
config.TASK_CONFIG.TASK.MEASUREMENTS.append("COLLISIONS")
config.freeze()
logger.info(f"env config: {config}")
self.envs = construct_envs(config, get_env_class(config.ENV_NAME))
# pass in aux config if we're doing attention
aux_cfg = self.config.RL.AUX_TASKS
self._setup_actor_critic_agent(ppo_cfg, task_cfg, aux_cfg)
# Check if we accidentally recorded `visual_resnet` in our checkpoint and drop it (it's redundant with `visual_encoder`)
ckpt_dict['state_dict'] = {
k: v
for k, v in ckpt_dict['state_dict'].items()
if 'visual_resnet' not in k
}
self.agent.load_state_dict(ckpt_dict["state_dict"])
logger.info("agent number of trainable parameters: {}".format(
sum(param.numel() for param in self.agent.parameters()
if param.requires_grad)))
self.actor_critic = self.agent.actor_critic
observations = self.envs.reset()
batch = batch_obs(observations, device=self.device)
current_episode_reward = torch.zeros(self.envs.num_envs,
1,
device=self.device)
test_recurrent_hidden_states = torch.zeros(
self.actor_critic.net.num_recurrent_layers,
self.config.NUM_PROCESSES,
ppo_cfg.hidden_size,
device=self.device,
)
_, num_recurrent_memories, _ = self._setup_auxiliary_tasks(
aux_cfg, ppo_cfg, task_cfg, is_eval=True)
if self.config.RL.PPO.policy in MULTIPLE_BELIEF_CLASSES:
aux_tasks = self.config.RL.AUX_TASKS.tasks
num_recurrent_memories = len(self.config.RL.AUX_TASKS.tasks)
test_recurrent_hidden_states = test_recurrent_hidden_states.unsqueeze(
2).repeat(1, 1, num_recurrent_memories, 1)
prev_actions = torch.zeros(self.config.NUM_PROCESSES,
1,
device=self.device,
dtype=torch.long)
not_done_masks = torch.zeros(self.config.NUM_PROCESSES,
1,
device=self.device)
stats_episodes = dict() # dict of dicts that stores stats per episode
rgb_frames = [[] for _ in range(self.config.NUM_PROCESSES)
] # type: List[List[np.ndarray]]
if len(self.config.VIDEO_OPTION) > 0:
os.makedirs(self.config.VIDEO_DIR, exist_ok=True)
number_of_eval_episodes = self.config.TEST_EPISODE_COUNT
if number_of_eval_episodes == -1:
number_of_eval_episodes = sum(self.envs.number_of_episodes)
else:
total_num_eps = sum(self.envs.number_of_episodes)
if total_num_eps < number_of_eval_episodes:
logger.warn(
f"Config specified {number_of_eval_episodes} eval episodes"
", dataset only has {total_num_eps}.")
logger.warn(f"Evaluating with {total_num_eps} instead.")
number_of_eval_episodes = total_num_eps
videos_cap = 2 # number of videos to generate per checkpoint
if len(log_diagnostics) > 0:
videos_cap = 10
# video_indices = random.sample(range(self.config.TEST_EPISODE_COUNT),
# min(videos_cap, self.config.TEST_EPISODE_COUNT))
video_indices = range(10)
print(f"Videos: {video_indices}")
total_stats = []
dones_per_ep = dict()
# Logging more extensive evaluation stats for analysis
if len(log_diagnostics) > 0:
d_stats = {}
for d in log_diagnostics:
d_stats[d] = [
[] for _ in range(self.config.NUM_PROCESSES)
] # stored as nested list envs x timesteps x k (# tasks)
pbar = tqdm.tqdm(total=number_of_eval_episodes * num_eval_runs)
self.agent.eval()
while (len(stats_episodes) < number_of_eval_episodes * num_eval_runs
and self.envs.num_envs > 0):
current_episodes = self.envs.current_episodes()
with torch.no_grad():
weights_output = None
if self.config.RL.PPO.policy in MULTIPLE_BELIEF_CLASSES:
weights_output = torch.empty(self.envs.num_envs,
len(aux_tasks))
(
_,
actions,
_,
test_recurrent_hidden_states,
) = self.actor_critic.act(batch,
test_recurrent_hidden_states,
prev_actions,
not_done_masks,
deterministic=False,
weights_output=weights_output)
prev_actions.copy_(actions)
for i in range(self.envs.num_envs):
if Diagnostics.actions in log_diagnostics:
d_stats[Diagnostics.actions][i].append(
prev_actions[i].item())
if Diagnostics.weights in log_diagnostics:
aux_weights = None if weights_output is None else weights_output[
i]
if aux_weights is not None:
d_stats[Diagnostics.weights][i].append(
aux_weights.half().tolist())
outputs = self.envs.step([a[0].item() for a in actions])
observations, rewards, dones, infos = [
list(x) for x in zip(*outputs)
]
batch = batch_obs(observations, device=self.device)
not_done_masks = torch.tensor(
[[0.0] if done else [1.0] for done in dones],
dtype=torch.float,
device=self.device,
)
rewards = torch.tensor(rewards,
dtype=torch.float,
device=self.device).unsqueeze(1)
current_episode_reward += rewards
next_episodes = self.envs.current_episodes()
envs_to_pause = []
n_envs = self.envs.num_envs
for i in range(n_envs):
next_k = (
next_episodes[i].scene_id,
next_episodes[i].episode_id,
)
if dones_per_ep.get(next_k, 0) == num_eval_runs:
envs_to_pause.append(i) # wait for the rest
if not_done_masks[i].item() == 0:
episode_stats = dict()
episode_stats["reward"] = current_episode_reward[i].item()
episode_stats.update(
self._extract_scalars_from_info(infos[i]))
current_episode_reward[i] = 0
# use scene_id + episode_id as unique id for storing stats
k = (
current_episodes[i].scene_id,
current_episodes[i].episode_id,
)
dones_per_ep[k] = dones_per_ep.get(k, 0) + 1
if dones_per_ep.get(k, 0) == 1 and len(
self.config.VIDEO_OPTION) > 0 and len(
stats_episodes) in video_indices:
logger.info(f"Generating video {len(stats_episodes)}")
category = getattr(current_episodes[i],
"object_category", "")
if category != "":
category += "_"
try:
generate_video(
video_option=self.config.VIDEO_OPTION,
video_dir=self.config.VIDEO_DIR,
images=rgb_frames[i],
episode_id=current_episodes[i].episode_id,
checkpoint_idx=checkpoint_index,
metrics=self._extract_scalars_from_info(
infos[i]),
tag=f"{category}{label}",
tb_writer=writer,
)
except Exception as e:
logger.warning(str(e))
rgb_frames[i] = []
stats_episodes[(
current_episodes[i].scene_id,
current_episodes[i].episode_id,
dones_per_ep[k],
)] = episode_stats
if len(log_diagnostics) > 0:
diagnostic_info = dict()
for metric in log_diagnostics:
diagnostic_info[metric] = d_stats[metric][i]
d_stats[metric][i] = []
if Diagnostics.top_down_map in log_diagnostics:
top_down_map = torch.tensor([])
if len(self.config.VIDEO_OPTION) > 0:
top_down_map = infos[i]["top_down_map"]["map"]
top_down_map = maps.colorize_topdown_map(
top_down_map, fog_of_war_mask=None)
diagnostic_info.update(
dict(top_down_map=top_down_map))
total_stats.append(
dict(
stats=episode_stats,
did_stop=bool(prev_actions[i] == 0),
episode_info=attr.asdict(current_episodes[i]),
info=diagnostic_info,
))
pbar.update()
# episode continues
else:
if len(self.config.VIDEO_OPTION) > 0:
aux_weights = None if weights_output is None else weights_output[
i]
frame = observations_to_image(
observations[i], infos[i],
current_episode_reward[i].item(), aux_weights,
aux_tasks)
rgb_frames[i].append(frame)
if Diagnostics.gps in log_diagnostics:
d_stats[Diagnostics.gps][i].append(
observations[i]["gps"].tolist())
if Diagnostics.heading in log_diagnostics:
d_stats[Diagnostics.heading][i].append(
observations[i]["heading"].tolist())
(
self.envs,
test_recurrent_hidden_states,
not_done_masks,
current_episode_reward,
prev_actions,
batch,
rgb_frames,
) = self._pause_envs(
envs_to_pause,
self.envs,
test_recurrent_hidden_states,
not_done_masks,
current_episode_reward,
prev_actions,
batch,
rgb_frames,
)
num_episodes = len(stats_episodes)
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)
for k, v in aggregated_stats.items():
logger.info(f"Average episode {k}: {v:.4f}")
step_id = checkpoint_index
if "extra_state" in ckpt_dict and "step" in ckpt_dict["extra_state"]:
step_id = ckpt_dict["extra_state"]["step"]
writer.add_scalars(
"eval_reward",
{"average reward": aggregated_stats["reward"]},
step_id,
)
metrics = {k: v for k, v in aggregated_stats.items() if k != "reward"}
if len(metrics) > 0:
writer.add_scalars("eval_metrics", metrics, step_id)
logger.info("eval_metrics")
logger.info(metrics)
if len(log_diagnostics) > 0:
os.makedirs(output_dir, exist_ok=True)
eval_fn = f"{label}.json"
with open(os.path.join(output_dir, eval_fn), 'w',
encoding='utf-8') as f:
json.dump(total_stats, f, ensure_ascii=False, indent=4)
self.envs.close()
def _eval_checkpoint(
self,
checkpoint_path: str,
writer: TensorboardWriter,
cur_ckpt_idx: int = 0,
) -> None:
r"""
Evaluates a single checkpoint
Args:
checkpoint_path: path of checkpoint
writer: tensorboard writer object for logging to tensorboard
cur_ckpt_idx: index of cur checkpoint for logging
Returns:
None
"""
ckpt_dict = self.load_checkpoint(checkpoint_path,
map_location=self.device)
ckpt_config = ckpt_dict["config"]
config = self.config.clone()
ckpt_cmd_opts = ckpt_config.CMD_TRAILING_OPTS
eval_cmd_opts = config.CMD_TRAILING_OPTS
# config merge priority: eval_opts > ckpt_opts > eval_cfg > ckpt_cfg
# first line for old checkpoint compatibility
config.merge_from_other_cfg(ckpt_config)
config.merge_from_other_cfg(self.config)
config.merge_from_list(ckpt_cmd_opts)
config.merge_from_list(eval_cmd_opts)
ppo_cfg = config.TRAINER.RL.PPO
config.TASK_CONFIG.defrost()
config.TASK_CONFIG.DATASET.SPLIT = "val"
agent_sensors = ppo_cfg.sensors.strip().split(",")
config.TASK_CONFIG.SIMULATOR.AGENT_0.SENSORS = agent_sensors
if self.video_option:
config.TASK_CONFIG.TASK.MEASUREMENTS.append("TOP_DOWN_MAP")
config.TASK_CONFIG.TASK.MEASUREMENTS.append("COLLISIONS")
config.freeze()
logger.info(f"env config: {config}")
self.envs = construct_envs(config, NavRLEnv)
self._setup_actor_critic_agent(ppo_cfg)
self.agent.load_state_dict(ckpt_dict["state_dict"])
self.actor_critic = self.agent.actor_critic
observations = self.envs.reset()
batch = batch_obs(observations)
for sensor in batch:
batch[sensor] = batch[sensor].to(self.device)
current_episode_reward = torch.zeros(self.envs.num_envs,
1,
device=self.device)
test_recurrent_hidden_states = torch.zeros(ppo_cfg.num_processes,
ppo_cfg.hidden_size,
device=self.device)
not_done_masks = torch.zeros(ppo_cfg.num_processes,
1,
device=self.device)
stats_episodes = dict() # dict of dicts that stores stats per episode
rgb_frames = [[]
] * ppo_cfg.num_processes # type: List[List[np.ndarray]]
if self.video_option:
os.makedirs(ppo_cfg.video_dir, exist_ok=True)
while (len(stats_episodes) < ppo_cfg.count_test_episodes
and self.envs.num_envs > 0):
current_episodes = self.envs.current_episodes()
with torch.no_grad():
_, actions, _, test_recurrent_hidden_states = self.actor_critic.act(
batch,
test_recurrent_hidden_states,
not_done_masks,
deterministic=False,
)
outputs = self.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(self.device)
not_done_masks = torch.tensor(
[[0.0] if done else [1.0] for done in dones],
dtype=torch.float,
device=self.device,
)
rewards = torch.tensor(rewards,
dtype=torch.float,
device=self.device).unsqueeze(1)
current_episode_reward += rewards
next_episodes = self.envs.current_episodes()
envs_to_pause = []
n_envs = self.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 self.video_option:
generate_video(
ppo_cfg,
rgb_frames[i],
current_episodes[i].episode_id,
cur_ckpt_idx,
infos[i]["spl"],
writer,
)
rgb_frames[i] = []
# episode continues
elif self.video_option:
frame = observations_to_image(observations[i], infos[i])
rgb_frames[i].append(frame)
# pausing self.envs with no new episode
if len(envs_to_pause) > 0:
state_index = list(range(self.envs.num_envs))
for idx in reversed(envs_to_pause):
state_index.pop(idx)
self.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 self.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,
)
def _eval_checkpoint(
self,
checkpoint_path: str,
writer: TensorboardWriter,
checkpoint_index: int = 0,
) -> None:
r"""Evaluates a single checkpoint.
Args:
checkpoint_path: path of checkpoint
writer: tensorboard writer object for logging to tensorboard
checkpoint_index: index of cur checkpoint for logging
Returns:
None
"""
NUM_PROCESSES = 3
# Map location CPU is almost always better than mapping to a CUDA device.
ckpt_dict = self.load_checkpoint(checkpoint_path, map_location="cpu")
if self.config.EVAL.USE_CKPT_CONFIG:
config = self._setup_eval_config(ckpt_dict["config"])
else:
config = self.config.clone()
ppo_cfg = config.RL.PPO
config.defrost()
config.TASK_CONFIG.DATASET.SPLIT = config.EVAL.SPLIT
config.NUM_PROCESSES = NUM_PROCESSES
config.NUM_VAL_PROCESSES = 0
config.TASK_CONFIG.ENVIRONMENT.ITERATOR_OPTIONS.MAX_SCENE_REPEAT_EPISODES = 10
config.TEST_EPISODE_COUNT = 500
config.RL.PPO.pretrained = False
config.RL.PPO.pretrained_encoder = False
if torch.cuda.device_count() <= 1:
config.TORCH_GPU_ID = 0
config.SIMULATOR_GPU_ID = 0
else:
config.TORCH_GPU_ID = 0
config.SIMULATOR_GPU_ID = 1
config.VIDEO_DIR += '_eval'
config.freeze()
if len(self.config.VIDEO_OPTION) > 0:
config.defrost()
config.TASK_CONFIG.TASK.MEASUREMENTS.append("TOP_DOWN_MAP")
config.TASK_CONFIG.TASK.MEASUREMENTS.append("COLLISIONS")
config.freeze()
self.config = config
logger.info(f"env config: {config}")
self.envs = construct_envs(config,
eval(self.config.ENV_NAME),
run_type='eval')
self._setup_actor_critic_agent(ppo_cfg)
print(config.POLICY)
# if 'SMT' in config.POLICY:
# sd = torch.load('visual_embedding18.pth')
# self.actor_critic.net.visual_encoder.load_state_dict(sd['visual_encoder'])
# self.actor_critic.net.prev_action_embedding.load_state_dict(sd['prev_action_embedding'])
# self.actor_critic.net.visual_encoder.cuda()
# self.actor_critic.net.prev_action_embedding.cuda()
# self.envs.setup_embedding_network(self.actor_critic.net.visual_encoder, self.actor_critic.net.prev_action_embedding)
# print('-----------------------------setup pretrained visual embedding network')
try:
self.agent.load_state_dict(ckpt_dict["state_dict"])
except:
raise
initial_state_dict = self.actor_critic.state_dict()
initial_state_dict.update({
k[len("actor_critic."):]: v
for k, v in ckpt_dict['state_dict'].items()
if k[len("actor_critic."):] in initial_state_dict and v.shape
== initial_state_dict[k[len("actor_critic."):]].shape
})
print({
k[len("actor_critic."):]: v
for k, v in ckpt_dict['state_dict'].items()
if k[len("actor_critic."):] in initial_state_dict and v.shape
== initial_state_dict[k[len("actor_critic."):]].shape
}.keys())
self.actor_critic.load_state_dict(initial_state_dict)
self.actor_critic = self.agent.actor_critic
batch = self.envs.reset()
#batch = batch_obs(observations, device=self.device)
current_episode_reward = torch.zeros(self.envs.num_envs,
1,
device=self.device)
test_recurrent_hidden_states = torch.zeros(
self.actor_critic.net.num_recurrent_layers,
NUM_PROCESSES,
ppo_cfg.hidden_size,
device=self.device,
)
prev_actions = torch.zeros(NUM_PROCESSES,
1,
device=self.device,
dtype=torch.long)
not_done_masks = torch.zeros(NUM_PROCESSES, 1, device=self.device)
stats_episodes = dict() # dict of dicts that stores stats per episode
rgb_frames = [[] for _ in range(NUM_PROCESSES)
] # type: List[List[np.ndarray]]
if len(self.config.VIDEO_OPTION) > 0:
os.makedirs(self.config.VIDEO_DIR, exist_ok=True)
pbar = tqdm.tqdm(total=self.config.TEST_EPISODE_COUNT)
self.actor_critic.eval()
while (len(stats_episodes) < self.config.TEST_EPISODE_COUNT
and self.envs.num_envs > 0):
#print(len(stats_episodes), self.config.TEST_EPISODE_COUNT, self.envs.num_envs)
current_episodes = self.envs.current_episodes()
with torch.no_grad():
(
_,
actions,
_,
test_recurrent_hidden_states,
) = self.actor_critic.act(
batch,
test_recurrent_hidden_states,
prev_actions,
not_done_masks,
deterministic=False,
)
actions = actions.unsqueeze(1)
prev_actions.copy_(actions)
batch, rewards, dones, infos = self.envs.step(
[a[0].item() for a in actions])
not_done_masks = torch.tensor(
[[0.0] if done else [1.0] for done in dones],
dtype=torch.float,
device=self.device,
)
rewards = torch.tensor(rewards,
dtype=torch.float,
device=self.device).unsqueeze(1)
current_episode_reward += rewards
next_episodes = self.envs.current_episodes()
envs_to_pause = []
n_envs = self.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:
pbar.update()
episode_stats = dict()
episode_stats["reward"] = current_episode_reward[i].item()
episode_stats.update(
self._extract_scalars_from_info(infos[i]))
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
# episode continues
#elif len(self.config.VIDEO_OPTION) > 0:
# frame = self.envs.call_at(i, 'render', {'mode': 'rgb_array'})#observations_to_image(observations[i], infos[i])
# rgb_frames[i].append(frame)
num_episodes = len(stats_episodes)
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)
for k, v in aggregated_stats.items():
logger.info(f"Average episode {k}: {v:.4f}")
step_id = checkpoint_index
if "extra_state" in ckpt_dict and "step" in ckpt_dict["extra_state"]:
step_id = ckpt_dict["extra_state"]["step"]
writer.add_scalars(
"eval_reward",
{"average reward": aggregated_stats["reward"]},
step_id,
)
metrics = {k: v for k, v in aggregated_stats.items() if k != "reward"}
if len(metrics) > 0:
writer.add_scalars("eval_metrics", metrics, step_id)
self.envs.close()
def _eval_checkpoint(
self,
checkpoint_path: str,
writer: TensorboardWriter,
checkpoint_index: int = 0,
) -> None:
r"""Evaluates a single checkpoint.
Args:
checkpoint_path: path of checkpoint
writer: tensorboard writer object for logging to tensorboard
checkpoint_index: index of cur checkpoint for logging
Returns:
None
"""
self.add_new_based_on_cfg()
# Map location CPU is almost always better than mapping to a CUDA device.
ckpt_dict = self.load_checkpoint(checkpoint_path, map_location="cpu")
# ==========================================================================================
# -- Update config for eval
if self.config.EVAL.USE_CKPT_CONFIG:
config = self._setup_eval_config(ckpt_dict["config"])
else:
config = self.config.clone()
ppo_cfg = config.RL.PPO
# # Mostly for visualization
# config.defrost()
# config.TASK_CONFIG.SIMULATOR.HABITAT_SIM_V0.GPU_GPU = False
# config.freeze()
split = config.TASK_CONFIG.DATASET.SPLIT
config.defrost()
config.TASK_CONFIG.TASK.MEASUREMENTS.append("TOP_DOWN_MAP")
config.TASK_CONFIG.TASK.MEASUREMENTS.append("COLLISIONS")
config.freeze()
# ==========================================================================================
num_procs = self.config.NUM_PROCESSES
device = self.device
cfg = self.config
logger.info(f"env config: {config}")
self.envs = construct_envs(config, get_env_class(self.config.ENV_NAME))
num_envs = self.envs.num_envs
self._setup_actor_critic_agent(ppo_cfg, train=False)
self.agent.load_state_dict(ckpt_dict["state_dict"])
self.actor_critic = self.agent.actor_critic
self.r_policy = self.agent.actor_critic.reachability_policy
aux_models = self.actor_critic.net.aux_models
other_losses = dict({
k: torch.zeros(num_envs, 1, device=device)
for k in aux_models.keys()
})
other_losses_action = dict({
k: torch.zeros(num_envs,
self.envs.action_spaces[0].n,
device=device)
for k in aux_models.keys()
})
num_steps = torch.zeros(num_envs, 1, device=device)
# Config aux models for eval per item in batch
for k, maux in aux_models.items():
maux.set_per_element_loss()
total_loss = 0
if config.EVAL_MODE:
self.agent.eval()
self.r_policy.eval()
# get name of performance metric, e.g. "spl"
metric_name = cfg.TASK_CONFIG.TASK.MEASUREMENTS[0]
metric_cfg = getattr(cfg.TASK_CONFIG.TASK, metric_name)
measure_type = baseline_registry.get_measure(metric_cfg.TYPE)
assert measure_type is not None, "invalid measurement type {}".format(
metric_cfg.TYPE)
self.metric_uuid = measure_type(sim=None, task=None,
config=None)._get_uuid()
observations = self.envs.reset()
batch = batch_obs_augment_aux(observations, self.envs.get_shared_mem())
info_data_keys = ["discovered", "collisions_wall", "collisions_prox"]
log_data_keys = [
"current_episode_reward", "current_episode_go_reward"
] + info_data_keys
log_data = dict({
k: torch.zeros(num_envs, 1, device=device)
for k in log_data_keys
})
info_data = dict({k: log_data[k] for k in info_data_keys})
test_recurrent_hidden_states = torch.zeros(
self.actor_critic.net.num_recurrent_layers,
num_procs,
ppo_cfg.hidden_size,
device=device,
)
prev_actions = torch.zeros(num_procs,
1,
device=device,
dtype=torch.long)
not_done_masks = torch.zeros(num_procs, 1, device=device)
stats_episodes = dict() # dict of dicts that stores stats per episode
stats_episodes_scenes = dict(
) # dict of number of collected stats from
# each scene
max_test_ep_count = cfg.TEST_EPISODE_COUNT
# TODO this should depend on number of scenes :(
# TODO But than envs shouldn't be paused but fast-fwd to next scene
# TODO We consider num envs == num scenes
max_ep_per_env = max_test_ep_count / float(num_envs)
rgb_frames = [[] for _ in range(num_procs)
] # type: List[List[np.ndarray]]
if len(cfg.VIDEO_OPTION) > 0:
os.makedirs(cfg.VIDEO_DIR, exist_ok=True)
video_log_int = cfg.VIDEO_OPTION_INTERVAL
num_frames = 0
plot_pos = -1
prev_true_pos = []
prev_pred_pos = []
while (len(stats_episodes) <= cfg.TEST_EPISODE_COUNT and num_envs > 0):
current_episodes = self.envs.current_episodes()
with torch.no_grad():
prev_hidden = test_recurrent_hidden_states
_, actions, _, test_recurrent_hidden_states, aux_out \
= self.actor_critic.act(
batch,
test_recurrent_hidden_states,
prev_actions,
not_done_masks,
deterministic=False
)
prev_actions.copy_(actions)
if 'action' in batch:
prev_actions = batch['action'].unsqueeze(1).to(
actions.device).long()
for k, v in aux_out.items():
loss = aux_models[k].calc_loss(v, batch, prev_hidden,
prev_actions,
not_done_masks, actions)
total_loss += loss
if other_losses[k] is None:
other_losses[k] = loss
else:
other_losses[k] += loss.unsqueeze(1)
if len(prev_actions) == 1:
other_losses_action[k][0, prev_actions.item()] += \
loss.item()
# ==================================================================================
# - Hacky logs
if plot_pos >= 0:
prev_true_pos.append(batch["gps_compass_start"]
[plot_pos].data[:2].cpu().numpy())
prev_pred_pos.append(aux_out["rel_start_pos_reg"]
[plot_pos].data.cpu().numpy() * 15)
if num_frames % 10 == 0:
xx, yy = [], []
for x, y in prev_true_pos:
xx.append(x)
yy.append(y)
plt.scatter(xx, yy, label="true_pos")
xx, yy = [], []
for x, y in prev_pred_pos:
xx.append(x)
yy.append(y)
plt.scatter(xx, yy, label="pred_pos")
plt.legend()
plt.show()
plt.waitforbuttonpress()
plt.close()
# ==================================================================================
num_steps += 1
outputs = self.envs.step([a[0].item() for a in actions])
observations, rewards, dones, infos = [
list(x) for x in zip(*outputs)
]
not_done_masks = torch.tensor(
[[0.0] if done else [1.0] for done in dones],
dtype=torch.float,
device=device,
)
map_values = self._get_mapping(observations, aux_out)
batch = batch_obs_augment_aux(observations,
self.envs.get_shared_mem(),
device=device,
map_values=map_values,
masks=not_done_masks)
valid_map_size = [
float(ifs["top_down_map"]["valid_map"].sum()) for ifs in infos
]
discovered_factor = [
infos[ix]["top_down_map"]["explored_map"].sum() /
valid_map_size[ix] for ix in range(len(infos))
]
seen_factor = [
infos[ix]["top_down_map"]["ful_fog_of_war_mask"].sum() /
valid_map_size[ix] for ix in range(len(infos))
]
rewards = torch.tensor(rewards, dtype=torch.float,
device=device).unsqueeze(1)
log_data["current_episode_reward"] += rewards
# -- Add intrinsic Reward
if self.only_intrinsic_reward:
rewards.zero_()
if self.r_enabled:
ir_rewards = self._add_intrinsic_reward(
batch, actions, rewards, not_done_masks)
log_data["current_episode_go_reward"] += ir_rewards
rewards += ir_rewards
# Log other info from infos dict
for iii, info in enumerate(infos):
for k_info, v_info in info_data.items():
v_info[iii] += info[k_info]
next_episodes = self.envs.current_episodes()
envs_to_pause = []
n_envs = num_envs
for i in range(n_envs):
scene = next_episodes[i].scene_id
if scene not in stats_episodes_scenes:
stats_episodes_scenes[scene] = 0
if stats_episodes_scenes[scene] >= max_ep_per_env:
envs_to_pause.append(i)
# episode ended
if not_done_masks[i].item() == 0:
episode_stats = dict()
episode_stats[self.metric_uuid] = infos[i][
self.metric_uuid]
episode_stats["success"] = int(
infos[i][self.metric_uuid] > 0)
for kk, vv in log_data.items():
episode_stats[kk] = vv[i].item()
vv[i] = 0
episode_stats["map_discovered"] = discovered_factor[i]
episode_stats["map_seen"] = seen_factor[i]
for k, v in other_losses.items():
episode_stats[k] = v[i].item() / num_steps[i].item()
other_losses_action[k][i].fill_(0)
other_losses[k][i] = 0
num_steps[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
print(f"Episode {len(stats_episodes)} stats:",
episode_stats)
stats_episodes_scenes[current_episodes[i].scene_id] += 1
if len(cfg.VIDEO_OPTION
) > 0 and checkpoint_index % video_log_int == 0:
generate_video(
video_option=cfg.VIDEO_OPTION,
video_dir=cfg.VIDEO_DIR,
images=rgb_frames[i],
episode_id=current_episodes[i].episode_id,
checkpoint_idx=checkpoint_index,
metric_name=self.metric_uuid,
metric_value=infos[i][self.metric_uuid],
tb_writer=writer,
)
rgb_frames[i] = []
# episode continues
elif len(cfg.VIDEO_OPTION) > 0:
for k, v in observations[i].items():
if isinstance(v, torch.Tensor):
observations[i][k] = v.cpu().numpy()
frame = observations_to_image(observations[i], infos[i])
rgb_frames[i].append(frame)
# Pop done envs:
if len(envs_to_pause) > 0:
s_index = list(range(num_envs))
for idx in reversed(envs_to_pause):
s_index.pop(idx)
for k, v in other_losses.items():
other_losses[k] = other_losses[k][s_index]
for k, v in log_data.items():
log_data[k] = log_data[k][s_index]
(
self.envs,
test_recurrent_hidden_states,
not_done_masks,
current_episode_reward,
prev_actions,
batch,
rgb_frames,
) = self._pause_envs(
envs_to_pause,
self.envs,
test_recurrent_hidden_states,
not_done_masks,
None,
prev_actions,
batch,
rgb_frames,
)
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)
episodes_agg_stats = dict()
for k, v in aggregated_stats.items():
episodes_agg_stats[k] = v / num_episodes
logger.info(f"Average episode {k}: {episodes_agg_stats[k]:.6f}")
for k, v in episodes_agg_stats.items():
writer.add_scalars(f"eval_{k}", {f"{split}_average {k}": v},
checkpoint_index)
print(f"[{checkpoint_index}] average {k}", v)
self.envs.close()