import os import numpy as np import torch from tensorboardX import SummaryWriter from envs.macf.env_wrappers import DummyVecEnv from utils.obs_utils import flat_obs_space, flat_state_space def _t2n(x): """Convert torch tensor to a numpy array.""" return x.detach().cpu().numpy() class Runner(object): """ Base class for training recurrent policies. :param config: (dict) Config dictionary containing parameters for training. """ def __init__(self, config): self.config = config self.use_gat = config.use_gat self.envs: DummyVecEnv = config.envs self.eval_envs = config.eval_envs self.device = config.device self.num_agents = config.num_agents # if config.__contains__("render_envs"): # self.render_envs = config['render_envs'] # parameters self.env_name = config.env_name self.algorithm_name = config.algorithm_name self.experiment_name = config.experiment_name self.use_centralized_V = config.use_centralized_V self.use_obs_instead_of_state = config.use_obs_instead_of_state self.num_env_steps = config.num_env_steps self.episode_length = config.episode_length self.n_rollout_threads = config.n_rollout_threads self.n_eval_rollout_threads = config.n_eval_rollout_threads # self.n_render_rollout_threads = config.n_render_rollout_threads self.use_linear_lr_decay = config.use_linear_lr_decay self.hidden_size = config.hidden_size self.use_render = config.use_render self.recurrent_N = config.recurrent_N # interval self.save_interval = config.save_interval self.use_eval = config.use_eval self.eval_interval = config.eval_interval self.eval_episodes = config.eval_episodes self.eval_seed = config.eval_seed self.log_interval = config.log_interval self.log_info_interval = config.log_info_interval # dir self.model_dir = config.model_dir self.run_dir = config.run_dir self.log_dir = str(getattr(config, "eval_output_dir", self.run_dir)) self.tb_dir = str(getattr(config, "tb_dir", self.run_dir / 'tb')) if not os.path.exists(self.log_dir): os.makedirs(self.log_dir) if not os.path.exists(self.tb_dir): os.makedirs(self.tb_dir) self.writter = SummaryWriter(self.tb_dir) self.save_dir = str(self.run_dir / 'models') if not os.path.exists(self.save_dir): os.makedirs(self.save_dir) if self.use_gat: from algorithms.algorithm.ga_rmappo import GARMAPPO as TrainAlgo from algorithms.algorithm.gaRMAPPOPolicy import GAMAPPOPolicy as Policy from utils.shared_gat_buffer import SharedGatReplayBuffer as Buffer else: from algorithms.algorithm.r_mappo import RMAPPO as TrainAlgo from algorithms.algorithm.rMAPPOPolicy import RMAPPOPolicy as Policy from utils.shared_buffer import SharedReplayBuffer as Buffer obs_space = self.envs.observation_space[0] # Dict space share_observation_space = self.envs.share_observation_space[0] if self.use_centralized_V else self.envs.observation_space[0] # Spaces provide dimensions; HATT keeps Dict observations until its actor, # while standard MAPPO and legacy GAT flatten/fuse them inside the policy. obs_box = flat_obs_space(obs_space) share_box = flat_state_space(share_observation_space) if hasattr(share_observation_space, 'spaces') else share_observation_space self.policy = Policy(config, obs_box, share_box, self.envs.action_space[0], device = self.device) if self.model_dir is not None: self.restore() # algorithm self.trainer = TrainAlgo(config, self.policy, device = self.device) # buffer — 传入 Dict space，buffer 内部处理 dict obs self.buffer = Buffer(config, self.num_agents, obs_space, share_observation_space, self.envs.action_space[0]) def run(self): """Collect training data, perform training updates, and evaluate policy.""" raise NotImplementedError def warmup(self): """Collect warmup pre-training data.""" raise NotImplementedError def collect(self, step): """Collect rollouts for training.""" raise NotImplementedError def insert(self, data): """ Insert data into buffer. :param data: (Tuple) data to insert into training buffer. """ raise NotImplementedError @torch.no_grad() def compute(self): """Calculate returns for the collected data.""" self.trainer.prep_rollout() next_share_obs = { k: np.concatenate(v[-1]) for k, v in self.buffer.share_obs.items() } if isinstance(self.buffer.share_obs, dict) else np.concatenate(self.buffer.share_obs[-1]) next_values = self.trainer.policy.get_values( next_share_obs, np.concatenate(self.buffer.rnn_states_critic[-1]), np.concatenate(self.buffer.masks[-1]), ) next_values = np.array(np.split(_t2n(next_values), self.n_rollout_threads)) self.buffer.compute_returns(next_values, self.trainer.value_normalizer) def train(self): """Train policies with data in buffer. """ self.trainer.prep_training() train_infos = self.trainer.train(self.buffer) self.buffer.after_update() return train_infos def save(self, tag=None): """Save policy's actor and critic networks. Args: tag: optional suffix appended to the filename (e.g. ``"ep000500"`` produces ``actor_ep000500.pt``). When *None* (the default) the canonical ``actor.pt`` / ``critic.pt`` files are written. """ actor_suffix = f"_{tag}" if tag else "" critic_suffix = f"_{tag}" if tag else "" actor_name = f"actor{actor_suffix}.pt" critic_name = f"critic{critic_suffix}.pt" valuenorm_name = f"valuenorm{critic_suffix}.pt" if hasattr(self.trainer.policy, "get_actor_checkpoint"): actor_state = self.trainer.policy.get_actor_checkpoint() else: actor_state = self.trainer.policy.actor.state_dict() torch.save(actor_state, os.path.join(str(self.save_dir), actor_name)) policy_critic = self.trainer.policy.critic torch.save(policy_critic.state_dict(), os.path.join(str(self.save_dir), critic_name)) value_normalizer = self.trainer.value_normalizer if value_normalizer is not None: torch.save( value_normalizer.state_dict(), os.path.join(str(self.save_dir), valuenorm_name), ) def restore(self): """Restore policy's networks from a saved model.""" policy_actor_state_dict = torch.load( str(self.model_dir) + '/actor.pt', map_location=self.device ) if hasattr(self.policy, "load_actor_checkpoint"): self.policy.load_actor_checkpoint(policy_actor_state_dict) else: self.policy.actor.load_state_dict(policy_actor_state_dict) if not self.config.use_render: policy_critic_state_dict = torch.load(str(self.model_dir) + '/critic.pt') self.policy.critic.load_state_dict(policy_critic_state_dict) def log_train(self, train_infos, total_num_steps): """ Log training info. :param train_infos: (dict) information about training update. :param total_num_steps: (int) total number of training env steps. """ for k, v in train_infos.items(): self.writter.add_scalar(k, v, total_num_steps) def log_env(self, env_infos, total_num_steps): """ Log env info. :param env_infos: (dict) information about env state. :param total_num_steps: (int) total number of training env steps. """ for k, v in env_infos.items(): if len(v)>0: self.writter.add_scalar(k, np.mean(v), total_num_steps)