返回代码索引

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)