def train(self) -> None:
r"""Main method for DD-PPO.
Returns:
None
"""
self.local_rank, tcp_store = init_distrib_slurm(
self.config.RL.DDPPO.distrib_backend)
add_signal_handlers()
# Stores the number of workers that have finished their rollout
num_rollouts_done_store = distrib.PrefixStore("rollout_tracker",
tcp_store)
num_rollouts_done_store.set("num_done", "0")
self.world_rank = distrib.get_rank()
self.world_size = distrib.get_world_size()
random.seed(self.config.TASK_CONFIG.SEED + self.world_rank)
np.random.seed(self.config.TASK_CONFIG.SEED + self.world_rank)
self.config.defrost()
self.config.TORCH_GPU_ID = self.local_rank
self.config.SIMULATOR_GPU_ID = self.local_rank
self.config.freeze()
if torch.cuda.is_available():
self.device = torch.device("cuda", self.local_rank)
torch.cuda.set_device(self.device)
else:
self.device = torch.device("cpu")
self.envs = construct_envs(self.config,
get_env_class(self.config.ENV_NAME))
ppo_cfg = self.config.RL.PPO
if (not os.path.isdir(self.config.CHECKPOINT_FOLDER)
and self.world_rank == 0):
os.makedirs(self.config.CHECKPOINT_FOLDER)
self._setup_actor_critic_agent(ppo_cfg)
self.agent.init_distributed(find_unused_params=True)
if self.world_rank == 0:
logger.info("agent number of trainable parameters: {}".format(
sum(param.numel() for param in self.agent.parameters()
if param.requires_grad)))
observations = self.envs.reset()
batch = batch_obs(observations, device=self.device)
obs_space = self.envs.observation_spaces[0]
if self._static_encoder:
self._encoder = self.actor_critic.net.embedd
obs_space = SpaceDict({
"visual_features":
spaces.Box(
low=np.finfo(np.float32).min,
high=np.finfo(np.float32).max,
shape=self.actor_critic.net.embedding_shape,
dtype=np.float32,
),
**obs_space.spaces,
})
with torch.no_grad():
batch["visual_features"] = self._encoder(batch)
memory_info = {
'embedding_size': self.config.memory.embedding_size,
'memory_size': self.config.memory.memory_size,
'pose_size': self.config.memory.pose_dim
}
rollouts = MemoryRolloutStorage(
ppo_cfg.num_steps,
self.envs.num_envs,
obs_space,
self.envs.action_spaces[0],
ppo_cfg.hidden_size,
memory_info,
num_recurrent_layers=self.actor_critic.net.num_recurrent_layers,
)
rollouts.to(self.device)
for sensor in rollouts.observations:
rollouts.observations[sensor][0].copy_(batch[sensor])
rollouts.pre_embedding[0, :, 0].copy_(batch['visual_features'])
rollouts.memory_masks[0, :, 0].copy_(
torch.ones_like(rollouts.memory_masks[0, :, 0]))
# batch and observations may contain shared PyTorch CUDA
# tensors. We must explicitly clear them here otherwise
# they will be kept in memory for the entire duration of training!
batch = None
observations = None
current_episode_reward = torch.zeros(self.envs.num_envs,
1,
device=self.device)
running_episode_stats = dict(count=torch.zeros(self.envs.num_envs,
1,
device=self.device),
reward=torch.zeros(self.envs.num_envs,
1,
device=self.device),
episode_num=torch.zeros(
self.envs.num_envs,
1,
device=self.device))
window_episode_stats = defaultdict(
lambda: deque(maxlen=ppo_cfg.reward_window_size))
t_start = time.time()
env_time = 0
pth_time = 0
count_steps = 0 if not hasattr(self,
'resume_steps') else self.resume_steps
count_checkpoints = 0
start_update = 0
prev_time = 0
lr_scheduler = LambdaLR(
optimizer=self.agent.optimizer,
lr_lambda=lambda x: linear_decay(x, self.config.NUM_UPDATES),
)
interrupted_state = load_interrupted_state()
if interrupted_state is not None:
self.agent.load_state_dict(interrupted_state["state_dict"])
self.agent.optimizer.load_state_dict(
interrupted_state["optim_state"])
lr_scheduler.load_state_dict(interrupted_state["lr_sched_state"])
requeue_stats = interrupted_state["requeue_stats"]
env_time = requeue_stats["env_time"]
pth_time = requeue_stats["pth_time"]
count_steps = requeue_stats["count_steps"]
count_checkpoints = requeue_stats["count_checkpoints"]
start_update = requeue_stats["start_update"]
prev_time = requeue_stats["prev_time"]
with (TensorboardWriter(self.config.TENSORBOARD_DIR,
flush_secs=self.flush_secs)
if self.world_rank == 0 else contextlib.suppress()) as writer:
for update in range(start_update, self.config.NUM_UPDATES):
if ppo_cfg.use_linear_lr_decay:
lr_scheduler.step()
if ppo_cfg.use_linear_clip_decay:
self.agent.clip_param = ppo_cfg.clip_param * linear_decay(
update, self.config.NUM_UPDATES)
if EXIT.is_set():
self.envs.close()
if REQUEUE.is_set() and self.world_rank == 0:
requeue_stats = dict(
env_time=env_time,
pth_time=pth_time,
count_steps=count_steps,
count_checkpoints=count_checkpoints,
start_update=update,
prev_time=(time.time() - t_start) + prev_time,
)
save_interrupted_state(
dict(
state_dict=self.agent.state_dict(),
optim_state=self.agent.optimizer.state_dict(),
lr_sched_state=lr_scheduler.state_dict(),
config=self.config,
requeue_stats=requeue_stats,
))
requeue_job()
return
count_steps_delta = 0
self.agent.eval()
debug_prev_time = time.time()
for step in range(ppo_cfg.num_steps):
if TIME_DEBUG: prevv_time = time.time()
(
delta_pth_time,
delta_env_time,
delta_steps,
) = self._collect_rollout_step(rollouts,
current_episode_reward,
running_episode_stats)
pth_time += delta_pth_time
env_time += delta_env_time
count_steps_delta += delta_steps
# This is where the preemption of workers happens. If a
# worker detects it will be a straggler, it preempts itself!
if (step >=
ppo_cfg.num_steps * self.SHORT_ROLLOUT_THRESHOLD
) and int(num_rollouts_done_store.get("num_done")) > (
self.config.RL.DDPPO.sync_frac * self.world_size):
break
if TIME_DEBUG:
prevv_time = log_time(prevv_time, 'each step done')
num_rollouts_done_store.add("num_done", 1)
if TIME_DEBUG:
debug_prev_time = log_time(debug_prev_time,
'collect rollout step done')
self.agent.train()
if self._static_encoder:
self.actor_critic.net.eval_embed_network()
(
delta_pth_time,
value_loss,
action_loss,
dist_entropy,
) = self._update_agent(ppo_cfg, rollouts)
pth_time += delta_pth_time
if TIME_DEBUG:
debug_prev_time = log_time(debug_prev_time, 'update done')
stats_ordering = list(sorted(running_episode_stats.keys()))
stats = torch.stack(
[running_episode_stats[k] for k in stats_ordering], 0)
distrib.all_reduce(stats)
for i, k in enumerate(stats_ordering):
window_episode_stats[k].append(stats[i].clone())
stats = torch.tensor(
[value_loss, action_loss, count_steps_delta],
device=self.device,
)
distrib.all_reduce(stats)
count_steps += stats[2].item()
if self.world_rank == 0:
num_rollouts_done_store.set("num_done", "0")
losses = [
stats[0].item() / self.world_size,
stats[1].item() / self.world_size,
]
deltas = {
k: ((v[-1] - v[0]).sum().item()
if len(v) > 1 else v[0].sum().item())
for k, v in window_episode_stats.items()
}
deltas["count"] = max(deltas["count"], 1.0)
writer.add_scalar(
"reward",
deltas["reward"] / deltas["count"],
count_steps,
)
# Check to see if there are any metrics
# that haven't been logged yet
metrics = {
k: v / deltas["count"]
for k, v in deltas.items()
if k not in {"reward", "count", "lengths"}
}
if len(metrics) > 0:
writer.add_scalars("metrics", metrics, count_steps)
writer.add_scalars(
"losses",
{k: l
for l, k in zip(losses, ["value", "policy"])},
count_steps,
)
writer.add_scalars(
"metrics",
{'length': deltas['length'] / deltas['episode_num']},
count_steps)
# log stats
if update > 0 and update % self.config.LOG_INTERVAL == 0:
logger.info("update: {}\tfps: {:.3f}\t".format(
update,
count_steps /
((time.time() - t_start) + prev_time),
))
logger.info(
"update: {}\tenv-time: {:.3f}s\tpth-time: {:.3f}s\t"
"frames: {}".format(update, env_time, pth_time,
count_steps))
logger.info("Average window size: {} {}".format(
len(window_episode_stats["count"]),
" ".join(
"{}: {:.3f}".format(k, v / deltas["count"])
for k, v in deltas.items() if k != "count"),
))
# checkpoint model
if update % self.config.CHECKPOINT_INTERVAL == 0:
self.save_checkpoint(
f"ckpt.{count_checkpoints}.pth",
dict(step=count_steps),
)
count_checkpoints += 1
self.envs.close()
from env_utils.make_env_utils import construct_envs
from IL_configs.default import get_config
import numpy as np
import os
import time
os.environ['CUDA_VISIBLE_DEVICES'] = "0"
config = get_config('IL_configs/lgmt.yaml')
config.defrost()
config.DIFFICULTY = 'hard'
config.TASK_CONFIG.ENVIRONMENT.ITERATOR_OPTIONS.MAX_SCENE_REPEAT_EPISODES = 10
config.NUM_PROCESSES = 1
config.NUM_VAL_PROCESSES = 0
config.freeze()
action_list = config.TASK_CONFIG.TASK.POSSIBLE_ACTIONS
env = construct_envs(config, eval(config.ENV_NAME), mode='single')
obs = env.reset()
img = env.render('rgb')
done = False
fps = {}
reset_time = {}
scene = env.env.current_episode.scene_id.split('/')[-2]
fps[scene] = []
reset_time[scene] = []
imgs = [img]
while True:
action = env.env.get_best_action()
#action = env.action_space.sample()
#action = action_list.index(action['action'])
img = env.render('rgb')
def benchmark(self) -> None:
if TIME_DEBUG: s = time.time()
#self.config.defrost()
#self.config.TASK_CONFIG.ENVIRONMENT.ITERATOR_OPTIONS.MAX_SCENE_REPEAT_EPISODES = 10
#self.config.freeze()
if torch.cuda.device_count() <= 1:
self.config.defrost()
self.config.TORCH_GPU_ID = 0
self.config.SIMULATOR_GPU_ID = 0
self.config.freeze()
self.envs = construct_envs(self.config, eval(self.config.ENV_NAME))
if ADD_IL:
self.il_envs = construct_envs(self.config,
eval(self.config.ENV_NAME),
no_val=True)
self.collect_mode = 'RL'
if TIME_DEBUG: s = log_time(s, 'construct envs')
ppo_cfg = self.config.RL.PPO
self.device = (torch.device("cuda", self.config.TORCH_GPU_ID)
if torch.cuda.is_available() else torch.device("cpu"))
if not os.path.isdir(self.config.CHECKPOINT_FOLDER):
os.makedirs(self.config.CHECKPOINT_FOLDER)
self._setup_actor_critic_agent(ppo_cfg)
# if 'SMT' in self.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)
logger.info("agent number of parameters: {}".format(
sum(param.numel() for param in self.agent.parameters())))
num_train_processes, num_val_processes = self.config.NUM_PROCESSES, self.config.NUM_VAL_PROCESSES
total_processes = num_train_processes + num_val_processes
OBS_LIST = self.config.OBS_TO_SAVE
self.num_processes = num_train_processes
rollouts = RolloutStorage(ppo_cfg.num_steps,
num_train_processes,
self.envs.observation_spaces[0],
self.envs.action_spaces[0],
ppo_cfg.hidden_size,
self.actor_critic.net.num_recurrent_layers,
OBS_LIST=OBS_LIST)
rollouts.to(self.device)
batch = self.envs.reset()
for sensor in rollouts.observations:
rollouts.observations[sensor][0].copy_(
batch[sensor][:num_train_processes])
self.last_observations = batch
self.last_recurrent_hidden_states = torch.zeros(
self.actor_critic.net.num_recurrent_layers, total_processes,
ppo_cfg.hidden_size).to(self.device)
self.last_prev_actions = torch.zeros(
total_processes, rollouts.prev_actions.shape[-1]).to(self.device)
self.last_masks = torch.zeros(total_processes, 1).to(self.device)
# batch and observations may contain shared PyTorch CUDA
# tensors. We must explicitly clear them here otherwise
# they will be kept in memory for the entire duration of training!
batch = None
observations = None
if ADD_IL:
rollouts2 = RolloutStorage(
ppo_cfg.num_steps,
num_train_processes,
self.il_envs.observation_spaces[0],
self.il_envs.action_spaces[0],
ppo_cfg.hidden_size,
self.actor_critic.net.num_recurrent_layers,
OBS_LIST=OBS_LIST)
rollouts2.to(self.device)
batch2 = self.il_envs.reset()
for sensor in rollouts2.observations:
rollouts2.observations[sensor][0].copy_(
batch2[sensor][:num_train_processes])
self.saved_last_obs = batch2
self.saved_last_recurrent_hidden_states = torch.zeros(
self.actor_critic.net.num_recurrent_layers, total_processes,
ppo_cfg.hidden_size).to(self.device)
self.saved_last_prev_actions = torch.zeros(
total_processes,
rollouts2.prev_actions.shape[-1]).to(self.device)
self.saved_last_masks = torch.zeros(total_processes,
1).to(self.device)
batch2 = None
else:
rollouts2 = None
current_episode_reward = torch.zeros(self.envs.num_envs, 1)
running_episode_stats = dict(
count=torch.zeros(self.envs.num_envs, 1),
reward=torch.zeros(self.envs.num_envs, 1),
)
window_episode_stats = defaultdict(
lambda: deque(maxlen=ppo_cfg.reward_window_size))
t_start = time.time()
env_time = 0
pth_time = 0
count_steps = 0 if not hasattr(self,
'resume_steps') else self.resume_steps
count_checkpoints = 0
lr_scheduler = LambdaLR(
optimizer=self.agent.optimizer,
lr_lambda=lambda x: linear_decay(x, self.config.NUM_UPDATES),
)
if TIME_DEBUG: s = log_time(s, 'setup all')
for update in range(100):
if ppo_cfg.use_linear_lr_decay:
lr_scheduler.step()
if ppo_cfg.use_linear_clip_decay:
self.agent.clip_param = ppo_cfg.clip_param * linear_decay(
update, self.config.NUM_UPDATES)
if TIME_DEBUG: s = log_time(s, 'collect rollout start')
modes = ['RL']
if ADD_IL:
modes += ['IL']
for collect_mode in modes:
self.collect_mode = collect_mode
use_rollouts = rollouts if self.collect_mode == 'RL' else rollouts2
if ADD_IL: self.exchange_lasts()
for step in range(ppo_cfg.num_steps):
(
delta_pth_time,
delta_env_time,
delta_steps,
) = self._collect_rollout_step(use_rollouts,
current_episode_reward,
running_episode_stats)
pth_time += delta_pth_time
env_time += delta_env_time
count_steps += delta_steps
#print(delta_env_time, delta_pth_time)
if TIME_DEBUG: s = log_time(s, 'collect rollout done')
(delta_pth_time, value_loss, action_loss, dist_entropy,
il_loss) = self._update_agent(ppo_cfg, rollouts, rollouts2)
#print(delta_pth_time)
pth_time += delta_pth_time
use_rollouts.after_update()
if TIME_DEBUG: s = log_time(s, 'update agent')
for k, v in running_episode_stats.items():
window_episode_stats[k].append(v.clone())
deltas = {
k: ((v[-1][:self.num_processes] -
v[0][:self.num_processes]).sum().item()
if len(v) > 1 else v[0][:self.num_processes].sum().item())
for k, v in window_episode_stats.items()
}
deltas["count"] = max(deltas["count"], 1.0)
#self.write_tb('train', writer, deltas, count_steps, losses)
eval_deltas = {
k: ((v[-1][self.num_processes:] -
v[0][self.num_processes:]).sum().item()
if len(v) > 1 else v[0][self.num_processes:].sum().item())
for k, v in window_episode_stats.items()
}
eval_deltas["count"] = max(eval_deltas["count"], 1.0)
#self.write_tb('val', writer, eval_deltas, count_steps)
# log stats
if update > 0 and update % self.config.LOG_INTERVAL == 0:
logger.info("update: {}\tfps: {:.3f}\t".format(
update, count_steps / (time.time() - t_start)))
logger.info(
"update: {}\tenv-time: {:.3f}s\tpth-time: {:.3f}s\t"
"frames: {}".format(update, env_time, pth_time,
count_steps))
logger.info("Average window size: {} {}".format(
len(window_episode_stats["count"]),
" ".join("{}: {:.3f}".format(k, v / deltas["count"])
for k, v in deltas.items() if k != "count"),
))
logger.info("validation metrics: {}".format(
" ".join("{}: {:.3f}".format(k, v / eval_deltas["count"])
for k, v in eval_deltas.items()
if k != "count"), ))
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
"""
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
"""
# 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.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)
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)
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):
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, 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):
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 train(self) -> None:
r"""Main method for training PPO.
Returns:
None
"""
self.envs = construct_envs(self.config,
get_env_class(self.config.ENV_NAME))
ppo_cfg = self.config.RL.PPO
self.device = (torch.device("cuda", self.config.TORCH_GPU_ID)
if torch.cuda.is_available() else torch.device("cpu"))
if not os.path.isdir(self.config.CHECKPOINT_FOLDER):
os.makedirs(self.config.CHECKPOINT_FOLDER)
self._setup_actor_critic_agent(ppo_cfg)
logger.info("agent number of parameters: {}".format(
sum(param.numel() for param in self.agent.parameters())))
num_train_processes, num_val_processes = self.config.NUM_PROCESSES, self.config.NUM_VAL_PROCESSES
total_processes = num_train_processes + num_val_processes
self.num_processes = num_train_processes
rollouts = RolloutStorage(ppo_cfg.num_steps, num_train_processes,
self.envs.observation_spaces[0],
self.envs.action_spaces[0],
ppo_cfg.hidden_size,
self.actor_critic.net.num_recurrent_layers)
rollouts.to(self.device)
observations = self.envs.reset()
batch = batch_obs(observations, device=self.device)
for sensor in rollouts.observations:
rollouts.observations[sensor][0].copy_(
batch[sensor][:num_train_processes])
self.last_observations = batch
self.last_recurrent_hidden_states = torch.zeros(
self.actor_critic.net.num_recurrent_layers, total_processes,
ppo_cfg.hidden_size).to(self.device)
self.last_prev_actions = torch.zeros(
total_processes, rollouts.prev_actions.shape[-1]).to(self.device)
self.last_masks = torch.zeros(total_processes, 1).to(self.device)
# batch and observations may contain shared PyTorch CUDA
# tensors. We must explicitly clear them here otherwise
# they will be kept in memory for the entire duration of training!
batch = None
observations = None
current_episode_reward = torch.zeros(self.envs.num_envs, 1)
running_episode_stats = dict(
count=torch.zeros(self.envs.num_envs, 1),
reward=torch.zeros(self.envs.num_envs, 1),
)
window_episode_stats = defaultdict(
lambda: deque(maxlen=ppo_cfg.reward_window_size))
t_start = time.time()
env_time = 0
pth_time = 0
count_steps = 0
count_checkpoints = 0
lr_scheduler = LambdaLR(
optimizer=self.agent.optimizer,
lr_lambda=lambda x: linear_decay(x, self.config.NUM_UPDATES),
)
with TensorboardWriter(self.config.TENSORBOARD_DIR,
flush_secs=self.flush_secs) as writer:
for update in range(self.config.NUM_UPDATES):
if ppo_cfg.use_linear_lr_decay:
lr_scheduler.step()
if ppo_cfg.use_linear_clip_decay:
self.agent.clip_param = ppo_cfg.clip_param * linear_decay(
update, self.config.NUM_UPDATES)
for step in range(ppo_cfg.num_steps):
(
delta_pth_time,
delta_env_time,
delta_steps,
) = self._collect_rollout_step(rollouts,
current_episode_reward,
running_episode_stats)
pth_time += delta_pth_time
env_time += delta_env_time
count_steps += delta_steps
(
delta_pth_time,
value_loss,
action_loss,
dist_entropy,
) = self._update_agent(ppo_cfg, rollouts)
pth_time += delta_pth_time
for k, v in running_episode_stats.items():
window_episode_stats[k].append(v.clone())
deltas = {
k: ((v[-1][:self.num_processes] -
v[0][:self.num_processes]).sum().item() if len(v) > 1
else v[0][:self.num_processes].sum().item())
for k, v in window_episode_stats.items()
}
deltas["count"] = max(deltas["count"], 1.0)
losses = [value_loss, action_loss]
self.write_tb('train', writer, deltas, count_steps, losses)
eval_deltas = {
k: ((v[-1][self.num_processes:] -
v[0][self.num_processes:]).sum().item() if len(v) > 1
else v[0][self.num_processes:].sum().item())
for k, v in window_episode_stats.items()
}
eval_deltas["count"] = max(eval_deltas["count"], 1.0)
self.write_tb('val', writer, eval_deltas, count_steps)
# log stats
if update > 0 and update % self.config.LOG_INTERVAL == 0:
logger.info("update: {}\tfps: {:.3f}\t".format(
update, count_steps / (time.time() - t_start)))
logger.info(
"update: {}\tenv-time: {:.3f}s\tpth-time: {:.3f}s\t"
"frames: {}".format(update, env_time, pth_time,
count_steps))
logger.info("Average window size: {} {}".format(
len(window_episode_stats["count"]),
" ".join("{}: {:.3f}".format(k, v / deltas["count"])
for k, v in deltas.items() if k != "count"),
))
logger.info("validation metrics: {}".format(
" ".join("{}: {:.3f}".format(k, v /
eval_deltas["count"])
for k, v in eval_deltas.items()
if k != "count"), ))
# checkpoint model
if update % self.config.CHECKPOINT_INTERVAL == 0:
self.save_checkpoint(f"ckpt.{count_checkpoints}.pth",
dict(step=count_steps))
count_checkpoints += 1
self.envs.close()