"""高层 MaskablePPO 训练（P1）。 架构与边界： - 环境 = :class:`mcm26b.hrl.hierarchical_env.HierarchicalEnv`，每一步是一个高层 宏动作（选择 0–119 槽位：clear/locate/station/fallback/保留）； - 训练/部署共用 :mod:`mcm26b.hrl.scheduler` 的候选、观测和动作掩码； - 低层默认使用几何选点；``--local-selector rl`` 仅作为接口预留并记录， 当前 HTTP 执行路径仍由高层宏动作终点决定，不把八方向推演当正式收益。 P1 smoke 示例： python scripts/train_hrl_high.py \ --problem 3 --stage 1 --timesteps 2048 --seed 320001 \ --envs 1 --vec-env dummy --device mps \ --config-profile production --model output/hrl/p1-smoke/q3.zip \ --eval-episodes 2 --verbose 0 """ from __future__ import annotations import argparse import hashlib import json import os import platform import random import sys import time from pathlib import Path from typing import Any from mcm26b.config import ProblemConfig from mcm26b.hrl.curriculum import STAGES, training_config from mcm26b.hrl.constants import (ACTION_LAYOUT_VERSION, FEATURE_VERSION, HIGH_OBS_DIM, N_CANDIDATES, REWARD_VERSION) def _stage_from_index(index: int): if not 1 <= int(index) <= len(STAGES): raise ValueError(f"stage 必须在 1..{len(STAGES)}") return STAGES[int(index) - 1] def _profile_fidelity(profile: str) -> str: if profile == "production": return "production" if profile == "coarse-pretrain": return "fast" raise ValueError(f"未知 config-profile: {profile}") def _reward_from_profile(profile: str): from mcm26b.hrl.hierarchical_env import HRLReward if profile == "time": return HRLReward.time_only() if profile == "prompt": return HRLReward() raise ValueError(f"未知 reward-profile: {profile}") def _make_env(problem: int, stage_index: int, seed: int, monitor_path: str | None, profile: str, reward_profile: str): """顶层工厂：供 DummyVecEnv / SubprocVecEnv 调用（必须是可 pickle 的）。""" from stable_baselines3.common.monitor import Monitor from mcm26b.hrl.hierarchical_env import HierarchicalEnv stage = _stage_from_index(stage_index) env = HierarchicalEnv( problem=problem, stage=stage, fidelity=_profile_fidelity(profile), max_episode_steps=300, real_budget_s=1800.0, reward=_reward_from_profile(reward_profile), ) return Monitor(env, filename=monitor_path) def _build_vec_env(args, monitor_paths: list[str]): from stable_baselines3.common.vec_env import (DummyVecEnv, SubprocVecEnv) from functools import partial factories = [ partial(_make_env, args.problem, args.stage, args.seed + rank, monitor_paths[rank], args.config_profile, args.reward_profile) for rank in range(args.envs) ] if args.vec_env == "subproc": # 每个 worker 独立 Monitor 文件；spawn 入口在 __main__ 保护内。 return SubprocVecEnv(factories, start_method="spawn") return DummyVecEnv(factories) def _hash_file(path: Path) -> str: return hashlib.sha256(path.read_bytes()).hexdigest() def _model_metadata(*, args, wall_s: float, n_envs: int, actual_timesteps: int) -> dict[str, Any]: import numpy import stable_baselines3 import sb3_contrib import torch stage = _stage_from_index(args.stage) return { "level": "high", "problem": int(args.problem), "stage": int(args.stage), "stage_name": stage.name, "algorithm": "maskable_ppo", "feature_version": FEATURE_VERSION, "action_layout_version": ACTION_LAYOUT_VERSION, "reward_version": REWARD_VERSION, "obs_dim": HIGH_OBS_DIM, "act_dim": N_CANDIDATES, "seed": int(args.seed), "requested_timesteps": int(args.timesteps), "timesteps_semantics": ("additional" if (args.resume or args.init) else "total"), "actual_timesteps": int(actual_timesteps), "n_envs": int(n_envs), "vec_env": args.vec_env, "config_profile": args.config_profile, "fidelity": _profile_fidelity(args.config_profile), "reward_profile": args.reward_profile, "reward_version": REWARD_VERSION, "local_selector": args.local_selector, "local_model": args.local_model, "device": args.device, "wall_seconds": round(wall_s, 3), "max_wall_s": float(args.max_wall_s), "created_unix_s": time.time(), "python": sys.version.split()[0], "platform": platform.platform(), "torch": torch.__version__, "numpy": numpy.__version__, "stable_baselines3": stable_baselines3.__version__, "sb3_contrib": sb3_contrib.__version__, "mps_available": bool(getattr(getattr(torch.backends, "mps", None), "is_available", lambda: False)()), } def _evaluate_env(model, *, problem: int, stage_index: int, episodes: int, seed: int, profile: str, reward_profile: str) -> dict[str, Any]: from mcm26b.hrl.hierarchical_env import HierarchicalEnv stage = _stage_from_index(stage_index) env = HierarchicalEnv(problem=problem, stage=stage, fidelity=_profile_fidelity(profile), max_episode_steps=300, real_budget_s=1800.0, reward=_reward_from_profile(reward_profile)) wins = 0 rows = [] if episodes <= 0: env.close() return {"episodes": 0, "successes": 0, "success_rate": None, "rows": []} for episode in range(episodes): obs, info = env.reset(seed=seed + episode) done = False steps = 0 while not done: action, _state = model.predict( obs, action_masks=env.action_masks(), deterministic=True) obs, reward, terminated, truncated, info = env.step(int(action)) done = bool(terminated or truncated) steps += 1 # 课程阶段的小场景只需证明全部频道 accounted；完整阶段至少应清掉 # 场景最少的题面源数。这里不读取真值，只使用环境自身统计。 success = bool( info.get("accounted", 0) >= 20 and info.get("cleared", 0) >= stage.sources_min and not info.get("anomalies") ) wins += int(success) rows.append({ "episode": episode, "seed": seed + episode, "steps": steps, "cleared": info.get("cleared", 0), "accounted": info.get("accounted", 0), "virtual_time_s": info.get("episode_virtual_cost_s"), "stop_reason": info.get("stop_reason"), "anomalies": info.get("anomalies"), "success": success, }) env.close() return { "episodes": len(rows), "successes": wins, "success_rate": wins / len(rows) if rows else 0.0, "rows": rows, } def _evaluate_manifest(model_path: Path, manifest_path: str, *, limit: int, problem: int) -> dict[str, Any]: """可选：用 benchmark_hrl 的同一套真值隔离评测在固定 manifest 上跑已训模型。""" scripts_dir = Path(__file__).resolve().parent if str(scripts_dir) not in sys.path: sys.path.insert(0, str(scripts_dir)) import benchmark_hrl as bench # scripts/benchmark_hrl.py manifest = json.loads(Path(manifest_path).read_text(encoding="utf-8")) scenes = bench.scenes_from_manifest(manifest) if limit > 0: scenes = scenes[:limit] rows = [] for scene in scenes: run_dir = Path(model_path).parent / "eval-manifest" / scene.scene_id result = bench.run_scene_variant( scene, "hrl", config=ProblemConfig(), hrl_model=str(model_path), run_dir=run_dir) rows.append(result["row"]) aggregate = bench._aggregate(rows) return aggregate def run_training(args) -> int: import numpy as np import torch from sb3_contrib import MaskablePPO from stable_baselines3.common.callbacks import BaseCallback from stable_baselines3.common.utils import set_random_seed torch.set_num_threads(1) set_random_seed(args.seed) random.seed(args.seed) np.random.seed(args.seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(args.seed) model_path = Path(args.model) if model_path.exists() and not (args.resume or args.force): print(f"模型已存在：{model_path}；如需覆盖请加 --force，续训请加 --resume。", file=sys.stderr) return 2 model_path.parent.mkdir(parents=True, exist_ok=True) run_dir = model_path.parent monitor_paths = [ str(run_dir / f"high_curve_{model_path.stem}_env{rank}.monitor.csv") for rank in range(args.envs) ] for path in monitor_paths: Path(path).parent.mkdir(parents=True, exist_ok=True) vec_env = _build_vec_env(args, monitor_paths) if args.resume: if not model_path.exists(): print(f"--resume 指定的模型不存在：{model_path}", file=sys.stderr) return 2 model = MaskablePPO.load(str(model_path), env=vec_env, device=args.device, print_system_info=False) elif args.init: init_path = Path(args.init) if not init_path.exists(): print(f"--init 指定的模型不存在：{init_path}", file=sys.stderr) return 2 model = MaskablePPO.load(str(init_path), env=vec_env, device=args.device, print_system_info=False) else: model = MaskablePPO( "MlpPolicy", vec_env, policy_kwargs={"net_arch": [128, 128]}, n_steps=args.rollout_steps, batch_size=args.batch_size, n_epochs=args.n_epochs, gamma=1.0, gae_lambda=0.95, learning_rate=3e-4, clip_range=0.2, ent_coef=0.005, vf_coef=0.5, max_grad_norm=0.5, seed=args.seed, device=args.device, verbose=args.verbose, ) if args.init: pass # 已在上面处理；此分支仅保持语义清晰 class _WallClockStop(BaseCallback): def __init__(self, max_wall_s: float): super().__init__() self.max_wall_s = max_wall_s self.start = time.monotonic() def _on_step(self) -> bool: if time.monotonic() - self.start >= self.max_wall_s: print(json.dumps({"event": "wall_clock_stop", "elapsed_s": time.monotonic() - self.start})) return False return True class _ProgressJSON(BaseCallback): """向训练前端输出单行 JSON 进度；不影响 SB3 自身 logger。 前端首个 curve 点越早出现越好，因此开头几步按更小间隔发送（前 8 个 进度点逐步放大到 ``every_steps``），之后稳定在 ``every_steps``。 这样短程 smoke 也能看到完整曲线，长训不会刷屏。 """ def __init__(self, every_steps: int): super().__init__() self.every_steps = max(1, int(every_steps)) self.last_emit = int(model.num_timesteps) self.emitted = 0 def _next_gap(self) -> int: warmup = max(0.2, 1.0 - self.emitted / 8.0) return max(1, min(self.every_steps, int(round(self.every_steps * warmup)))) def _on_step(self) -> bool: current = int(self.model.num_timesteps) if current - self.last_emit < self._next_gap(): return True self.last_emit = current self.emitted += 1 values = getattr(self.model.logger, "name_to_value", {}) or {} def _num(key, default=None): value = values.get(key) if value is None: return default try: value = float(value) except (TypeError, ValueError): return default if value != value or value in (float("inf"), float("-inf")): return default return value ep_rewards = [] ep_lengths = [] for item in getattr(self.model, "ep_info_buffer", []) or []: try: ep_rewards.append(float(item.get("r"))) ep_lengths.append(float(item.get("l"))) except (TypeError, ValueError, AttributeError): continue payload = { "event": "progress", "timesteps": current, "n_updates": int(getattr(self.model, "_n_updates", 0) or 0), "ep_rew_mean": (sum(ep_rewards) / len(ep_rewards) if ep_rewards else None), "ep_len_mean": (sum(ep_lengths) / len(ep_lengths) if ep_lengths else None), "value_loss": _num("train/value_loss"), "policy_gradient_loss": _num("train/policy_gradient_loss"), "approx_kl": _num("train/approx_kl"), "entropy_loss": _num("train/entropy_loss"), "explained_variance": _num("train/explained_variance"), "fps": _num("time/fps"), } print(json.dumps(payload, ensure_ascii=False, allow_nan=False), flush=True) return True progress_every = (args.progress_every if args.progress_every > 0 else args.rollout_steps * args.envs) callbacks = [_WallClockStop(args.max_wall_s)] if args.progress_json: callbacks.insert(0, _ProgressJSON(progress_every)) started = time.monotonic() actual_timesteps = int(model.num_timesteps) try: model.learn(total_timesteps=args.timesteps, callback=callbacks if len(callbacks) > 1 else callbacks[0], progress_bar=False, reset_num_timesteps=not (args.resume or args.init)) except Exception as exc: # 训练失败也必须保存失败上下文，不静默 failure = { "event": "training_failed", "error": f"{type(exc).__name__}: {exc}", "model": str(model_path), } (run_dir / "training_failure.json").write_text( json.dumps(failure, ensure_ascii=False, indent=2), encoding="utf-8") raise wall_s = time.monotonic() - started actual_timesteps = int(model.num_timesteps) model.save(str(model_path)) metadata = _model_metadata(args=args, wall_s=wall_s, n_envs=args.envs, actual_timesteps=actual_timesteps) metadata["model_sha256"] = _hash_file(model_path) metadata["monitor_files"] = monitor_paths metadata["reward_note"] = ( "reward_profile=time 使用 HRLReward.time_only（P0.4 纯时间口径）；" "reward_profile=prompt 使用原始 §5.2 塑形。模型元数据已固化口径。") metadata_path = model_path.with_suffix(".json") metadata_path.write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8") vec_env.close() evaluation = _evaluate_env( model, problem=args.problem, stage_index=args.stage, episodes=args.eval_episodes, seed=args.seed + 1_000_000, profile=args.config_profile, reward_profile=args.reward_profile) manifest_eval = None if args.eval_manifest: try: manifest_eval = _evaluate_manifest( model_path, args.eval_manifest, limit=args.eval_manifest_limit, problem=args.problem) except Exception as exc: manifest_eval = {"error": f"{type(exc).__name__}: {exc}"} metadata["evaluation"] = evaluation metadata["eval_manifest"] = manifest_eval metadata_path.write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8") print(json.dumps(metadata, ensure_ascii=False, indent=2)) return 0 def main(argv=None) -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--problem", type=int, choices=(3, 4), required=True) parser.add_argument("--stage", type=int, choices=(1, 2, 3, 4), default=4) parser.add_argument("--timesteps", type=int, required=True) parser.add_argument("--seed", type=int, default=320001) parser.add_argument("--envs", type=int, default=1) parser.add_argument("--vec-env", choices=("dummy", "subproc"), default="dummy") parser.add_argument("--device", choices=("cpu", "cuda", "mps", "auto"), default="auto") parser.add_argument("--config-profile", choices=("production", "coarse-pretrain"), default="production") parser.add_argument("--model", required=True) parser.add_argument("--resume", action="store_true") parser.add_argument("--init", default=None, help="可选：Q3→Q4 或继续预训练权重路径") parser.add_argument("--force", action="store_true") parser.add_argument("--eval-manifest", default=None) parser.add_argument("--eval-manifest-limit", type=int, default=0) parser.add_argument("--eval-episodes", type=int, default=5) parser.add_argument("--local-selector", choices=("geometry", "rl"), default="geometry") parser.add_argument("--local-model", default=None) parser.add_argument("--rollout-steps", type=int, default=256) parser.add_argument("--batch-size", type=int, default=256) parser.add_argument("--n-epochs", type=int, default=10) parser.add_argument("--max-wall-s", type=float, default=1200.0, help="训练墙钟硬上限；默认 20 分钟") parser.add_argument("--reward-profile", choices=("time", "prompt"), default="time", help="time=P0.4 纯时间口径；prompt=原始 §5.2 塑形对照") parser.add_argument("--progress-json", action="store_true", help="向训练前端输出单行 JSON 进度（UI 自动启用）") parser.add_argument("--progress-every", type=int, default=0, help="JSON 进度间隔环境步数；0=每个 rollout") parser.add_argument("--verbose", type=int, default=1) args = parser.parse_args(argv) if args.timesteps <= 0 or args.envs <= 0: parser.error("timesteps 与 envs 必须为正") if args.envs * args.rollout_steps < 2: parser.error("PPO rollout 至少需要 2 个样本") if args.eval_episodes < 0: parser.error("eval-episodes 不能为负") return run_training(args) if __name__ == "__main__": raise SystemExit(main())