macf_runner.py
值分解 rollout 与评估
原文件 · 398 行 · 18828 字节
runner/qmix/macf_runner.py
SHA-256: efd11997e57e38d9ba59938e8dc0c350e93894f1d1ab71464fbbb79c3b536f6f
1"""Current-environment runner for recurrent QMIX, VDN and HATT-QMIX."""23import os4import time5from pathlib import Path67import numpy as np8import torch9from tensorboardX import SummaryWriter1011from algorithms.qmix.algorithm.QMixPolicy import QMixPolicy12from algorithms.qmix.qmix import QMix13from runner.shared.macf_runner import MACFEnvRunner14from utils.obs_utils import flatten_obs_dict, flatten_state_dict15from utils.qmix_buffer import EpisodeReplayBuffer161718def _env_item(tree, index):19 if isinstance(tree, dict):20 return {key: np.asarray(value[index], dtype=np.float32).copy() for key, value in tree.items()}21 return np.asarray(tree[index], dtype=np.float32).copy()222324def _append_obs(storage, obs):25 if isinstance(obs, dict):26 for key, value in obs.items():27 storage.setdefault(key, []).append(value)28 else:29 storage.append(obs)303132def _stack_obs(storage):33 if isinstance(storage, dict):34 return {key: np.stack(values).astype(np.float32) for key, values in storage.items()}35 return np.stack(storage).astype(np.float32)363738class MACFRunner(MACFEnvRunner):39 """Off-policy runner aligned with ``DummyVecEnv`` and sweep checkpoints."""4041 def __init__(self, config):42 self.config = config43 self.envs = config.envs44 self.eval_envs = config.eval_envs45 self.device = config.device46 self.num_agents = int(config.num_agents)47 self.env_name = config.env_name48 self.algorithm_name = config.algorithm_name49 self.experiment_name = config.experiment_name50 self.num_env_steps = int(config.num_env_steps)51 self.episode_length = int(config.episode_length)52 self.n_rollout_threads = int(config.n_rollout_threads)53 self.n_eval_rollout_threads = int(config.n_eval_rollout_threads)54 self.recurrent_N = 155 self.hidden_size = int(config.qmix_agent_hidden_dim)56 self.use_eval = bool(config.use_eval)57 self.eval_interval = config.eval_interval58 self.eval_episodes = int(config.eval_episodes)59 self.eval_seed = int(config.eval_seed)60 self.save_interval = config.save_interval61 self.log_interval = int(config.log_interval)62 self.log_info_interval = int(config.log_info_interval)63 self.model_dir = config.model_dir64 self.run_dir = Path(config.run_dir)65 self.log_dir = str(getattr(config, "eval_output_dir", self.run_dir))66 self.tb_dir = str(getattr(config, "tb_dir", self.run_dir / "tb"))67 os.makedirs(self.log_dir, exist_ok=True)68 os.makedirs(self.tb_dir, exist_ok=True)69 self.writter = SummaryWriter(self.tb_dir)70 self.writer = self.writter71 self.save_dir = str(self.run_dir / "models")72 os.makedirs(self.save_dir, exist_ok=True)7374 obs_space = self.envs.observation_space[0]75 action_space = self.envs.action_space[0]76 # The PPO shared-state space also declares a per-focal one-hot field.77 # QMIX consumes the environment's single joint state directly.78 state_spaces = self.envs.share_observation_space[0].spaces79 state_dim = sum(80 int(np.prod(space.shape))81 for key, space in state_spaces.items()82 if key != "focal"83 )84 self.policy = QMixPolicy(85 config, obs_space, action_space, self.num_agents, self.device86 )87 self.trainer = QMix(88 config, self.policy, state_dim, self.num_agents, self.device89 )90 self.buffer = EpisodeReplayBuffer(config.qmix_buffer_size, config.seed)91 self.env_episode = 092 self.episode = 093 self.best_eval_score = None94 self.best_eval_summary = None95 self.last_save_env_episode = 096 self.last_eval_env_episode = 097 self.last_train_env_episode = 098 self.last_target_env_episode = 099 if self.model_dir:100 self.restore()101102 def _policy_obs(self, obs):103 if self.policy.use_hatt:104 return obs105 return flatten_obs_dict(obs)106107 def _new_episode(self, obs, state, available, active, env_index):108 initial_obs = _env_item(self._policy_obs(obs), env_index)109 obs_storage = {} if isinstance(initial_obs, dict) else []110 _append_obs(obs_storage, initial_obs)111 return {112 "obs": obs_storage,113 "states": [np.asarray(state[env_index], dtype=np.float32).copy()],114 "available_actions": [np.asarray(available[env_index], dtype=np.float32).copy()],115 "active_masks": [np.asarray(active[env_index], dtype=np.float32).copy()],116 "alive_masks": [np.ones((self.num_agents, 1), dtype=np.float32)],117 "actions": [], "rewards": [], "dones": [], "held_actions": [],118 "return": 0.0, "sent": 0.0, "sent_goal": 0.0,119 }120121 def _finish_episode(self, episode):122 packed = {123 "obs": _stack_obs(episode["obs"]),124 "states": np.stack(episode["states"]).astype(np.float32),125 "available_actions": np.stack(episode["available_actions"]).astype(np.float32),126 "actions": np.stack(episode["actions"]).astype(np.int64),127 "rewards": np.asarray(episode["rewards"], dtype=np.float32).reshape(-1, 1),128 "dones": np.asarray(episode["dones"], dtype=np.float32).reshape(-1, 1),129 "active_masks": np.stack(episode["active_masks"]).astype(np.float32),130 "alive_masks": np.stack(episode["alive_masks"]).astype(np.float32),131 "held_actions": np.stack(episode["held_actions"]).astype(np.int64),132 }133 packed = self._pack_transition_mode(packed)134 self.buffer.insert(packed)135136 def _pack_transition_mode(self, packed):137 """Attach physical discounts or compress an episode to decision events."""138 gamma = float(self.config.gamma)139 transitions = len(packed["actions"])140 if self.config.qmix_transition_mode != "event_smdp":141 packed["discounts"] = np.full(142 (transitions, 1), gamma, dtype=np.float32143 )144 return packed145146 starts = [147 step for step in range(transitions)148 if np.any(packed["active_masks"][step] > 0.5)149 ]150 if not starts:151 starts = [0]152 endpoints = starts[1:] + [transitions]153 obs = packed["obs"]154 event = {155 "obs": ({key: [value[starts[0]]] for key, value in obs.items()}156 if isinstance(obs, dict) else [obs[starts[0]]]),157 "states": [packed["states"][starts[0]]],158 "available_actions": [packed["available_actions"][starts[0]]],159 "active_masks": [packed["active_masks"][starts[0]]],160 "alive_masks": [packed["alive_masks"][starts[0]]],161 "actions": [], "rewards": [], "dones": [],162 "held_actions": [], "discounts": [],163 }164 for start, end in zip(starts, endpoints):165 duration = end - start166 weights = gamma ** np.arange(duration, dtype=np.float32)167 reward = float(np.sum(packed["rewards"][start:end, 0] * weights))168 event["actions"].append(packed["actions"][start])169 event["rewards"].append([reward])170 event["dones"].append(packed["dones"][end - 1])171 event["held_actions"].append(packed["held_actions"][end - 1])172 event["discounts"].append([gamma ** duration])173 if isinstance(obs, dict):174 for key, value in obs.items():175 event["obs"][key].append(value[end])176 else:177 event["obs"].append(obs[end])178 for key in (179 "states", "available_actions", "active_masks", "alive_masks"180 ):181 event[key].append(packed[key][end])182 for key, value in event.items():183 if key == "obs" and isinstance(value, dict):184 event[key] = {name: np.stack(items) for name, items in value.items()}185 elif key in ("actions", "held_actions"):186 event[key] = np.asarray(value, dtype=np.int64)187 else:188 event[key] = np.asarray(value, dtype=np.float32)189 return event190191 def _epsilon(self, total_steps):192 fraction = min(1.0, total_steps / max(1, int(self.config.qmix_epsilon_anneal_steps)))193 return float(self.config.qmix_epsilon_start + fraction * (194 self.config.qmix_epsilon_finish - self.config.qmix_epsilon_start195 ))196197 @staticmethod198 def _advance_rollout_hidden(next_hidden, agent_dones, active_masks=None,199 previous_hidden=None, event_mode=False):200 """Advance recurrent state on every physical transition.201202 Event-mode active masks control whether a newly proposed action is203 executed. They do not pause observation history: replay training and204 deterministic evaluation both feed every physical transition through205 the recurrent agent network. Keep collection aligned with those paths206 and reset only environments whose episodes ended.207 """208 hidden = np.asarray(next_hidden, dtype=np.float32).copy()209 if event_mode:210 update = np.any(211 np.asarray(active_masks) > 0.5, axis=tuple(range(1, np.asarray(active_masks).ndim))212 )[:, None, None]213 hidden = np.where(update, hidden, np.asarray(previous_hidden))214 hidden[np.asarray(agent_dones, dtype=bool)] = 0.0215 return hidden216217 @staticmethod218 def _completed_interval_due(current_episode, last_episode, interval):219 intervals = interval if isinstance(interval, (list, tuple)) else [interval]220 return current_episode > last_episode and any(221 int(value) > 0 and current_episode % int(value) == 0222 for value in intervals223 )224225 def run(self):226 obs, active_masks, available = self.envs.reset(0)227 state = flatten_state_dict(self.envs.get_global_state())228 episodes = [229 self._new_episode(obs, state, available, active_masks, i)230 for i in range(self.n_rollout_threads)231 ]232 hidden = np.zeros(233 (self.n_rollout_threads, self.num_agents, self.hidden_size), dtype=np.float32234 )235 last_actions = np.zeros((self.n_rollout_threads, self.num_agents, 1), dtype=np.int64)236 max_episodes = getattr(self.config, "max_train_episodes", None)237 total_steps = 0238 updates = 0239 last_info = {}240 last_episode_return = 0.0241 last_sent_percentage = 0.0242 last_log_episode = -1243 start = time.time()244 while total_steps < self.num_env_steps and (245 max_episodes is None or self.env_episode < int(max_episodes)246 ):247 epsilon = self._epsilon(total_steps)248 self.trainer.prep_rollout()249 with torch.no_grad():250 actions_t, next_hidden_t = self.policy.select_actions(251 obs, hidden, available, epsilon=epsilon252 )253 actions = actions_t.cpu().numpy().astype(np.int64)254 next_hidden = next_hidden_t.cpu().numpy()255 actions = active_masks.astype(np.int64) * actions + (256 1 - active_masks.astype(np.int64)257 ) * last_actions258 next_obs, rewards, dones, infos, next_active, _, next_available = self.envs.step(259 actions, self.env_episode260 )261 next_state = flatten_state_dict(self.envs.get_global_state())262 env_dones = np.all(dones, axis=1)263 completed_this_step = int(np.count_nonzero(env_dones))264 for index in range(self.n_rollout_threads):265 episode = episodes[index]266 episode["actions"].append(actions[index, :, 0].copy())267 episode["held_actions"].append(actions[index, :, 0].copy())268 team_reward = float(np.mean(rewards[index])) * float(self.config.qmix_reward_scale)269 episode["rewards"].append(team_reward)270 episode["dones"].append(float(env_dones[index]))271 episode["return"] += team_reward272 episode["sent"] += float(infos[index].get("sent", 0.0))273 episode["sent_goal"] = float(infos[index].get("sent_goal", 0.0))274 _append_obs(episode["obs"], _env_item(self._policy_obs(next_obs), index))275 episode["states"].append(np.asarray(next_state[index], dtype=np.float32).copy())276 episode["available_actions"].append(np.asarray(next_available[index], dtype=np.float32).copy())277 episode["active_masks"].append(278 np.asarray(next_active[index], dtype=np.float32).copy()279 )280 episode["alive_masks"].append(281 (~np.asarray(dones[index], dtype=bool)).astype(np.float32)[:, None]282 )283 if env_dones[index]:284 self._finish_episode(episode)285 self.env_episode += 1286 last_info = infos[index]287 last_episode_return = episode["return"]288 last_sent_percentage = 100.0 * episode["sent"] / max(289 1.0, episode["sent_goal"]290 )291 self.writer.add_scalar("reward/episode", episode["return"], self.env_episode)292 self.writer.add_scalar("train/epsilon", epsilon, self.env_episode)293 episodes[index] = self._new_episode(294 next_obs, next_state, next_available, next_active, index295 )296 hidden = self._advance_rollout_hidden(297 next_hidden, dones, active_masks=active_masks,298 previous_hidden=hidden,299 event_mode=self.config.qmix_transition_mode == "event_smdp",300 )301 last_actions = actions.copy()302 for index, done in enumerate(env_dones):303 if done:304 last_actions[index] = 0305 obs, state, active_masks, available = next_obs, next_state, next_active, next_available306 total_steps += self.n_rollout_threads307308 train_interval = int(self.config.qmix_train_interval)309 learning_ready = len(self.buffer) >= int(self.config.qmix_learning_starts)310 if completed_this_step > 0 and not learning_ready:311 # Warm-up episodes fill replay without accumulating a gradient312 # update debt that would be paid all at once at learning start.313 self.last_train_env_episode = self.env_episode314 elif (completed_this_step > 0 and learning_ready315 and self.env_episode - self.last_train_env_episode >= train_interval):316 due_groups = (317 self.env_episode - self.last_train_env_episode318 ) // train_interval319 for _ in range(due_groups * int(self.config.qmix_updates_per_train)):320 batch = self.buffer.sample(321 self.config.qmix_batch_size,322 self.config.qmix_sequence_length,323 self.config.qmix_burn_in,324 )325 train_info = self.trainer.train_batch(batch)326 updates += 1327 for key, value in train_info.items():328 self.writer.add_scalar(f"train/{key}", value, total_steps)329 self.last_train_env_episode = self.env_episode330331 target_interval = int(self.config.qmix_target_update_interval)332 if (completed_this_step > 0 and target_interval > 0333 and self.env_episode - self.last_target_env_episode >= target_interval):334 self.trainer.hard_target_updates()335 self.last_target_env_episode = self.env_episode336 self.writer.add_scalar("train/target_updates", 1, self.env_episode)337338 save_due = completed_this_step > 0 and self._completed_interval_due(339 self.env_episode, self.last_save_env_episode, self.save_interval340 )341 if save_due:342 self.save()343 if getattr(self.config, "save_episode_checkpoints", False):344 self.save(tag=f"ep{self.env_episode:06d}")345 self.last_save_env_episode = self.env_episode346 eval_due = completed_this_step > 0 and self._completed_interval_due(347 self.env_episode, self.last_eval_env_episode, self.eval_interval348 )349 if eval_due and self.use_eval:350 tag = f"ep{self.env_episode:06d}"351 self.save(tag=tag)352 summary = self.periodic_eval(353 total_steps,354 Path(self.save_dir) / f"actor_{tag}.pt",355 )356 self._save_eval_checkpoints(summary, total_steps)357 self.last_eval_env_episode = self.env_episode358 if (self.env_episode != last_log_episode and self.env_episode359 and self.env_episode % max(1, self.log_interval) == 0):360 elapsed = max(time.time() - start, 1e-6)361 expected = elapsed * self.num_env_steps / max(1, total_steps)362 completed = 0363 if last_info:364 completed = sum(365 int(task.get("completed", False))366 for task in last_info.get("tasks", [])367 if int(task.get("type", -1)) == 1368 )369 target_episodes = int(max_episodes) if max_episodes else max(370 1, self.num_env_steps // max(1, self.episode_length)371 )372 print(373 f"\n rollout {self.env_episode}/{target_episodes}, "374 f"env_episode {self.env_episode}, {last_sent_percentage:.1f}% sent, "375 f"{completed} completed, ep_reward: {last_episode_return:.2f}, "376 f"timesteps {total_steps}/{self.num_env_steps}, "377 f"FPS {int(total_steps / elapsed)}, time cost {elapsed:.1f}, "378 f"expected time {expected:.1f}.\n"379 )380 last_log_episode = self.env_episode381 self.save()382383 def save(self, tag=None):384 suffix = f"_{tag}" if tag else ""385 torch.save(386 self.policy.get_actor_checkpoint(),387 os.path.join(self.save_dir, f"actor{suffix}.pt"),388 )389 torch.save(390 self.trainer.state_dict(),391 os.path.join(self.save_dir, f"trainer{suffix}.pt"),392 )393394 def restore(self):395 checkpoint = torch.load(396 os.path.join(str(self.model_dir), "actor.pt"), map_location=self.device397 )398 self.policy.load_actor_checkpoint(checkpoint)