1919810

scripts/train_hrl_high.py

@@ -7,6 +7,15 @@
 - 低层默认使用几何选点；``--local-selector rl`` 仅作为接口预留并记录，
   当前 HTTP 执行路径仍由高层宏动作终点决定，不把八方向推演当正式收益。
 
+落盘保证（长训必需）：
+- ``--checkpoint-every`` 每 N 个环境步在 ``<模型目录>/checkpoints/`` 落一个可续训
+  的 checkpoint（zip + 同名 json 元数据），只保留最近 ``--checkpoint-keep`` 个；
+- 训练前端的「停止」是对**进程组**发 SIGTERM，因此 subproc worker 在 :func:`_make_env`
+  里显式忽略 SIGTERM/SIGINT，主进程收到信号后在 step 边界收尾并保存模型，
+  不会被 worker 先行退出打断；
+- 任何异常（含被强杀之外的失败）都会先落一份 ``*_emergency.zip`` 再抛出，
+  不再出现"跑了很久什么都没有"。
+
 P1 smoke 示例：
     python scripts/train_hrl_high.py \
         --problem 3 --stage 1 --timesteps 2048 --seed 320001 \
@@ -23,6 +32,7 @@ import json
 import os
 import platform
 import random
+import signal
 import sys
 import time
 from pathlib import Path
@@ -58,6 +68,23 @@ def _reward_from_profile(profile: str):
     raise ValueError(f"未知 reward-profile: {profile}")
 
 
+def _ignore_stop_signals_in_worker() -> None:
+    """subproc worker 忽略 SIGTERM/SIGINT，保证主进程能优雅收尾落盘。
+
+    训练前端的「停止」向整个进程组发 SIGTERM；若 worker 同时退出，主进程会在
+    保存模型之前就拿到断管/EOF 错误，整局训练白跑。worker 活在管道后面，主进程
+    收尾时通过 ``vec_env.close()`` 的 ``close`` 指令正常结束它们。
+    """
+    try:
+        import multiprocessing
+        if multiprocessing.current_process().name == "MainProcess":
+            return
+        for signum in (signal.SIGTERM, signal.SIGINT):
+            signal.signal(signum, signal.SIG_IGN)
+    except (ImportError, ValueError, OSError):  # 非主线程/平台不支持：不影响训练
+        return
+
+
 def _make_env(problem: int, stage_index: int, seed: int,
               monitor_path: str | None, profile: str,
               reward_profile: str):
