477 lines
19 KiB
Python
477 lines
19 KiB
Python
"""高层 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())
|