114514

scripts/train_hrl_high.py

@@ -0,0 +1,476 @@
+"""高层 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())