macf_runner.py
collect/insert、focal state、决策掩码、训练与评价
原文件 · 1318 行 · 54704 字节
runner/shared/macf_runner.py
SHA-256: 453dfeb849d3593e42d58515bdf4038dfb2cb299ad0c6a8c572d84f53df7d719
1"""2# @Time : 2021/7/1 7:15 下午3# @Author : hezhiqiang014# @Email : hezhiqiang01@baidu.com5# @File : env_runner.py6"""78# 原 import9import numpy as np10import torch11from .base_runner import Runner1213# 数据处理14import os15import csv16import json17import math18import pickle19import subprocess20import sys21from datetime import datetime22from pathlib import Path23import time2425import yaml2627# from copy import deepcopy282930def _t2n(x):31 return x.detach().cpu().numpy()323334def _concat_obs_dict(obs):35 if not isinstance(obs, dict):36 return np.concatenate(obs)37 return {k: np.concatenate(v) for k, v in obs.items()}383940class MACFEnvRunner(Runner):41 """Runner class to perform training, evaluation. and data collection for the MPEs. See parent class for details."""4243 def __init__(self, config):44 super(MACFEnvRunner, self).__init__(config)4546 self.episode = 047 self.env_episode = 04849 self.info_buffer = []5051 self.writer = self.writter52 print(getattr(self.writer, "logdir", self.tb_dir))53 self.best_eval_score = None54 self.best_eval_summary = None55 self.last_save_env_episode = -156 self.last_eval_env_episode = -15758 @staticmethod59 def _episode_interval_due(current_episode, last_episode, interval):60 """Check whether an episode-based trigger is due.6162 Supports a single interval or a list of intervals. A trigger fires at63 episode 0 and every ``interval`` episodes afterwards (i.e. at64 ``interval - 1, 2 * interval - 1, ...``) for each configured interval.65 This matches 0-indexed episode counting.66 """67 if current_episode < 0:68 return False69 intervals = interval if isinstance(interval, (list, tuple)) else [interval]70 intervals = [int(i) for i in intervals if int(i) > 0]71 if not intervals:72 return False7374 def _next_node_after(ep, n):75 # Nodes are 0, n-1, 2n-1, 3n-1, ...76 if ep < 0:77 return 078 k = (ep + 1) // n79 candidate = k * n - 180 if candidate > ep:81 return max(0, candidate)82 return (k + 1) * n - 18384 next_nodes = [_next_node_after(last_episode, n) for n in intervals]85 return current_episode >= min(next_nodes)8687 def _info_log_interval(self):88 if self.env_episode < 10:89 return 190 if self.env_episode < 100:91 return 1092 return int(self.log_info_interval)9394 def _publish_eval_progress(self, completed, state="running"):95 progress_file = getattr(self.config, "eval_progress_file", None)96 if not progress_file:97 return98 path = Path(progress_file)99 path.parent.mkdir(parents=True, exist_ok=True)100 payload = {101 "state": str(state),102 "completed_episodes": int(completed),103 "total_episodes": int(self.eval_episodes),104 "updated_at": datetime.now().astimezone().isoformat(timespec="seconds"),105 }106 temporary = path.with_name(f"{path.name}.{os.getpid()}.tmp")107 with temporary.open("w", encoding="utf-8") as file:108 json.dump(payload, file, ensure_ascii=False, indent=2)109 temporary.replace(path)110111 @staticmethod112 def _print_eval_summary(summary):113 print(114 "eval: episodes={}, success={:.1%}, reward={:.2f}±{:.2f}, "115 "length={:.1f}±{:.1f}, successful_length={}".format(116 summary["num_episodes"],117 summary["success"],118 summary["episode_reward"],119 summary["episode_reward_std"],120 summary["episode_length"],121 summary["episode_length_std"],122 "NA"123 if summary["successful_episode_length"] is None124 else f"{summary['successful_episode_length']:.1f}",125 ),126 flush=True,127 )128129 def _log_eval_scalars(self, summary, total_num_steps):130 for key, value in summary.items():131 if value is not None and isinstance(value, (int, float)):132 self.writer.add_scalar(f"eval/{key}", value, total_num_steps)133134 def _external_periodic_eval(self, total_num_steps, checkpoint_path, eval_cuda):135 """Evaluate a frozen checkpoint on another CUDA device.136137 The training process waits at the checkpoint boundary, preserving the138 existing best/latest selection order. Other jobs on the training GPU139 continue, while this job's evaluation model and environments live in a140 short-lived subprocess on ``eval_cuda``.141 """142143 config_file = getattr(self.config, "config_file", None)144 if not config_file or not Path(config_file).is_file():145 raise RuntimeError(146 "training_eval_cuda requires training to use a readable config_file"147 )148 root = Path(self.run_dir) / "training_eval_external"149 step_dir = root / f"step_{int(total_num_steps):012d}"150 step_dir.mkdir(parents=True, exist_ok=True)151 with Path(config_file).open("r", encoding="utf-8") as file:152 eval_config = yaml.safe_load(file) or {}153 eval_config.update(154 cuda=int(eval_cuda),155 use_eval=True,156 model_dir=None,157 eval_progress_file=None,158 )159 eval_config_path = step_dir / "config.yaml"160 with eval_config_path.open("w", encoding="utf-8") as file:161 yaml.safe_dump(eval_config, file, sort_keys=False, allow_unicode=True)162163 manifest = {164 "name": f"training-eval-step-{int(total_num_steps):012d}",165 "output_root": str(step_dir / "results"),166 "models": [167 {168 "name": "checkpoint",169 "checkpoint": str(Path(checkpoint_path).resolve()),170 "overrides": {},171 }172 ],173 }174 manifest_path = step_dir / "manifest.yaml"175 with manifest_path.open("w", encoding="utf-8") as file:176 yaml.safe_dump(manifest, file, sort_keys=False, allow_unicode=True)177 progress_path = step_dir / "progress.json"178 command = [179 sys.executable,180 str(Path(__file__).resolve().parents[2] / "train" / "eval.py"),181 "--config_file",182 str(eval_config_path),183 "--models_file",184 str(manifest_path),185 "--progress_file",186 str(progress_path),187 ]188 guard = eval_config.get("gpu_thermal_guard", {}) or {}189 if guard.get("enabled", False):190 command = [191 sys.executable,192 str(193 Path(__file__).resolve().parents[2]194 / "scripts"195 / "gpu_thermal_guard.py"196 ),197 "--gpu", str(int(eval_cuda)),198 "--pause-at", str(int(guard.get("pause_at_c", 93))),199 "--resume-at", str(int(guard.get("resume_at_c", 85))),200 "--poll-seconds", str(float(guard.get("poll_seconds", 2))),201 "--safe-readings", str(int(guard.get("safe_readings", 3))),202 "--query-failures", str(int(guard.get("query_failures", 2))),203 "--status-file", str(step_dir / "thermal_status.json"),204 "--",205 *command,206 ]207 environment = os.environ.copy()208 environment["PYTHONUNBUFFERED"] = "1"209 process = subprocess.Popen(210 command,211 cwd=Path(__file__).resolve().parents[2],212 env=environment,213 stdin=subprocess.DEVNULL,214 stdout=subprocess.PIPE,215 stderr=subprocess.STDOUT,216 text=True,217 bufsize=1,218 )219 output_lines = []220 suite_dir = None221 assert process.stdout is not None222 for line in process.stdout:223 output_lines.append(line)224 stripped = line.strip()225 if stripped.startswith("eval progress:"):226 print(stripped, flush=True)227 if "Evaluation suite written to" in stripped:228 suite_dir = stripped.split("Evaluation suite written to", 1)[1].strip()229 return_code = process.wait()230 if return_code != 0 or not suite_dir:231 tail = "".join(output_lines[-30:])232 raise RuntimeError(233 f"external periodic eval failed on cuda:{eval_cuda} "234 f"with return code {return_code}\n{tail}"235 )236 summary_path = Path(suite_dir) / "summary.json"237 with summary_path.open("r", encoding="utf-8") as file:238 payload = json.load(file)239 records = payload.get("models") or []240 if len(records) != 1:241 raise RuntimeError(f"Unexpected external eval summary: {summary_path}")242 summary = {243 key: value244 for key, value in records[0].items()245 if key not in ("name", "checkpoint")246 }247 self._log_eval_scalars(summary, total_num_steps)248 self._print_eval_summary(summary)249 return summary250251 def periodic_eval(self, total_num_steps, checkpoint_path):252 eval_cuda = getattr(self.config, "training_eval_cuda", None)253 if eval_cuda is None or int(eval_cuda) == int(self.config.cuda):254 return self.eval(total_num_steps)255 return self._external_periodic_eval(256 total_num_steps, checkpoint_path, int(eval_cuda)257 )258259 def run(self):260 # 初始化261 self.warmup()262263 start = datetime.now()264 episodes = int(self.num_env_steps) // self.episode_length // self.n_rollout_threads265 max_train_episodes = getattr(self.config, "max_train_episodes", None)266 if max_train_episodes is not None:267 max_train_episodes = int(max_train_episodes)268 if max_train_episodes <= 0:269 raise ValueError("max_train_episodes must be positive or None.")270 episodes = min(episodes, max_train_episodes)271272 # 在文件夹不存在时创建文件夹273 os.makedirs(os.path.join(self.log_dir, 'infos'), exist_ok=True)274275 # 时间步循环276 episode_reward = 0277 sent = 0278 sent_goal = 0279 last_episode_sent_percentage = 0.0280281 for self.episode in range(episodes):282 episode = self.episode283 # self.warmup()284285 rollout_reward = 0286 success_flags = 0287288 if self.use_linear_lr_decay:289 self.trainer.policy.lr_decay(episode, episodes)290291 for step in range(self.episode_length):292 # Sample actions 动作采样293 (294 values,295 actions,296 action_log_probs,297 gate_actions,298 rnn_states,299 rnn_states_critic,300 ) = self.collect(step)301302 # Obser reward and next obs 环境交互303 obs, rewards, dones, infos, active_masks, bad_masks, available_actions = self.envs.step(actions)304 state = self.envs.get_global_state()305306 sent += infos[0]['sent']307 sent_goal = infos[0]['sent_goal']308 episode_reward += rewards.mean()309 rollout_reward += rewards.mean()310311 data = (312 obs,313 state,314 rewards,315 dones,316 infos,317 values,318 actions,319 action_log_probs,320 gate_actions,321 rnn_states,322 rnn_states_critic,323 active_masks,324 bad_masks,325 available_actions,326 )327328 # insert data into buffer 插入经验回放缓冲区329 self.insert(data)330331 # buffer 写入后再记录 scalar332 if infos[0]['end']:333 tc = sum(334 1335 for t in infos[0]['tasks']336 if int(t.get("type", -1)) == 1 and t["completed"]337 )338 sp = sent / sent_goal if sent_goal > 0 else 0339 last_episode_sent_percentage = sp340 self.writer.add_scalar("reward/episode", episode_reward, self.env_episode)341 self.writer.add_scalar("reward/sent_percentage", sp, self.env_episode)342 self.writer.add_scalar("reward/task_completed", tc, self.env_episode)343 self.writer.add_scalar("reward/episode_length", infos[0]['end'], self.env_episode)344 reconfiguration = infos[0].get("reconfiguration", {})345 reconfiguration_decisions = int(346 reconfiguration.get("total_decisions", 0)347 )348 reconfiguration_changes = int(349 reconfiguration.get("total_coalition_changes", 0)350 )351 self.writer.add_scalar(352 "reconfiguration/decisions",353 reconfiguration_decisions,354 self.env_episode,355 )356 self.writer.add_scalar(357 "reconfiguration/coalition_changes",358 reconfiguration_changes,359 self.env_episode,360 )361 self.writer.add_scalar(362 "reconfiguration/change_per_decision",363 reconfiguration_changes / max(1, reconfiguration_decisions),364 self.env_episode,365 )366 alive_agent_steps = int(367 reconfiguration.get("total_alive_agent_steps", 0)368 )369 self.writer.add_scalar(370 "reconfiguration/decisions_per_100_alive_steps",371 100.0 * reconfiguration_decisions / max(1, alive_agent_steps),372 self.env_episode,373 )374 self.writer.add_scalar(375 "reconfiguration/changes_per_100_alive_steps",376 100.0 * reconfiguration_changes / max(1, alive_agent_steps),377 self.env_episode,378 )379 self.writer.add_scalar(380 "reconfiguration/mean_response_time",381 float(reconfiguration.get("mean_response_time", 0.0)),382 self.env_episode,383 )384 self.writer.add_scalar(385 "reconfiguration/urgency_pressure",386 float(reconfiguration.get("urgency_pressure", 0.0)),387 self.env_episode,388 )389 # ===== 记录零初始化门控 =====390 if (391 self.use_gat392 and getattr(self.trainer.policy, "graph_model", None) == "legacy_gat"393 and hasattr(394 self.trainer.policy,395 "get_gat_gate_values"396 )397 ):398 actor_alpha, critic_alpha = (399 self.trainer.policy.get_gat_gate_values()400 )401402 self.writer.add_scalar(403 "gat/actor_alpha",404 actor_alpha,405 self.env_episode406 )407408 self.writer.add_scalar(409 "gat/critic_alpha",410 critic_alpha,411 self.env_episode412 )413 if self.use_gat and hasattr(414 self.trainer.policy, "get_graph_diagnostics"415 ):416 for tag, value in self.trainer.policy.get_graph_diagnostics().items():417 self.writer.add_scalar(tag, value, self.env_episode)418 # ==========================419420 # 保存 info_buffer421 log_info_interval = self._info_log_interval()422423 if self.log_info_interval > 0:424 if self.env_episode % log_info_interval == 0:425 with open(os.path.join(self.log_dir, 'infos', f'{self.env_episode:04d}.pkl'), 'wb') as f:426 pickle.dump(self.info_buffer[1:], f)427 print(f'{self.env_episode:04d}.pkl generated')428 self.info_buffer = []429430 self.env_episode += 1431 episode_reward = 0432 sent = 0433 if (434 self.config.curriculum_switch_ep > 0435 and self.env_episode >= self.config.curriculum_switch_ep436 and getattr(self, "_curriculum_switched", False) is False437 ):438 for env in self.envs.envs:439 env.set_reset_mode("full_random")440 self._curriculum_switched = True441 print(442 f"curriculum: switched to full_random at env_episode {self.env_episode}"443 )444 if (445 max_train_episodes is not None446 and self.env_episode >= max_train_episodes447 and step + 1 < self.episode_length448 ):449 # A partial rollout cannot be used by the fixed-size PPO buffer.450 # Preserve the latest trained weights and stop at the exact env-episode cap.451 self.save()452 return453454 # if np.all(dones):455 # break456457 t0 = t = datetime.now().timestamp()458 ta = []459460 # info_buffer 可能已在 env episode 结束时清空，用最后一条 info461 last_info = self.info_buffer[-1][0] if self.info_buffer else infos[0]462 sent_goal = last_info['sent_goal']463 task_completed = sum(464 1465 for t in last_info['tasks']466 if int(t.get("type", -1)) == 1 and t["completed"]467 )468469 t = datetime.now().timestamp()470 ta.append(t - t0)471 t0 = t472473 # compute return and update network474 self.compute()475476 t = datetime.now().timestamp()477 ta.append(t - t0)478 t0 = t479480 train_infos = self.train()481482 t = datetime.now().timestamp()483 ta.append(t - t0)484 t0 = t485486 # post process487 total_num_steps = (episode + 1) * self.episode_length * self.n_rollout_threads488489 # save model490 save_due = self._episode_interval_due(491 self.env_episode,492 self.last_save_env_episode,493 self.save_interval,494 )495 if save_due or episode == episodes - 1:496 self.save()497 if getattr(self.config, "save_episode_checkpoints", False):498 self.save(tag=f"ep{self.env_episode:06d}")499 self.last_save_env_episode = self.env_episode500501 # log information502 if episode % self.log_interval == 0:503 end = datetime.now()504505 print(506 "\n rollout {}/{}, env_episode {}, {}% sent, {} completed, ep_reward: {:>7.2f}, timesteps {}/{}, FPS {}, time cost {}, expected time {}.\n".format(507 episode,508 episodes,509 self.env_episode,510 last_episode_sent_percentage * 100,511 task_completed,512 rollout_reward,513 total_num_steps,514 self.num_env_steps,515 int(total_num_steps / (end - start).total_seconds()),516 (end - start).total_seconds(),517 float(episodes)/(1 if episode == 0 else episode)*((end - start).total_seconds())518 )519 )520521 train_infos["average_episode_rewards"] = np.mean(self.buffer.rewards) * self.episode_length522 print("average episode rewards is {}".format(train_infos["average_episode_rewards"]))523 self.log_train(train_infos, total_num_steps)524 # self.log_env(env_infos, total_num_steps)525526 # eval527 eval_due = self._episode_interval_due(528 self.env_episode,529 self.last_eval_env_episode,530 self.eval_interval,531 )532 if eval_due and self.use_eval:533 tag = f"ep{self.env_episode:06d}"534 self.save(tag=tag)535 summary = self.periodic_eval(536 total_num_steps,537 Path(self.save_dir) / f"actor_{tag}.pt",538 )539 self._save_eval_checkpoints(summary, total_num_steps)540 self.last_eval_env_episode = self.env_episode541542 if (543 max_train_episodes is not None544 and self.env_episode >= max_train_episodes545 ):546 break547548 t = datetime.now().timestamp()549 ta.append(t - t0)550 t0 = t551552553554 def warmup(self):555 # reset env — obs is dict556 obs, active_masks, available_actions = self.envs.reset(self.episode) # obs: dict[str, ndarray]557 state = self.envs.get_global_state()558559 self.info_buffer = []560561 share_obs = self._build_share_obs(state)562563 for k in self.buffer.share_obs:564 self.buffer.share_obs[k][0] = share_obs[k].copy()565 for k in self.buffer.obs:566 self.buffer.obs[k][0] = obs[k].copy()567 self.buffer.active_masks[0] = active_masks.copy()568 self.buffer.available_actions[0] = available_actions.copy()569570 def _build_share_obs(self, state):571 if not self.use_centralized_V:572 raise NotImplementedError("Structured local state for use_centralized_V=False is not implemented.")573574 repeated = {575 k: np.repeat(v[:, None, ...], self.num_agents, axis=1)576 for k, v in state.items()577 }578 focal = np.eye(self.num_agents, dtype=np.float32)[None, :, :, None]579 repeated["focal"] = np.repeat(focal, self.n_rollout_threads, axis=0)580 return repeated581582 def _get_actions(self, step, available_actions):583 obs_batch = _concat_obs_dict({k: v[step] for k, v in self.buffer.obs.items()})584 if self.use_gat:585 return self.trainer.policy.get_actions(586 {k: np.concatenate(v[step]) for k, v in self.buffer.share_obs.items()},587 obs_batch,588 np.concatenate(self.buffer.rnn_states[step]),589 np.concatenate(self.buffer.rnn_states_critic[step]),590 np.concatenate(self.buffer.masks[step]),591 available_actions=np.concatenate(available_actions)592 )593 else:594 return self.trainer.policy.get_actions(595 {k: np.concatenate(v[step]) for k, v in self.buffer.share_obs.items()},596 obs_batch,597 np.concatenate(self.buffer.rnn_states[step]),598 np.concatenate(self.buffer.rnn_states_critic[step]),599 np.concatenate(self.buffer.masks[step]),600 available_actions=np.concatenate(available_actions)601 )602603 @torch.no_grad()604 def collect(self, step):605606 available_actions = self.buffer.available_actions[step]607 active_masks = self.buffer.active_masks[step]608609 self.trainer.prep_rollout()610611 action_outputs = self._get_actions(step, available_actions)612 if len(action_outputs) == 6:613 (614 value,615 action,616 action_log_prob,617 gate_action,618 rnn_states,619 rnn_states_critic,620 ) = action_outputs621 else:622 # Standard MAPPO has no explicit hierarchical gate action. Keep a623 # zero placeholder in the runner tuple so the environment loop can624 # remain shared, but do not pass it into SharedReplayBuffer below.625 (626 value,627 action,628 action_log_prob,629 rnn_states,630 rnn_states_critic,631 ) = action_outputs632 gate_action = torch.zeros_like(action, dtype=torch.float32)633 # [self.envs, agents, dim]634 values = np.array(np.split(_t2n(value), self.n_rollout_threads)) # [env_num, agent_num, 1]635 actions = np.array(np.split(_t2n(action), self.n_rollout_threads)) # [env_num, agent_num, action_dim]636 action_log_probs = np.array(637 np.split(_t2n(action_log_prob), self.n_rollout_threads)638 ) # [env_num, agent_num, 1]639 gate_actions = np.array(640 np.split(_t2n(gate_action), self.n_rollout_threads)641 )642 rnn_states = np.array(643 np.split(_t2n(rnn_states), self.n_rollout_threads)644 ) # [env_num, agent_num, 1, hidden_size]645 rnn_states_critic = np.array(646 np.split(_t2n(rnn_states_critic), self.n_rollout_threads)647 ) # [env_num, agent_num, 1, hidden_size]648649 if step == 0:650 last_actions = self.buffer.actions[-1]651 last_action_log_probs = self.buffer.action_log_probs[-1]652 else:653 last_actions = self.buffer.actions[step-1]654 last_action_log_probs = self.buffer.action_log_probs[step-1]655656 # 添加 active_masks 掩码657 actions = active_masks * actions + (1 - active_masks) * last_actions658 action_log_probs = active_masks * action_log_probs + (1 - active_masks) * last_action_log_probs659 # inactive agent 沿用上一动作，不产生新的 gate 决策；否则 PPO 会把660 # 环境未执行的 redecision 误当作 rollout 样本。661 gate_actions = active_masks * gate_actions662663 return (664 values,665 actions,666 action_log_probs,667 gate_actions,668 rnn_states,669 rnn_states_critic,670 # active agents671 )672673 def insert(self, data):674 (675 obs,676 state,677 rewards,678 dones,679 infos,680 values,681 actions,682 action_log_probs,683 gate_actions,684 rnn_states,685 rnn_states_critic,686 active_masks,687 bad_masks,688 available_actions,689 ) = data690691 rnn_states[dones == True] = np.zeros(692 ((dones == True).sum(), self.recurrent_N, self.hidden_size),693 dtype=np.float32,694 )695 rnn_states_critic[dones == True] = np.zeros(696 ((dones == True).sum(), *self.buffer.rnn_states_critic.shape[3:]),697 dtype=np.float32,698 )699 masks = np.ones((self.n_rollout_threads, self.num_agents, 1), dtype=np.float32)700 masks[dones == True] = np.zeros(((dones == True).sum(), 1), dtype=np.float32)701702 share_obs = self._build_share_obs(state)703704 if bool(getattr(self.config, "mappo_team_mean_reward", False)):705 rewards = np.repeat(706 rewards.mean(axis=1, keepdims=True), self.num_agents, axis=1707 )708709 # rewards = rewards[:, None] # (T,) → (T,1)710 # masks = masks[:, None] # 同理711712 common_args = (713 share_obs,714 obs,715 rnn_states,716 rnn_states_critic,717 actions,718 action_log_probs,719 )720 trailing_args = (values, rewards, masks)721 if self.use_gat:722 self.buffer.insert(723 *common_args,724 gate_actions,725 *trailing_args,726 active_masks=active_masks,727 bad_masks=bad_masks,728 available_actions=available_actions,729 )730 else:731 self.buffer.insert(732 *common_args,733 *trailing_args,734 active_masks=active_masks,735 bad_masks=bad_masks,736 available_actions=available_actions,737 )738739 self.info_buffer.append(infos)740741 @torch.no_grad()742 def _run_eval_episode(self, env, scenario_id):743 scenario_seed = self.eval_seed + scenario_id744 numpy_state = np.random.get_state()745 np.random.seed(scenario_seed)746 try:747 env.seed(scenario_seed)748 obs, active_masks, available_actions = env.reset(scenario_id)749 rnn_states = np.zeros(750 (self.num_agents, self.recurrent_N, self.hidden_size),751 dtype=np.float32,752 )753 masks = np.ones((self.num_agents, 1), dtype=np.float32)754 last_actions = np.zeros((self.num_agents, 1), dtype=np.float32)755 previous_coalitions = np.array(756 [agent.coalition for agent in env.env.agents], dtype=np.int64757 )758759 episode_reward = 0.0760 sent = 0.0761 coalition_changes = 0762 premature_no_explorer_steps = 0763 reward_components = {}764 decision_type_counts = {task_type: 0 for task_type in range(4)}765 executed_type_counts = {task_type: 0 for task_type in range(4)}766 decision_count = 0767 executed_count = 0768 gate_decision_count = 0769 gate_open_count = 0770 gate_mandatory_count = 0771 gate_mandatory_open_count = 0772 gate_event_trigger_count = 0773 gate_event_trigger_open_count = 0774 gate_learned_count = 0775 gate_learned_open_count = 0776 num_explore_actions = int(777 getattr(env.env, "num_explore_actions", 1)778 )779 sector_decision_counts = np.zeros(780 (self.num_agents, num_explore_actions), dtype=np.int64781 )782 explore_legs_planned = 0783 explore_planned_distance = 0.0784 short_explore_legs = 0785 explore_distance_moved = 0.0786 unavailable_sector_ratio_sum = 0.0787 alive_agent_steps = 0788 harmful_idle_agent_steps = 0789 useful_idle_agent_steps = 0790 committed_payload_shortage_ratio_sum = 0.0791 committed_payload_surplus_ratio_sum = 0.0792 max_idle_duration = 0.0793 final_info = None794795 self.trainer.prep_rollout()796 eval_step_limit = int(797 getattr(self.config, "episode_time_limit", None)798 or self.config.max_steps799 )800 for step in range(eval_step_limit):801 previous_rnn_states = rnn_states802 decision_active_masks = np.asarray(active_masks).copy()803 actions, next_rnn_states = self.trainer.policy.act(804 obs,805 rnn_states,806 masks,807 available_actions=available_actions,808 deterministic=getattr(self.config, "eval_deterministic", True),809 )810 gate_actions = getattr(811 getattr(self.trainer.policy, "actor", None),812 "last_gate_actions",813 None,814 )815 actions = _t2n(actions)816 next_rnn_states = _t2n(next_rnn_states)817 proposed_actions = actions.astype(np.int64).reshape(-1)818 decision_mask = np.asarray(active_masks).reshape(-1) > 0.5819 if gate_actions is not None:820 # hierarchical gate 的 z=1 表示本次 active 机会真正交给任务821 # decoder；z=0 表示保持 current。这里按 eval episode 口径统计，822 # 区分环境强制开门和 learned gate 自主开门。823 gate_values = _t2n(gate_actions).reshape(-1) > 0.5824 active_gate_values = gate_values[decision_mask]825 gate_decision_count += int(active_gate_values.size)826 gate_open_count += int(np.count_nonzero(active_gate_values))827 reconfig_obs = (828 obs.get("reconfiguration")829 if isinstance(obs, dict)830 else None831 )832 if reconfig_obs is not None:833 reconfig_values = np.asarray(reconfig_obs)834 context_dim = reconfig_values.shape[-1]835 # Match HATTV20Actor's event override exactly. The836 # previous metric treated only context[10] as forced,837 # so event-only evaluation incorrectly reported its838 # context[0/1/3] openings as "learned" openings.839 event_trigger = np.zeros(840 reconfig_values.shape[:-1], dtype=bool841 )842 if context_dim >= 4:843 event_trigger = (844 (reconfig_values[..., 0] > 0.5)845 | (reconfig_values[..., 1] > 0.5)846 | (reconfig_values[..., 3] > 0.5)847 )848 mandatory_all = np.zeros_like(event_trigger)849 if context_dim >= 11:850 mandatory_all = reconfig_values[..., 10] > 0.5851 event_trigger = event_trigger | mandatory_all852 event_trigger = event_trigger.reshape(-1)[decision_mask]853 mandatory = mandatory_all.reshape(-1)[decision_mask]854 gate_mandatory_count += int(np.count_nonzero(mandatory))855 gate_mandatory_open_count += int(856 np.count_nonzero(active_gate_values & mandatory)857 )858 gate_event_trigger_count += int(859 np.count_nonzero(event_trigger)860 )861 gate_event_trigger_open_count += int(862 np.count_nonzero(active_gate_values & event_trigger)863 )864 # Keep the historical learned metric for downstream865 # CSV compatibility, but give it the correct meaning:866 # openings outside every forced event trigger.867 learned = ~event_trigger868 gate_learned_count += int(np.count_nonzero(learned))869 gate_learned_open_count += int(870 np.count_nonzero(active_gate_values & learned)871 )872 for agent_idx in np.flatnonzero(decision_mask):873 action_id = proposed_actions[agent_idx]874 task_type = int(env.env.tasks[int(action_id)].type)875 decision_type_counts[task_type] += 1876 decision_count += 1877 if (878 task_type == 0879 and getattr(env.env, "explore_action_mode", None)880 == "sector"881 ):882 sector = int(action_id) - int(env.env.explore_task_start)883 if 0 <= sector < num_explore_actions:884 sector_decision_counts[agent_idx, sector] += 1885 actions = active_masks * actions + (1.0 - active_masks) * last_actions886 for action_id in actions.astype(np.int64).reshape(-1):887 task_type = int(env.env.tasks[int(action_id)].type)888 executed_type_counts[task_type] += 1889 executed_count += 1890891 (892 obs,893 rewards,894 dones,895 info,896 active_masks,897 _,898 available_actions,899 ) = env.step(actions)900 final_info = info901 episode_reward += float(np.mean(rewards))902 sent += float(info["sent"])903 allocation = info.get("allocation_diagnostics", {})904 exploration = info.get("exploration_diagnostics", {})905 alive_agent_steps += int(sum(agent["living"] for agent in info["agents"]))906 harmful_idle_agent_steps += int(allocation.get("harmful_idle_agents", 0))907 useful_idle_agent_steps += int(allocation.get("useful_idle_agents", 0))908 committed_payload_shortage_ratio_sum += float(909 allocation.get("committed_payload_shortage_ratio", 0.0)910 )911 committed_payload_surplus_ratio_sum += float(912 allocation.get("committed_payload_surplus_ratio", 0.0)913 )914 explore_legs_planned += int(exploration.get("legs_planned", 0))915 explore_planned_distance += float(916 exploration.get("planned_distance", 0.0)917 )918 short_explore_legs += int(exploration.get("short_legs", 0))919 explore_distance_moved += float(920 exploration.get("distance_moved", 0.0)921 )922 unavailable_sector_ratio_sum += float(923 exploration.get("unavailable_sector_ratio", 0.0)924 )925 max_idle_duration = max(926 max_idle_duration,927 max(928 (929 float(agent.get("idle_wait_steps", 0.0))930 for agent in info["agents"]931 if int(info["tasks"][int(agent["coalition"])]["type"]) == 2932 ),933 default=0.0,934 ),935 )936937 component_values = np.asarray(info["rewards"], dtype=np.float64).mean(axis=1)938 for key, value in zip(info["reward_keys"], component_values):939 reward_components[key] = reward_components.get(key, 0.0) + float(value)940941 coalitions = np.array(942 [agent["coalition"] for agent in info["agents"]], dtype=np.int64943 )944 coalition_changes += int(np.count_nonzero(coalitions != previous_coalitions))945 previous_coalitions = coalitions946947 delivery_tasks = [task for task in info["tasks"] if int(task["type"]) == 1]948 discovered_ratio = (949 sum(task["discovered"] for task in delivery_tasks)950 / max(1, len(delivery_tasks))951 )952 any_explorer = any(953 agent["living"]954 and int(info["tasks"][int(agent["coalition"])]["type"]) == 0955 for agent in info["agents"]956 )957 if not any_explorer and discovered_ratio <= 0.9:958 premature_no_explorer_steps += 1959960 done_mask = np.asarray(dones, dtype=bool)961 if (962 str(getattr(self.config, "qmix_transition_mode", ""))963 == "event_smdp"964 ):965 if not np.any(decision_active_masks > 0.5):966 next_rnn_states = previous_rnn_states967 next_rnn_states[done_mask] = 0.0968 masks = np.ones((self.num_agents, 1), dtype=np.float32)969 masks[done_mask] = 0.0970 rnn_states = next_rnn_states971 last_actions = actions972973 if bool(np.all(done_mask)) or info.get("end", 0):974 break975 finally:976 np.random.set_state(numpy_state)977978 if final_info is None:979 raise RuntimeError("Evaluation episode ended without environment info.")980981 delivery_tasks = [982 task for task in final_info["tasks"] if int(task["type"]) == 1983 ]984 tasks_completed = sum(task["completed"] for task in delivery_tasks)985 tasks_discovered = sum(task["discovered"] for task in delivery_tasks)986 episode_length = step + 1987 result = {988 "scenario_id": int(scenario_id),989 "scenario_seed": int(scenario_seed),990 "success": int(tasks_completed == len(delivery_tasks)),991 "episode_reward": float(episode_reward),992 "episode_length": episode_length,993 "tasks_completed": int(tasks_completed),994 "task_completion_ratio": float(995 tasks_completed / max(1, len(delivery_tasks))996 ),997 "sent_percentage": float(sent / max(1.0, final_info["sent_goal"])),998 "discovered_ratio": float(tasks_discovered / max(1, len(delivery_tasks))),999 "dead_agents": int(sum(not agent["living"] for agent in final_info["agents"])),1000 "coalition_changes": int(coalition_changes),1001 "reconfiguration_decisions": int(1002 final_info.get("reconfiguration", {}).get(1003 "total_decisions", decision_count1004 )1005 ),1006 "coalition_changes_per_decision": float(1007 coalition_changes / max(1008 1,1009 final_info.get("reconfiguration", {}).get(1010 "total_decisions", decision_count1011 ),1012 )1013 ),1014 "reconfiguration_decisions_per_100_alive_steps": float(1015 100.01016 * final_info.get("reconfiguration", {}).get(1017 "total_decisions", decision_count1018 )1019 / max(1020 1,1021 final_info.get("reconfiguration", {}).get(1022 "total_alive_agent_steps", alive_agent_steps1023 ),1024 )1025 ),1026 "coalition_changes_per_100_alive_steps": float(1027 100.01028 * coalition_changes1029 / max(1030 1,1031 final_info.get("reconfiguration", {}).get(1032 "total_alive_agent_steps", alive_agent_steps1033 ),1034 )1035 ),1036 "mean_response_time": float(1037 final_info.get("reconfiguration", {}).get(1038 "mean_response_time", 0.01039 )1040 ),1041 "premature_no_explorer_ratio": float(1042 premature_no_explorer_steps / max(1, episode_length)1043 ),1044 "harmful_idle_agent_step_ratio": float(1045 harmful_idle_agent_steps / max(1, alive_agent_steps)1046 ),1047 "useful_idle_agent_step_ratio": float(1048 useful_idle_agent_steps / max(1, alive_agent_steps)1049 ),1050 "max_idle_duration": float(max_idle_duration),1051 "committed_payload_shortage_ratio": float(1052 committed_payload_shortage_ratio_sum / max(1, episode_length)1053 ),1054 "committed_payload_surplus_ratio": float(1055 committed_payload_surplus_ratio_sum / max(1, episode_length)1056 ),1057 "mean_explore_leg_distance": float(1058 explore_planned_distance / max(1, explore_legs_planned)1059 ),1060 "short_explore_leg_ratio": float(1061 short_explore_legs / max(1, explore_legs_planned)1062 ),1063 "explore_distance_per_alive_step": float(1064 explore_distance_moved / max(1, alive_agent_steps)1065 ),1066 "unavailable_sector_ratio": float(1067 unavailable_sector_ratio_sum / max(1, episode_length)1068 ),1069 "explored_area_ratio": float(1070 final_info.get("exploration_diagnostics", {}).get(1071 "explored_area_ratio", 0.01072 )1073 ),1074 }1075 if getattr(env.env, "explore_action_mode", None) == "sector":1076 agent_entropies = []1077 agent_dominant_ratios = []1078 agent_coverage_ratios = []1079 for counts in sector_decision_counts:1080 total = int(np.sum(counts))1081 if total <= 0:1082 continue1083 probabilities = counts[counts > 0].astype(np.float64) / total1084 entropy = -float(np.sum(probabilities * np.log(probabilities)))1085 if num_explore_actions > 1:1086 entropy /= float(np.log(num_explore_actions))1087 agent_entropies.append(entropy)1088 agent_dominant_ratios.append(float(np.max(counts) / total))1089 agent_coverage_ratios.append(1090 float(np.count_nonzero(counts) / num_explore_actions)1091 )1092 result["mean_agent_sector_entropy"] = float(1093 np.mean(agent_entropies) if agent_entropies else 0.01094 )1095 result["mean_agent_dominant_sector_ratio"] = float(1096 np.mean(agent_dominant_ratios)1097 if agent_dominant_ratios1098 else 0.01099 )1100 result["mean_agent_sector_coverage_ratio"] = float(1101 np.mean(agent_coverage_ratios)1102 if agent_coverage_ratios1103 else 0.01104 )1105 task_type_names = ("explore", "deliver", "idle", "resupply")1106 for task_type, task_type_name in enumerate(task_type_names):1107 result[f"decision_{task_type_name}_ratio"] = float(1108 decision_type_counts[task_type] / max(1, decision_count)1109 )1110 result[f"executed_{task_type_name}_ratio"] = float(1111 executed_type_counts[task_type] / max(1, executed_count)1112 )1113 result["decision_count"] = int(decision_count)1114 result["gate_decision_count"] = int(gate_decision_count)1115 result["gate_open_count"] = int(gate_open_count)1116 result["gate_open_rate"] = float(1117 gate_open_count / max(1, gate_decision_count)1118 )1119 result["gate_mandatory_open_rate"] = float(1120 gate_mandatory_open_count / max(1, gate_mandatory_count)1121 )1122 result["gate_event_trigger_open_rate"] = float(1123 gate_event_trigger_open_count / max(1, gate_event_trigger_count)1124 )1125 result["gate_off_trigger_open_rate"] = float(1126 gate_learned_open_count / max(1, gate_learned_count)1127 )1128 result["gate_learned_open_rate"] = float(1129 gate_learned_open_count / max(1, gate_learned_count)1130 )1131 result["coalition_changes_per_gate_open"] = float(1132 coalition_changes / max(1, gate_open_count)1133 )1134 result.update(1135 {f"reward_{key}": float(value) for key, value in reward_components.items()}1136 )1137 return result11381139 @staticmethod1140 def _summarize_eval(episodes):1141 numeric_keys = [1142 key1143 for key, value in episodes[0].items()1144 if key not in {"scenario_id", "scenario_seed"}1145 and isinstance(value, (int, float, np.integer, np.floating))1146 ]1147 summary = {1148 key: float(np.mean([episode[key] for episode in episodes]))1149 for key in numeric_keys1150 }1151 summary["episode_reward_std"] = float(1152 np.std([episode["episode_reward"] for episode in episodes])1153 )1154 summary["episode_length_std"] = float(1155 np.std([episode["episode_length"] for episode in episodes])1156 )1157 successful_lengths = [1158 episode["episode_length"] for episode in episodes if episode["success"]1159 ]1160 summary["successful_episode_length"] = (1161 float(np.mean(successful_lengths)) if successful_lengths else None1162 )1163 summary["successful_episode_length_median"] = (1164 float(np.median(successful_lengths)) if successful_lengths else None1165 )1166 summary["num_episodes"] = len(episodes)1167 return summary11681169 def _save_eval_results(self, episodes, summary, total_num_steps):1170 eval_dir = os.path.join(self.log_dir, "eval")1171 os.makedirs(eval_dir, exist_ok=True)1172 stem = f"step_{int(total_num_steps):012d}"1173 csv_path = os.path.join(eval_dir, stem + ".csv")1174 json_path = os.path.join(eval_dir, stem + ".json")11751176 with open(csv_path, "w", newline="", encoding="utf-8") as file:1177 writer = csv.DictWriter(file, fieldnames=list(episodes[0].keys()))1178 writer.writeheader()1179 writer.writerows(episodes)1180 with open(json_path, "w", encoding="utf-8") as file:1181 json.dump(1182 {"summary": summary, "episodes": episodes},1183 file,1184 ensure_ascii=False,1185 indent=2,1186 )11871188 @staticmethod1189 def _eval_score(summary):1190 """Lexicographic score for selecting the best evaluation checkpoint."""1191 successful_length = summary.get("successful_episode_length")1192 if successful_length is None:1193 successful_length_score = float("-inf")1194 else:1195 successful_length_score = -float(successful_length)1196 return (1197 float(summary.get("success", 0.0)),1198 float(summary.get("task_completion_ratio", 0.0)),1199 float(summary.get("episode_reward", float("-inf"))),1200 successful_length_score,1201 )12021203 def _save_eval_checkpoints(self, summary, total_num_steps):1204 self.save(tag="eval_latest")1205 score = self._eval_score(summary)1206 is_best = self.best_eval_score is None or score > self.best_eval_score1207 json_score = [value if math.isfinite(value) else None for value in score]1208 metadata = {1209 "total_num_steps": int(total_num_steps),1210 "env_episode": int(self.env_episode),1211 "score": json_score,1212 "summary": summary,1213 }1214 latest_path = os.path.join(self.save_dir, "eval_latest.json")1215 with open(latest_path, "w", encoding="utf-8") as file:1216 json.dump(metadata, file, ensure_ascii=False, indent=2)1217 if is_best:1218 self.best_eval_score = score1219 self.best_eval_summary = dict(summary)1220 self.save(tag="eval_best")1221 best_path = os.path.join(self.save_dir, "eval_best.json")1222 with open(best_path, "w", encoding="utf-8") as file:1223 json.dump(metadata, file, ensure_ascii=False, indent=2)12241225 @torch.no_grad()1226 def eval(self, total_num_steps):1227 if self.eval_envs is None:1228 raise RuntimeError("use_eval=True requires eval_envs.")12291230 episodes = []1231 self._publish_eval_progress(0)1232 envs = self.eval_envs.envs1233 for scenario_id in range(self.eval_episodes):1234 env = envs[scenario_id % len(envs)]1235 episodes.append(self._run_eval_episode(env, scenario_id))1236 completed = scenario_id + 11237 self._publish_eval_progress(1238 completed,1239 state="completed" if completed == self.eval_episodes else "running",1240 )1241 print(1242 f"eval progress: episodes={completed}/{self.eval_episodes}",1243 flush=True,1244 )12451246 summary = self._summarize_eval(episodes)1247 self._save_eval_results(episodes, summary, total_num_steps)1248 self._log_eval_scalars(summary, total_num_steps)1249 self._print_eval_summary(summary)1250 return summary12511252 @torch.no_grad()1253 def render(self):1254 """Visualize the env."""1255 envs = self.envs12561257 all_frames = []1258 for episode in range(self.config.render_episodes):1259 obs = envs.reset(episode)1260 if self.config.save_gifs:1261 # Maybe to fix1262 image = None #envs.render("rgb_array")[0][0]1263 all_frames.append(image)1264 else:1265 envs.render("human")12661267 rnn_states = np.zeros(1268 (1269 self.n_rollout_threads,1270 self.num_agents,1271 self.recurrent_N,1272 self.hidden_size,1273 ),1274 dtype=np.float32,1275 )1276 masks = np.ones((self.n_rollout_threads, self.num_agents, 1), dtype=np.float32)12771278 episode_rewards = []12791280 for step in range(self.episode_length):1281 calc_start = datetime.now()12821283 self.trainer.prep_rollout()1284 action, rnn_states = self.trainer.policy.act(1285 _concat_obs_dict(obs),1286 np.concatenate(rnn_states),1287 np.concatenate(masks),1288 deterministic=getattr(self.config, "eval_deterministic", True),1289 )1290 actions = np.array(np.split(_t2n(action), self.n_rollout_threads))1291 rnn_states = np.array(np.split(_t2n(rnn_states), self.n_rollout_threads))12921293 # Obser reward and next obs1294 obs, rewards, dones, infos, active_masks, bad_masks, available_agents = envs.step(actions)1295 episode_rewards.append(rewards)12961297 rnn_states[dones == True] = np.zeros(1298 ((dones == True).sum(), self.recurrent_N, self.hidden_size),1299 dtype=np.float32,1300 )1301 masks = np.ones((self.n_rollout_threads, self.num_agents, 1), dtype=np.float32)1302 masks[dones == True] = np.zeros(((dones == True).sum(), 1), dtype=np.float32)13031304 if self.config.save_gifs:1305 # Maybe to fix1306 image = None #envs.render("rgb_array")[0][0]1307 all_frames.append(image)1308 calc_end = datetime.now()1309 elapsed = calc_end - calc_start1310 if elapsed < self.config.ifi:1311 time.sleep(self.config.ifi - elapsed)1312 else:1313 envs.render("human")13141315 print("average episode rewards is: " + str(np.mean(np.sum(np.array(episode_rewards), axis=0))))13161317 # if self.all_args.save_gifs:1318 # imageio.mimsave(str(self.gif_dir) + '/render.gif', all_frames, duration=self.all_args.ifi)