@@ -66,6 +93,7 @@ def _make_env(problem: int, stage_index: int, seed: int,
 
     from mcm26b.hrl.hierarchical_env import HierarchicalEnv
 
+    _ignore_stop_signals_in_worker()
     stage = _stage_from_index(stage_index)
     env = HierarchicalEnv(
         problem=problem,
@@ -98,6 +126,41 @@ def _hash_file(path: Path) -> str:
     return hashlib.sha256(path.read_bytes()).hexdigest()
 
 
+def _policy_hyperparams(model) -> dict[str, Any]:
+    """记录**实际生效**的 PPO 超参。
+
+    ``--resume``/``--init`` 时 ``MaskablePPO.load`` 会一并恢复 checkpoint 内的
+    n_steps/batch_size/n_epochs 等，命令行同名参数被忽略；把它们写进元数据，
+    事后才能核对一局到底用什么配置训的。
+    """
+
+    def _scalar(value):
+        # SB3 2.9 的 learning_rate/clip_range 可能是 Schedule 对象
+        current = getattr(value, "current_value", None)
+        if callable(current):
+            try:
+                value = current()
+            except Exception:
+                return str(value)
+        try:
+            return float(value)
+        except (TypeError, ValueError):
+            return str(value)
+
+    return {
+        "n_steps": int(getattr(model, "n_steps", 0) or 0),
+        "batch_size": int(getattr(model, "batch_size", 0) or 0),
+        "n_epochs": int(getattr(model, "n_epochs", 0) or 0),
+        "gamma": _scalar(getattr(model, "gamma", 0.0)),
+        "gae_lambda": _scalar(getattr(model, "gae_lambda", 0.0)),
+        "learning_rate": _scalar(getattr(model, "learning_rate", 0.0)),
+        "clip_range": _scalar(getattr(model, "clip_range", 0.0)),
+        "ent_coef": _scalar(getattr(model, "ent_coef", 0.0)),
+        "vf_coef": _scalar(getattr(model, "vf_coef", 0.0)),
+        "max_grad_norm": _scalar(getattr(model, "max_grad_norm", 0.0)),
+    }
+
+
 def _model_metadata(*, args, wall_s: float, n_envs: int,
                     actual_timesteps: int) -> dict[str, Any]:
     import numpy
@@ -130,6 +193,9 @@ def _model_metadata(*, args, wall_s: float, n_envs: int,
         "reward_version": REWARD_VERSION,
         "local_selector": args.local_selector,
         "local_model": args.local_model,
+        "checkpoint_every": int(getattr(args, "checkpoint_every", 0) or 0),
+        "checkpoint_keep": int(getattr(args, "checkpoint_keep", 0) or 0),
+        "checkpoint_dir": getattr(args, "checkpoint_dir", None),
         "device": args.device,
         "wall_seconds": round(wall_s, 3),
         "max_wall_s": float(args.max_wall_s),
@@ -145,6 +211,40 @@ def _model_metadata(*, args, wall_s: float, n_envs: int,
     }
 
 
+def _write_metadata(*, model_path: Path, args, wall_s: float, n_envs: int,
+                    actual_timesteps: int, monitor_paths: list[str],
+                    model=None,
+                    extra: dict[str, Any] | None = None) -> dict[str, Any]:
+    """写模型旁挂元数据（含 sha256），训练前端/审计都读它。"""
+    metadata = _model_metadata(args=args, wall_s=wall_s, n_envs=n_envs,
+                               actual_timesteps=actual_timesteps)
+    if model is not None:
+        metadata["ppo_hyperparams"] = _policy_hyperparams(model)
+    metadata["model_sha256"] = _hash_file(model_path)
+    metadata["monitor_files"] = list(monitor_paths)
+    metadata["reward_note"] = (
+        "reward_profile=time 使用 HRLReward.time_only（P0.4 纯时间口径）；"
+        "reward_profile=prompt 使用原始 §5.2 塑形。模型元数据已固化口径。")
+    if extra:
+        metadata.update(extra)
+    model_path.with_suffix(".json").write_text(
+        json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8")
+    return metadata
+
+
+def _prune_checkpoints(directory: Path, stem: str, keep: int) -> list[str]:
+    """只保留最近 ``keep`` 个 checkpoint（zip+json 同删），返回被删的 zip 路径。"""
+    files = sorted(directory.glob(f"{stem}_*.zip"),
+                   key=lambda path: path.stat().st_mtime, reverse=True)
+    removed: list[str] = []
+    for stale in files[max(1, int(keep)):]:
+        for candidate in (stale, stale.with_suffix(".json")):
+            if candidate.exists():
+                candidate.unlink()
+        removed.append(str(stale))
+    return removed
+
+
 def _evaluate_env(model, *, problem: int, stage_index: int,
                   episodes: int, seed: int, profile: str,
                   reward_profile: str) -> dict[str, Any]:
@@ -288,6 +388,52 @@ def run_training(args) -> int:
         if args.init:
             pass  # 已在上面处理；此分支仅保持语义清晰
 
+    # resume/init 会恢复 checkpoint 内的 PPO 超参，命令行同名三项目前被忽略；
+    # 显式说出差异，避免"以为换了 n_steps 其实没换"。
+    if args.resume or args.init:
+        effective = _policy_hyperparams(model)
+        requested = {"n_steps": int(args.rollout_steps),
+                     "batch_size": int(args.batch_size),
+                     "n_epochs": int(args.n_epochs)}
+        if {key: effective[key] for key in requested} != requested:
+            print(json.dumps({
+                "event": "hyperparams_from_checkpoint",
+                "effective": {key: effective[key] for key in requested},
+                "requested": requested,
+                "note": ("resume/init 沿用 checkpoint 内的 PPO 超参，"
+                         "命令行 rollout-steps/batch-size/n-epochs 被忽略")},
+                ensure_ascii=False), file=sys.stderr, flush=True)
+
+    # ------------------------------------------------------------ 优雅停止
+    # 训练前端「停止」= killpg(SIGTERM)：worker 已在 _make_env 里忽略该信号，
+    # 主进程在此装处理器，等当前 step 结束后收尾落盘。
+    stop_state: dict[str, Any] = {"requested": False, "signal": None}
+    previous_handlers: dict[int, Any] = {}
+
+    def _on_stop_signal(signum, _frame):
+        stop_state["requested"] = True
+        stop_state["signal"] = signal.Signals(signum).name
+        print(json.dumps({"event": "stop_requested",
+                          "signal": signal.Signals(signum).name}), flush=True)
+
+    for _signum in (signal.SIGTERM, signal.SIGINT):
+        try:
+            previous_handlers[_signum] = signal.signal(_signum, _on_stop_signal)
+        except (ValueError, OSError):  # 非主线程/平台不支持：不阻塞训练
+            pass
+
+    class _GracefulStop(BaseCallback):
+        """收到停止信号后在最近一次 step 边界结束 rollout，让模型能落盘。"""
+
+        def _on_step(self) -> bool:
+            if stop_state["requested"]:
+                print(json.dumps({"event": "graceful_stop",
+                                  "signal": stop_state["signal"],
+                                  "timesteps": int(self.model.num_timesteps)}),
+                      flush=True)
+                return False
+            return True
+
     class _WallClockStop(BaseCallback):
         def __init__(self, max_wall_s: float):
             super().__init__()
@@ -312,7 +458,9 @@ def run_training(args) -> int:
         def __init__(self, every_steps: int):
             super().__init__()
             self.every_steps = max(1, int(every_steps))
-            self.last_emit = int(model.num_timesteps)
+            # init/resume 时计数器从上一阶段继续，前端要用增量算进度
+            self.start = int(model.num_timesteps)
+            self.last_emit = self.start
             self.emitted = 0
 
         def _next_gap(self) -> int:
@@ -350,6 +498,7 @@ def run_training(args) -> int:
             payload = {
                 "event": "progress",
                 "timesteps": current,
+                "start_timesteps": self.start,
                 "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),
@@ -366,43 +515,119 @@ def run_training(args) -> int:
                   flush=True)
             return True
 
+    started = time.monotonic()
+    checkpoints_dir = (Path(args.checkpoint_dir) if args.checkpoint_dir
+                       else run_dir / "checkpoints")
+    checkpoint_stem = f"{model_path.stem}_ckpt"
+
+    class _CheckpointJSON(BaseCallback):
+        """每 N 个环境步落一个可续训 checkpoint（zip + 同名元数据）。
+
+        只保留最近 ``keep`` 个；checkpoint 与最终模型共用同名 json 契约，
+        因此 ``--resume <checkpoint.zip>`` 直接可用。
+        """
+
+        def __init__(self, every_steps: int, keep: int):
+            super().__init__()
+            self.every_steps = max(1, int(every_steps))
+            self.keep = max(1, int(keep))
+            self.last_saved = int(model.num_timesteps)
+            self.saved: list[str] = []
+
+        def _on_step(self) -> bool:
+            current = int(self.model.num_timesteps)
+            if current - self.last_saved < self.every_steps:
+                return True
+            self.last_saved = current
+            self._write(current)
+            return True
+
+        def _write(self, current: int) -> None:
+            checkpoints_dir.mkdir(parents=True, exist_ok=True)
+            target = checkpoints_dir / f"{checkpoint_stem}_{current:09d}.zip"
+            self.model.save(str(target))
+            _write_metadata(
+                model_path=target, args=args,
+                wall_s=time.monotonic() - started, n_envs=args.envs,
+                actual_timesteps=current, monitor_paths=monitor_paths,
+                model=self.model,
+                extra={"checkpoint": True, "checkpoint_timesteps": current,
+                       "checkpoint_of": str(model_path)})
+            removed = _prune_checkpoints(checkpoints_dir, checkpoint_stem,
+                                         self.keep)
+            self.saved.append(str(target))
+            print(json.dumps({"event": "checkpoint", "timesteps": current,
+                              "model": str(target), "pruned": removed},
+                             ensure_ascii=False), flush=True)
+
     progress_every = (args.progress_every if args.progress_every > 0
                       else args.rollout_steps * args.envs)
-    callbacks = [_WallClockStop(args.max_wall_s)]
+    checkpoint_cb = None
+    if args.checkpoint_every > 0:
+        checkpoint_cb = _CheckpointJSON(args.checkpoint_every,
+                                        args.checkpoint_keep)
+    callbacks = [_GracefulStop(), _WallClockStop(args.max_wall_s)]
+    if checkpoint_cb is not None:
+        callbacks.append(checkpoint_cb)
     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:  # 训练失败也必须保存失败上下文，不静默
+    except BaseException as exc:  # 训练失败也必须留下可续训产物，不静默
         failure = {
             "event": "training_failed",
             "error": f"{type(exc).__name__}: {exc}",
             "model": str(model_path),
+            "timesteps": int(model.num_timesteps),
         }
+        emergency = run_dir / f"{model_path.stem}_emergency.zip"
+        try:
+            model.save(str(emergency))
+            _write_metadata(model_path=emergency, args=args,
+                            wall_s=time.monotonic() - started,
+                            n_envs=args.envs,
+                            actual_timesteps=int(model.num_timesteps),
+                            monitor_paths=monitor_paths, model=model,
+                            extra={"emergency": True,
+                                   "error": failure["error"]})
+            failure["emergency_model"] = str(emergency)
+        except Exception as save_exc:  # 落盘失败不能再掩盖原始异常
+            failure["emergency_save_error"] = (
+                f"{type(save_exc).__name__}: {save_exc}")
         (run_dir / "training_failure.json").write_text(
             json.dumps(failure, ensure_ascii=False, indent=2), encoding="utf-8")
+        print(json.dumps(failure, ensure_ascii=False), file=sys.stderr,
+              flush=True)
         raise
     wall_s = time.monotonic() - started
     actual_timesteps = int(model.num_timesteps)
+    stopped_by_signal = bool(stop_state["requested"])
     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 = _write_metadata(
+        model_path=model_path, args=args, wall_s=wall_s, n_envs=args.envs,
+        actual_timesteps=actual_timesteps, monitor_paths=monitor_paths,
+        model=model,
+        extra={"stopped_by_signal": stopped_by_signal,
+               "stop_signal": stop_state["signal"],
+               "checkpoint_files": (checkpoint_cb.saved
+                                    if checkpoint_cb is not None else [])})
     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()
+    # 收尾阶段恢复默认处理器：再次 Ctrl+C / 停止可立即强杀，不再等待评估。
+    for _signum, _handler in previous_handlers.items():
+        try:
+            signal.signal(_signum, _handler)
+        except (ValueError, OSError):
+            pass
+    if stopped_by_signal:
+        print(json.dumps({"event": "saved_after_stop",
+                          "model": str(model_path),
+                          "timesteps": actual_timesteps}), flush=True)
 
     evaluation = _evaluate_env(
         model, problem=args.problem, stage_index=args.stage,
@@ -439,7 +664,8 @@ def main(argv=None) -> int:
                         choices=("production", "coarse-pretrain"),
                         default="production")
     parser.add_argument("--model", required=True)
-    parser.add_argument("--resume", action="store_true")
+    parser.add_argument("--resume", action="store_true",
+                        help="从 --model 续训（沿用 checkpoint 内的 PPO 超参）")
     parser.add_argument("--init", default=None,
                         help="可选：Q3→Q4 或继续预训练权重路径")
     parser.add_argument("--force", action="store_true")
@@ -461,6 +687,12 @@ def main(argv=None) -> int:
                         help="向训练前端输出单行 JSON 进度（UI 自动启用）")
     parser.add_argument("--progress-every", type=int, default=0,
                         help="JSON 进度间隔环境步数；0=每个 rollout")
+    parser.add_argument("--checkpoint-every", type=int, default=20000,
+                        help="每 N 个环境步落一个 checkpoint；0=关闭（默认 20000）")
+    parser.add_argument("--checkpoint-keep", type=int, default=3,
+                        help="checkpoint 保留最近 K 个（默认 3）")
+    parser.add_argument("--checkpoint-dir", default=None,
+                        help="checkpoint 目录；默认 <模型目录>/checkpoints")
     parser.add_argument("--verbose", type=int, default=1)
     args = parser.parse_args(argv)
     if args.timesteps <= 0 or args.envs <= 0:
@@ -469,6 +701,10 @@ def main(argv=None) -> int:
         parser.error("PPO rollout 至少需要 2 个样本")
     if args.eval_episodes < 0:
         parser.error("eval-episodes 不能为负")
+    if args.checkpoint_every < 0:
+        parser.error("checkpoint-every 不能为负")
+    if args.checkpoint_keep < 1:
+        parser.error("checkpoint-keep 至少为 1")
     return run_training(args)