988 lines
43 KiB
Python
988 lines
43 KiB
Python
"""高层 MaskablePPO 训练（P1）。
|
||
|
||
架构与边界：
|
||
- 环境 = :class:`mcm26b.hrl.hierarchical_env.HierarchicalEnv`，每一步是一个高层
|
||
宏动作（选择 0–119 槽位：clear/locate/station/fallback/保留）；
|
||
- 训练/部署共用 :mod:`mcm26b.hrl.scheduler` 的候选、观测和动作掩码；
|
||
- 低层默认使用几何选点；``--local-selector rl`` 仅作为接口预留并记录，
|
||
当前 HTTP 执行路径仍由高层宏动作终点决定，不把八方向推演当正式收益。
|
||
|
||
落盘保证（长训必需）：
|
||
- **急停可恢复**：``--recovery-every``（默认 2000 环境步）把一份滚动恢复点
|
||
``<模型目录>/checkpoints/<stem>_latest.zip`` 原子覆盖写；训练开始前先写一份
|
||
起步基线。因此即使进程被 SIGKILL / ``taskkill /F`` 强杀（没有任何 Python
|
||
收尾机会），也**最多只丢 ``--recovery-every`` 个环境步**，且磁盘上始终有一份
|
||
完整、可加载、可 ``--resume`` 的产物；
|
||
- **原子落盘**：模型 zip、旁挂 json 与 ``<run_dir>/resume.json`` 都是"临时文件 +
|
||
``os.replace``"，保存途中被杀不会留下半截文件，也不会破坏上一份完整快照；
|
||
- ``--checkpoint-every`` 每 N 个环境步在 ``<模型目录>/checkpoints/`` 另留一份历史
|
||
checkpoint（``<stem>_<步数>.zip``），只保留最近 ``--checkpoint-keep`` 个，供
|
||
返祖对比与多档恢复；
|
||
- ``<run_dir>/resume.json`` 是恢复的**权威指针**：无论上一刻是滚动恢复点、历史
|
||
checkpoint、优雅收尾还是异常兜底，它总指向一份可续训产物，并记录增量语义
|
||
（``timesteps`` 是追加步数）与生效过的 PPO 超参；
|
||
- ``--control-file`` 是训练前端的**跨平台**控制通道：前端原子写 JSON
|
||
（``paused`` / ``stop``），本脚本每个 step 轮询。暂停时训练在 step 边界挂起，
|
||
停止时先保存模型与元数据再退出；
|
||
- 信号（SIGTERM/SIGINT，Windows 上另有 SIGBREAK）仍可用：直接 kill 本脚本或
|
||
前端发进程组信号时同样走优雅收尾；subproc worker 在 :func:`_make_env` 里忽略
|
||
这些信号，保证它们活到主进程 ``vec_env.close()`` 之后再退出；
|
||
- 任何异常都会先落一份 ``*_emergency.zip`` + ``resume.json`` 再抛出，不再出现
|
||
"跑了很久什么都没有"。
|
||
|
||
为什么暂停/停止不靠信号：Windows 没有 SIGSTOP，且 ``os.kill(pid, SIGTERM)`` 会
|
||
退化成 ``TerminateProcess`` 硬杀——收尾分支根本不会执行。控制文件两端都只是
|
||
普通文件 IO，POSIX 与 Windows 行为一致（见 :mod:`mcm26b.process_control`）。
|
||
|
||
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 math
|
||
import os
|
||
import platform
|
||
import random
|
||
import signal
|
||
import sys
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from mcm26b.config import ProblemConfig
|
||
from mcm26b.hrl.curriculum import STAGES
|
||
from mcm26b.hrl.constants import (ACTION_LAYOUT_VERSION, FEATURE_VERSION,
|
||
HIGH_OBS_DIM, N_CANDIDATES, REWARD_VERSION)
|
||
from mcm26b.process_control import ControlState, read_control
|
||
|
||
|
||
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 _stop_signals() -> list[int]:
|
||
"""本平台可用于"请求收尾"的信号集合。
|
||
|
||
Windows 上没有 SIGTERM 的软语义（``os.kill`` 直接 TerminateProcess），但
|
||
``SIGBREAK`` 可以由 ``CTRL_BREAK_EVENT`` 触发并进入 Python 处理器；POSIX 上
|
||
SIGTERM/SIGINT 均可捕获。控制文件才是主通道，这里只是让"直接 kill 脚本"
|
||
也尽量走优雅收尾。
|
||
"""
|
||
signals = [signal.SIGTERM, signal.SIGINT]
|
||
sigbreak = getattr(signal, "SIGBREAK", None)
|
||
if sigbreak is not None:
|
||
signals.append(sigbreak)
|
||
return signals
|
||
|
||
|
||
def _ignore_stop_signals_in_worker() -> None:
|
||
"""subproc worker 忽略收尾信号，保证主进程能优雅落盘。
|
||
|
||
无论信号来自进程组 SIGTERM 还是 Windows ``CTRL_BREAK_EVENT``，若 worker
|
||
同时退出，主进程会在保存模型之前就拿到断管/EOF 错误，整局训练白跑。worker
|
||
活在管道后面，主进程收尾时通过 ``vec_env.close()`` 的 ``close`` 指令正常
|
||
结束它们。
|
||
"""
|
||
try:
|
||
import multiprocessing
|
||
if multiprocessing.current_process().name == "MainProcess":
|
||
return
|
||
for signum in _stop_signals():
|
||
signal.signal(signum, signal.SIG_IGN)
|
||
except (ImportError, ValueError, OSError): # 非主线程/平台不支持：不影响训练
|
||
return
|
||
|
||
|
||
class _ControlPoll:
|
||
"""控制文件轮询状态机（纯标准库，可脱离 SB3 单测）。
|
||
|
||
``poll()`` 每 ``poll_every`` 次调用才真正读盘，避免每个 env step 都做一次
|
||
文件 IO；``wait_while_paused()`` 阻塞到暂停解除或收到 stop，并累计暂停秒数
|
||
（``--max-wall-s`` 要把这段扣掉，否则"挂起一小时回来就判超时"）。
|
||
"""
|
||
|
||
def __init__(self, path: Path, *, poll_every: int = 32,
|
||
poll_s: float = 0.2):
|
||
self.path = Path(path)
|
||
self.poll_every = max(1, int(poll_every))
|
||
self.poll_s = max(0.02, float(poll_s))
|
||
self.paused_total_s = 0.0
|
||
self._calls = 0
|
||
self._state = ControlState()
|
||
|
||
def poll(self) -> ControlState:
|
||
"""按抽样间隔返回最近一次控制状态。"""
|
||
self._calls += 1
|
||
if self._calls % self.poll_every:
|
||
return self._state
|
||
self._state = read_control(self.path)
|
||
return self._state
|
||
|
||
def wait_while_paused(self, on_pause, on_resume) -> bool:
|
||
"""阻塞到继续或停止；返回 True 表示期间收到了 stop。
|
||
|
||
``on_pause`` / ``on_resume`` 用来向前端回显事件（训练脚本的 JSON 行）。
|
||
"""
|
||
on_pause()
|
||
started = time.monotonic()
|
||
stopped = False
|
||
while True:
|
||
time.sleep(self.poll_s)
|
||
self._state = read_control(self.path)
|
||
if self._state.stop:
|
||
stopped = True
|
||
break
|
||
if not self._state.paused:
|
||
break
|
||
self.paused_total_s += time.monotonic() - started
|
||
on_resume()
|
||
return stopped
|
||
|
||
|
||
def _make_env(problem: int, stage_index: 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
|
||
|
||
_ignore_stop_signals_in_worker()
|
||
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,
|
||
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 _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 _progress_payload(model, start: int) -> dict[str, Any]:
|
||
"""组装 ``--progress-json`` 的单行负载。
|
||
|
||
纯函数（只读 model 的 logger / ep_info_buffer），便于脱离训练循环测试；
|
||
NaN / ±inf 一律折算成 None，保证 ``json.dumps(allow_nan=False)`` 不会炸。
|
||
"""
|
||
values = getattr(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 not math.isfinite(value):
|
||
return default
|
||
return value
|
||
|
||
ep_rewards: list[float] = []
|
||
ep_lengths: list[float] = []
|
||
for item in getattr(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
|
||
return {
|
||
"event": "progress",
|
||
"timesteps": int(model.num_timesteps),
|
||
"start_timesteps": int(start),
|
||
"n_updates": int(getattr(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"),
|
||
}
|
||
|
||
|
||
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,
|
||
"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),
|
||
"recovery_every": int(getattr(args, "recovery_every", 0) or 0),
|
||
"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 _write_json_atomic(path: Path, payload: dict[str, Any]) -> None:
|
||
"""原子写 JSON：同目录 ``.tmp-`` 前缀临时文件 + ``os.replace``。
|
||
|
||
急停（SIGKILL / ``taskkill /F``）会绕过所有 Python 清理逻辑，保存途中被杀
|
||
是完全可能的。原子替换保证**真实文件名下要么是旧的完整文件、要么是新的
|
||
完整文件**，不会出现半截 JSON 让续训读到一半就崩。
|
||
"""
|
||
path = Path(path)
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
tmp = path.with_name(f".tmp-{path.name}")
|
||
tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2),
|
||
encoding="utf-8")
|
||
os.replace(tmp, path)
|
||
|
||
|
||
def _save_model_atomic(model, target: Path) -> Path:
|
||
"""把 SB3 模型原子地写到 ``target``，返回实际落盘路径。
|
||
|
||
先写 ``.tmp-<name>``（同样以 ``.zip`` 结尾，SB3 不会再加后缀），完整写完后
|
||
``os.replace`` 覆盖目标。于是：
|
||
|
||
- 保存途中被强杀 → 目标路径仍是上一份**完整**的快照，可继续复训；
|
||
- 保存成功 → 目标路径立即变成新快照，不存在"看起来存在但打不开"的窗口。
|
||
|
||
临时名带 ``.tmp-`` 前缀，因此不会被 ``*_ckpt_*.zip`` 的清理逻辑误当成正式
|
||
checkpoint。
|
||
"""
|
||
target = Path(target)
|
||
target.parent.mkdir(parents=True, exist_ok=True)
|
||
tmp = target.with_name(f".tmp-{target.name}")
|
||
model.save(str(tmp))
|
||
if not tmp.exists(): # SB3 对未知后缀可能自行补 .zip
|
||
candidate = tmp.with_name(tmp.name + ".zip")
|
||
if candidate.exists():
|
||
tmp = candidate
|
||
os.replace(tmp, target)
|
||
return target
|
||
|
||
|
||
def _write_resume_hint(*, run_dir: Path, model_path: Path, timesteps: int,
|
||
reason: str, model=None) -> dict[str, Any]:
|
||
"""把"从哪继续"写进 ``<run_dir>/resume.json``。
|
||
|
||
这是急停恢复的**权威指针**：无论上一刻是周期 checkpoint、滚动恢复点、
|
||
优雅收尾还是异常兜底，这个文件总指向一份完整、可 ``--resume`` 的产物，
|
||
并带上续训语义（``timesteps`` 是增量步数）与生效过的 PPO 超参。
|
||
"""
|
||
hint: dict[str, Any] = {
|
||
"event": "resume_hint",
|
||
"model": str(model_path),
|
||
"timesteps": int(timesteps),
|
||
"timesteps_semantics": "additional",
|
||
"reason": str(reason),
|
||
"created_unix_s": time.time(),
|
||
"resume_args": ["--resume", "--model", str(model_path)],
|
||
}
|
||
if model is not None:
|
||
hint["ppo_hyperparams"] = _policy_hyperparams(model)
|
||
_write_json_atomic(run_dir / "resume.json", hint)
|
||
return hint
|
||
|
||
|
||
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)
|
||
_write_json_atomic(model_path.with_suffix(".json"), metadata)
|
||
return metadata
|
||
|
||
|
||
def _prune_checkpoints(directory: Path, stem: str, keep: int) -> list[str]:
|
||
"""只保留最近 ``keep`` 个 checkpoint（zip+json 同删），返回被删的 zip 路径。
|
||
|
||
只匹配 ``<stem>_<步数>.zip``；滚动恢复点 ``<stem>_latest.zip`` 与临时文件
|
||
``.tmp-*.zip`` 都不匹配，因此永远不会被历史清理误删。
|
||
"""
|
||
files = sorted((path for path in directory.glob(f"{stem}_[0-9]*.zip")
|
||
if path.stem.rsplit("_", 1)[-1].isdigit()),
|
||
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]:
|
||
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, _, 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) -> 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 # 已在上面处理；此分支仅保持语义清晰
|
||
|
||
# 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)
|
||
|
||
# ------------------------------------------------------ 优雅暂停 / 停止
|
||
# 主通道是控制文件（跨平台）；信号是兜底：直接 kill 本进程时也走同一段收尾。
|
||
stop_state: dict[str, Any] = {"requested": False, "signal": None}
|
||
previous_handlers: dict[int, Any] = {}
|
||
|
||
def _request_stop(source: str) -> None:
|
||
if stop_state["requested"]:
|
||
return
|
||
stop_state["requested"] = True
|
||
stop_state["signal"] = source
|
||
print(json.dumps({"event": "stop_requested", "signal": source}), flush=True)
|
||
|
||
def _on_stop_signal(signum, _frame):
|
||
_request_stop(signal.Signals(signum).name)
|
||
|
||
for _signum in _stop_signals():
|
||
try:
|
||
previous_handlers[_signum] = signal.signal(_signum, _on_stop_signal)
|
||
except (ValueError, OSError): # 非主线程/平台不支持：不阻塞训练
|
||
pass
|
||
|
||
class _ControlChannel(BaseCallback):
|
||
"""把 :class:`_ControlPoll` 接到 SB3 的 step 回调上。
|
||
|
||
为什么不用信号：Windows 没有 SIGSTOP，且 ``os.kill(SIGTERM)`` 会退化成
|
||
``TerminateProcess`` 硬杀，收尾分支根本不会执行。控制文件在 POSIX 与
|
||
Windows 上都只是普通文件 IO，语义完全一致。
|
||
|
||
暂停是**协作式**的：在 ``_on_step`` 内阻塞轮询，训练因此停在 step 边界，
|
||
进程仍存活、显存仍占用，但不消耗算力；收到 stop 立刻转入收尾落盘。
|
||
"""
|
||
|
||
def __init__(self, path: Path, *, poll_every: int = 32,
|
||
poll_s: float = 0.2):
|
||
super().__init__()
|
||
self.poller = _ControlPoll(path, poll_every=poll_every,
|
||
poll_s=poll_s)
|
||
|
||
def paused_seconds(self) -> float:
|
||
return self.poller.paused_total_s
|
||
|
||
def _event(self, name: str) -> None:
|
||
print(json.dumps({"event": name,
|
||
"timesteps": int(self.model.num_timesteps)}),
|
||
flush=True)
|
||
|
||
def _on_step(self) -> bool:
|
||
state = self.poller.poll()
|
||
if state.stop:
|
||
_request_stop("control_file")
|
||
return False
|
||
if state.paused:
|
||
stopped = self.poller.wait_while_paused(
|
||
lambda: self._event("paused"),
|
||
lambda: self._event("resumed"))
|
||
if stopped:
|
||
_request_stop("control_file")
|
||
return False
|
||
return not stop_state["requested"]
|
||
return True
|
||
|
||
control_path = getattr(args, "control_file", None)
|
||
control_cb = _ControlChannel(Path(control_path)) if control_path else None
|
||
|
||
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, paused_clock=None):
|
||
super().__init__()
|
||
self.max_wall_s = max_wall_s
|
||
self.start = time.monotonic()
|
||
self.paused_clock = paused_clock
|
||
|
||
def paused_seconds(self) -> float:
|
||
return float(self.paused_clock()) if self.paused_clock else 0.0
|
||
|
||
def _on_step(self) -> bool:
|
||
elapsed = time.monotonic() - self.start - self.paused_seconds()
|
||
if elapsed >= self.max_wall_s:
|
||
print(json.dumps({"event": "wall_clock_stop",
|
||
"elapsed_s": elapsed}))
|
||
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))
|
||
# init/resume 时计数器从上一阶段继续，前端要用增量算进度
|
||
self.start = int(model.num_timesteps)
|
||
self.last_emit = self.start
|
||
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
|
||
payload = _progress_payload(self.model, self.start)
|
||
print(json.dumps(payload, ensure_ascii=False, allow_nan=False),
|
||
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):
|
||
"""可续训落盘：滚动恢复点 + 周期 checkpoint + resume.json 指针。
|
||
|
||
- **滚动恢复点** ``<stem>_latest.zip``：每 ``recovery_every`` 个环境步
|
||
原子覆盖写一次，训练开始前也先写一份。它保证"急停"（SIGKILL /
|
||
``taskkill /F``，没有任何 Python 清理机会）之后**最多只丢
|
||
``recovery_every`` 个环境步**，而且文件始终完整可加载。
|
||
- **周期 checkpoint** ``<stem>_<步数>.zip``：每 ``every_steps`` 步留一份，
|
||
只保留最近 ``keep`` 个，用于返祖对比与多档恢复。
|
||
- 每次落盘都刷新 ``<run_dir>/resume.json``：无论怎么停，都有一份权威的
|
||
"从哪继续"指针（含增量语义与生效过的 PPO 超参）。
|
||
"""
|
||
|
||
def __init__(self, every_steps: int, keep: int, recovery_every: int):
|
||
super().__init__()
|
||
self.every_steps = max(0, int(every_steps))
|
||
self.keep = max(1, int(keep))
|
||
self.recovery_every = max(0, int(recovery_every))
|
||
self.last_saved = int(model.num_timesteps)
|
||
self.last_recovery = int(model.num_timesteps)
|
||
self.saved: list[str] = []
|
||
self.latest_path = checkpoints_dir / f"{checkpoint_stem}_latest.zip"
|
||
|
||
def write_recovery(self, reason: str) -> None:
|
||
"""立刻写一份滚动恢复点；训练开始前调用即"起步基线"。
|
||
|
||
用闭包里的 ``model`` 而不是 ``self.model``：``BaseCallback.model``
|
||
要等 ``learn()`` 时才被赋值，而起步基线恰恰在 ``learn()`` 之前写。
|
||
"""
|
||
current = int(model.num_timesteps)
|
||
self.last_recovery = current
|
||
self._write_recovery(current, reason)
|
||
|
||
def _on_step(self) -> bool:
|
||
current = int(model.num_timesteps)
|
||
if (self.recovery_every > 0
|
||
and current - self.last_recovery >= self.recovery_every):
|
||
self.last_recovery = current
|
||
self._write_recovery(current, "interval")
|
||
if (self.every_steps > 0
|
||
and current - self.last_saved >= self.every_steps):
|
||
self.last_saved = current
|
||
self._write_interval(current)
|
||
return True
|
||
|
||
def _persist(self, target: Path, current: int, reason: str,
|
||
extra: dict[str, Any]) -> None:
|
||
"""原子落盘 zip + 元数据，并刷新 resume.json 指针。"""
|
||
_save_model_atomic(model, 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=model,
|
||
extra={"checkpoint_timesteps": current,
|
||
"checkpoint_reason": reason,
|
||
"checkpoint_of": str(model_path), **extra})
|
||
_write_resume_hint(run_dir=run_dir, model_path=target,
|
||
timesteps=current, reason=reason,
|
||
model=model)
|
||
|
||
def _write_recovery(self, current: int, reason: str) -> None:
|
||
checkpoints_dir.mkdir(parents=True, exist_ok=True)
|
||
self._persist(self.latest_path, current, reason,
|
||
{"recovery": True})
|
||
print(json.dumps({"event": "recovery_checkpoint",
|
||
"reason": reason, "timesteps": current,
|
||
"model": str(self.latest_path)},
|
||
ensure_ascii=False), flush=True)
|
||
|
||
def _write_interval(self, current: int) -> None:
|
||
checkpoints_dir.mkdir(parents=True, exist_ok=True)
|
||
target = checkpoints_dir / f"{checkpoint_stem}_{current:09d}.zip"
|
||
self._persist(target, current, "interval", {"checkpoint": True})
|
||
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)
|
||
checkpoint_cb = _CheckpointJSON(args.checkpoint_every, args.checkpoint_keep,
|
||
args.recovery_every)
|
||
if args.recovery_every > 0:
|
||
# 起步基线：训练还没跑就急停，也至少有一份可加载、可续训的产物
|
||
checkpoint_cb.write_recovery("baseline")
|
||
# 顺序有语义：控制通道最先（暂停/停止要在本步立刻生效），进度次之，
|
||
# 墙钟上限要能读到控制通道累计的暂停秒数，checkpoint 放最后。
|
||
callbacks: list[Any] = []
|
||
if control_cb is not None:
|
||
callbacks.append(control_cb)
|
||
if args.progress_json:
|
||
callbacks.append(_ProgressJSON(progress_every))
|
||
callbacks.append(_GracefulStop())
|
||
callbacks.append(_WallClockStop(
|
||
args.max_wall_s,
|
||
paused_clock=(control_cb.paused_seconds if control_cb is not None
|
||
else None)))
|
||
callbacks.append(checkpoint_cb)
|
||
|
||
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 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:
|
||
_save_model_atomic(model, 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"]})
|
||
_write_resume_hint(run_dir=run_dir, model_path=emergency,
|
||
timesteps=int(model.num_timesteps),
|
||
reason="emergency", model=model)
|
||
failure["emergency_model"] = str(emergency)
|
||
except Exception as save_exc: # 落盘失败不能再掩盖原始异常
|
||
failure["emergency_save_error"] = (
|
||
f"{type(save_exc).__name__}: {save_exc}")
|
||
_write_json_atomic(run_dir / "training_failure.json", failure)
|
||
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"])
|
||
_save_model_atomic(model, model_path)
|
||
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,
|
||
"recovery_model": str(checkpoint_cb.latest_path)})
|
||
# 最终模型也刷新一次 resume.json：优雅收尾后"从哪继续"指向最终模型
|
||
_write_resume_hint(run_dir=run_dir, model_path=model_path,
|
||
timesteps=actual_timesteps,
|
||
reason=("stopped" if stopped_by_signal else "finished"),
|
||
model=model)
|
||
metadata_path = model_path.with_suffix(".json")
|
||
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,
|
||
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)
|
||
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",
|
||
help="从 --model 续训（沿用 checkpoint 内的 PPO 超参）")
|
||
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("--control-file", default=None,
|
||
help="训练前端控制文件（JSON：paused/stop）；"
|
||
"跨平台暂停与优雅停止的主通道")
|
||
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("--recovery-every", type=int, default=2000,
|
||
help="每 N 个环境步原子覆盖写一份滚动恢复点 "
|
||
"<模型目录>/checkpoints/<stem>_latest.zip；"
|
||
"急停后最多丢 N 步。0=关闭（默认 2000）")
|
||
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:
|
||
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 不能为负")
|
||
if args.checkpoint_every < 0:
|
||
parser.error("checkpoint-every 不能为负")
|
||
if args.recovery_every < 0:
|
||
parser.error("recovery-every 不能为负")
|
||
if args.checkpoint_keep < 1:
|
||
parser.error("checkpoint-keep 至少为 1")
|
||
return run_training(args)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|