713 lines
30 KiB
Python
713 lines
30 KiB
Python
"""高层 MaskablePPO 训练（P1）。
|
||
|
||
架构与边界：
|
||
- 环境 = :class:`mcm26b.hrl.hierarchical_env.HierarchicalEnv`，每一步是一个高层
|
||
宏动作（选择 0–119 槽位：clear/locate/station/fallback/保留）；
|
||
- 训练/部署共用 :mod:`mcm26b.hrl.scheduler` 的候选、观测和动作掩码；
|
||
- 低层默认使用几何选点；``--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 \
|
||
--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 signal
|
||
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 _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):
|
||
"""顶层工厂：供 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, 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 _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
|
||
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,
|
||
"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),
|
||
"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_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]:
|
||
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 # 已在上面处理；此分支仅保持语义清晰
|
||
|
||
# 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__()
|
||
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))
|
||
# 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
|
||
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,
|
||
"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),
|
||
"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
|
||
|
||
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)
|
||
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))
|
||
|
||
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:
|
||
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 = _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")
|
||
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,
|
||
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",
|
||
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("--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:
|
||
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.checkpoint_keep < 1:
|
||
parser.error("checkpoint-keep 至少为 1")
|
||
return run_training(args)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|