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()