base_runner.py
训练器/buffer 初始化及公共流程
原文件 · 203 行 · 8644 字节
runner/shared/base_runner.py
SHA-256: 089753fc171381140f5c2d34b98becd9072a8686e8f804eb928199c7d91c8be7
1import os2import numpy as np3import torch4from tensorboardX import SummaryWriter5from envs.macf.env_wrappers import DummyVecEnv6from utils.obs_utils import flat_obs_space, flat_state_space78def _t2n(x):9 """Convert torch tensor to a numpy array."""10 return x.detach().cpu().numpy()111213class Runner(object):14 """15 Base class for training recurrent policies.16 :param config: (dict) Config dictionary containing parameters for training.17 """18 def __init__(self, config):1920 self.config = config21 self.use_gat = config.use_gat22 self.envs: DummyVecEnv = config.envs23 self.eval_envs = config.eval_envs24 self.device = config.device25 self.num_agents = config.num_agents26 # if config.__contains__("render_envs"):27 # self.render_envs = config['render_envs'] 2829 # parameters30 self.env_name = config.env_name31 self.algorithm_name = config.algorithm_name32 self.experiment_name = config.experiment_name33 self.use_centralized_V = config.use_centralized_V34 self.use_obs_instead_of_state = config.use_obs_instead_of_state35 self.num_env_steps = config.num_env_steps36 self.episode_length = config.episode_length37 self.n_rollout_threads = config.n_rollout_threads38 self.n_eval_rollout_threads = config.n_eval_rollout_threads39 # self.n_render_rollout_threads = config.n_render_rollout_threads40 self.use_linear_lr_decay = config.use_linear_lr_decay41 self.hidden_size = config.hidden_size42 self.use_render = config.use_render43 self.recurrent_N = config.recurrent_N4445 # interval46 self.save_interval = config.save_interval47 self.use_eval = config.use_eval48 self.eval_interval = config.eval_interval49 self.eval_episodes = config.eval_episodes50 self.eval_seed = config.eval_seed51 self.log_interval = config.log_interval52 self.log_info_interval = config.log_info_interval5354 # dir55 self.model_dir = config.model_dir5657 self.run_dir = config.run_dir58 self.log_dir = str(getattr(config, "eval_output_dir", self.run_dir))59 self.tb_dir = str(getattr(config, "tb_dir", self.run_dir / 'tb'))60 if not os.path.exists(self.log_dir):61 os.makedirs(self.log_dir)62 if not os.path.exists(self.tb_dir):63 os.makedirs(self.tb_dir)64 self.writter = SummaryWriter(self.tb_dir)65 self.save_dir = str(self.run_dir / 'models')66 if not os.path.exists(self.save_dir):67 os.makedirs(self.save_dir)6869 if self.use_gat:70 from algorithms.algorithm.ga_rmappo import GARMAPPO as TrainAlgo71 from algorithms.algorithm.gaRMAPPOPolicy import GAMAPPOPolicy as Policy72 from utils.shared_gat_buffer import SharedGatReplayBuffer as Buffer73 else:74 from algorithms.algorithm.r_mappo import RMAPPO as TrainAlgo75 from algorithms.algorithm.rMAPPOPolicy import RMAPPOPolicy as Policy76 from utils.shared_buffer import SharedReplayBuffer as Buffer7778 obs_space = self.envs.observation_space[0] # Dict space79 share_observation_space = self.envs.share_observation_space[0] if self.use_centralized_V else self.envs.observation_space[0]8081 # Spaces provide dimensions; HATT keeps Dict observations until its actor,82 # while standard MAPPO and legacy GAT flatten/fuse them inside the policy.83 obs_box = flat_obs_space(obs_space)84 share_box = flat_state_space(share_observation_space) if hasattr(share_observation_space, 'spaces') else share_observation_space85 self.policy = Policy(config,86 obs_box,87 share_box,88 self.envs.action_space[0],89 device = self.device)9091 if self.model_dir is not None:92 self.restore()9394 # algorithm95 self.trainer = TrainAlgo(config, self.policy, device = self.device)9697 # buffer — 传入 Dict space，buffer 内部处理 dict obs98 self.buffer = Buffer(config,99 self.num_agents,100 obs_space,101 share_observation_space,102 self.envs.action_space[0])103104 def run(self):105 """Collect training data, perform training updates, and evaluate policy."""106 raise NotImplementedError107108 def warmup(self):109 """Collect warmup pre-training data."""110 raise NotImplementedError111112 def collect(self, step):113 """Collect rollouts for training."""114 raise NotImplementedError115116 def insert(self, data):117 """118 Insert data into buffer.119 :param data: (Tuple) data to insert into training buffer.120 """121 raise NotImplementedError122 123 @torch.no_grad()124 def compute(self):125 """Calculate returns for the collected data."""126 self.trainer.prep_rollout()127 next_share_obs = {128 k: np.concatenate(v[-1]) for k, v in self.buffer.share_obs.items()129 } if isinstance(self.buffer.share_obs, dict) else np.concatenate(self.buffer.share_obs[-1])130 next_values = self.trainer.policy.get_values(131 next_share_obs,132 np.concatenate(self.buffer.rnn_states_critic[-1]),133 np.concatenate(self.buffer.masks[-1]),134 )135 next_values = np.array(np.split(_t2n(next_values), self.n_rollout_threads))136 self.buffer.compute_returns(next_values, self.trainer.value_normalizer)137 138 def train(self):139 """Train policies with data in buffer. """140 self.trainer.prep_training()141 train_infos = self.trainer.train(self.buffer) 142 self.buffer.after_update()143 return train_infos144145 def save(self, tag=None):146 """Save policy's actor and critic networks.147148 Args:149 tag: optional suffix appended to the filename (e.g. ``"ep000500"``150 produces ``actor_ep000500.pt``). When *None* (the default)151 the canonical ``actor.pt`` / ``critic.pt`` files are written.152 """153 actor_suffix = f"_{tag}" if tag else ""154 critic_suffix = f"_{tag}" if tag else ""155 actor_name = f"actor{actor_suffix}.pt"156 critic_name = f"critic{critic_suffix}.pt"157 valuenorm_name = f"valuenorm{critic_suffix}.pt"158159 if hasattr(self.trainer.policy, "get_actor_checkpoint"):160 actor_state = self.trainer.policy.get_actor_checkpoint()161 else:162 actor_state = self.trainer.policy.actor.state_dict()163 torch.save(actor_state, os.path.join(str(self.save_dir), actor_name))164 policy_critic = self.trainer.policy.critic165 torch.save(policy_critic.state_dict(), os.path.join(str(self.save_dir), critic_name))166 value_normalizer = self.trainer.value_normalizer167 if value_normalizer is not None:168 torch.save(169 value_normalizer.state_dict(),170 os.path.join(str(self.save_dir), valuenorm_name),171 )172173 def restore(self):174 """Restore policy's networks from a saved model."""175 policy_actor_state_dict = torch.load(176 str(self.model_dir) + '/actor.pt', map_location=self.device177 )178 if hasattr(self.policy, "load_actor_checkpoint"):179 self.policy.load_actor_checkpoint(policy_actor_state_dict)180 else:181 self.policy.actor.load_state_dict(policy_actor_state_dict)182 if not self.config.use_render:183 policy_critic_state_dict = torch.load(str(self.model_dir) + '/critic.pt')184 self.policy.critic.load_state_dict(policy_critic_state_dict)185 186 def log_train(self, train_infos, total_num_steps):187 """188 Log training info.189 :param train_infos: (dict) information about training update.190 :param total_num_steps: (int) total number of training env steps.191 """192 for k, v in train_infos.items():193 self.writter.add_scalar(k, v, total_num_steps)194195 def log_env(self, env_infos, total_num_steps):196 """197 Log env info.198 :param env_infos: (dict) information about env state.199 :param total_num_steps: (int) total number of training env steps.200 """201 for k, v in env_infos.items():202 if len(v)>0:203 self.writter.add_scalar(k, np.mean(v), total_num_steps)