429 lines
18 KiB
Python
429 lines
18 KiB
Python
"""HRL/基线统一评测入口（P1）。
|
||
|
||
特点：
|
||
- 同一 manifest 中的真值场景被所有 variant 复用，保证配对比较；
|
||
- 真值仅在测试端用于构造 ``LocalSimulator`` 与最终核对，策略端只收到协议响应；
|
||
- 支持 ``hrl_heuristic``（高层启发式降级）、``bayes``、``directional``；
|
||
- 支持已训练高层模型 ``hrl``（需 ``--hrl-model`` 或 manifest 中的 variant 配置）。
|
||
|
||
输出目录（默认拒绝覆盖已有结果）：
|
||
<output>/manifest.json
|
||
<output>/runs/<scene_id>/<variant>/summary.json
|
||
<output>/runs/<scene_id>/<variant>/session.jsonl
|
||
<output>/results.json
|
||
<output>/aggregate.json
|
||
<output>/aggregate.csv
|
||
|
||
示例：
|
||
python scripts/benchmark_hrl.py \
|
||
--variants hrl_heuristic,bayes,directional \
|
||
--manifest output/hrl/manifests/dev.json \
|
||
--output output/hrl/p1-heuristic
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import csv
|
||
import json
|
||
import statistics
|
||
import sys
|
||
import time
|
||
from collections import defaultdict
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from mcm26b.config import ProblemConfig
|
||
from mcm26b.harness import LocalSimTransport, mixed_scene, random_omni_scene
|
||
from mcm26b.harness.scenes import referee_spec
|
||
from mcm26b.protocol import SimulatorClient
|
||
from mcm26b.protocol.log import SessionLog
|
||
|
||
try: # 测试端依赖；运行时代码不导入
|
||
from referee import LocalSimulator, Source
|
||
except ImportError as exc: # pragma: no cover
|
||
raise SystemExit("benchmark_hrl 需要 dev 依赖 referee。请先 `uv sync`。") from exc
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# manifest / scene
|
||
# ---------------------------------------------------------------------------
|
||
@dataclass(frozen=True)
|
||
class SceneSpec:
|
||
scene_id: str
|
||
seed: int
|
||
error_seed: int
|
||
problem: int
|
||
sources: list[dict[str, Any]]
|
||
|
||
@property
|
||
def truth_channels(self) -> tuple[int, ...]:
|
||
return tuple(sorted(int(src["channel"]) for src in self.sources))
|
||
|
||
|
||
def _source_dict(source: Source) -> dict[str, Any]:
|
||
return {
|
||
"channel": int(source.channel),
|
||
"x": float(source.x),
|
||
"y": float(source.y),
|
||
"r_eff_m": float(source.r_eff_m),
|
||
"directional": bool(source.directional),
|
||
"heading_deg": float(source.heading_deg),
|
||
}
|
||
|
||
|
||
def _sources_to_dicts(sources: list[Source]) -> list[dict[str, Any]]:
|
||
return [_source_dict(source) for source in sources]
|
||
|
||
|
||
def _source_from_dict(data: dict[str, Any]) -> Source:
|
||
return Source(
|
||
channel=int(data["channel"]),
|
||
x=float(data["x"]),
|
||
y=float(data["y"]),
|
||
r_eff_m=float(data["r_eff_m"]),
|
||
directional=bool(data.get("directional", False)),
|
||
heading_deg=float(data.get("heading_deg", 0.0)),
|
||
)
|
||
|
||
|
||
def make_default_manifest(*, problem: int, scenes: int, sources: int | None,
|
||
directional: int, start_seed: int,
|
||
config: ProblemConfig) -> dict[str, Any]:
|
||
"""按 CLI 参数生成可复现的固定场景清单。"""
|
||
entries: list[dict[str, Any]] = []
|
||
for index in range(scenes):
|
||
seed = start_seed + index
|
||
if problem == 4:
|
||
scene_sources = mixed_scene(seed=seed, count=sources,
|
||
directional=directional, config=config)
|
||
else:
|
||
scene_sources = random_omni_scene(seed=seed, count=sources,
|
||
config=config)
|
||
entries.append({
|
||
"scene_id": f"scene-{seed}",
|
||
"seed": seed,
|
||
"error_seed": seed + 10_000,
|
||
"problem": problem,
|
||
"sources": _sources_to_dicts(scene_sources),
|
||
})
|
||
return {"version": 1, "problem": problem, "scenes": entries}
|
||
|
||
|
||
def load_or_generate_manifest(path: Path, *, problem: int, scenes: int,
|
||
sources: int | None, directional: int,
|
||
start_seed: int,
|
||
config: ProblemConfig) -> dict[str, Any]:
|
||
if path.exists():
|
||
manifest = json.loads(path.read_text(encoding="utf-8"))
|
||
else:
|
||
manifest = make_default_manifest(
|
||
problem=problem, scenes=scenes, sources=sources,
|
||
directional=directional, start_seed=start_seed, config=config)
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2),
|
||
encoding="utf-8")
|
||
if "scenes" not in manifest or not manifest["scenes"]:
|
||
raise ValueError(f"manifest 无 scenes: {path}")
|
||
for raw in manifest["scenes"]:
|
||
if not {"scene_id", "error_seed", "problem", "sources"} <= set(raw):
|
||
raise ValueError(f"manifest scene 字段缺失: {raw!r}")
|
||
if not raw["sources"]:
|
||
raise ValueError(f"scene {raw.get('scene_id')} 没有来源真值")
|
||
return manifest
|
||
|
||
|
||
def scenes_from_manifest(manifest: dict[str, Any]) -> list[SceneSpec]:
|
||
result: list[SceneSpec] = []
|
||
for raw in manifest["scenes"]:
|
||
result.append(SceneSpec(
|
||
scene_id=str(raw["scene_id"]),
|
||
seed=int(raw.get("seed", 0)),
|
||
error_seed=int(raw["error_seed"]),
|
||
problem=int(raw["problem"]),
|
||
sources=[dict(item) for item in raw["sources"]],
|
||
))
|
||
return result
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# variants
|
||
# ---------------------------------------------------------------------------
|
||
def _build_loop(variant: str, client: SimulatorClient, config: ProblemConfig,
|
||
problem: int, *, hrl_model: str | None,
|
||
hrl_device: str = "auto"):
|
||
"""延迟导入变体实现，避免 benchmark 顶层拉起重依赖。"""
|
||
if variant == "bayes":
|
||
from mcm26b.planner.bayes_loop import BayesLoop
|
||
return BayesLoop(client, config)
|
||
if variant == "directional":
|
||
from mcm26b.planner import DirectionalExtension
|
||
return DirectionalExtension(client, config)
|
||
if variant == "hrl_heuristic":
|
||
from mcm26b.planner.hrl_loop import HRLLoop
|
||
return HRLLoop(client, config, problem=problem, high_policy=None)
|
||
if variant == "hrl":
|
||
if not hrl_model:
|
||
raise RuntimeError("variant=hrl 需要 --hrl-model")
|
||
from mcm26b.hrl.policy import HighLevelPolicy
|
||
from mcm26b.planner.hrl_loop import HRLLoop
|
||
policy = HighLevelPolicy(hrl_model, device=hrl_device, problem=problem)
|
||
return HRLLoop(client, config, problem=problem, high_policy=policy)
|
||
raise ValueError(f"未知 variant: {variant!r}")
|
||
|
||
|
||
def _validity(summary: dict[str, Any], truth_channels: tuple[int, ...]) -> dict[str, Any]:
|
||
"""测试端有效性审计。真值只在此处进入。"""
|
||
cleared = set(int(c) for c in summary.get("cleared_channels", []))
|
||
absent = set(int(c) for c in summary.get("absent_channels", []))
|
||
truth = set(truth_channels)
|
||
declared_absent = truth & absent
|
||
cleared_matches_truth = cleared == truth
|
||
no_source_declared_absent = not declared_absent
|
||
completed = bool(summary.get("completed"))
|
||
anomalies = list(summary.get("anomalies", []))
|
||
# 本地 benchmark 要求：清洗真值且无守恒性异常；缺席集合不得含真值。
|
||
valid = completed and cleared_matches_truth and no_source_declared_absent \
|
||
and not anomalies
|
||
return {
|
||
"valid": valid,
|
||
"completed": completed,
|
||
"cleared_matches_truth": cleared_matches_truth,
|
||
"no_source_declared_absent": no_source_declared_absent,
|
||
"truth_channels": sorted(truth),
|
||
"declared_absent_truth": sorted(declared_absent),
|
||
"anomaly_count": len(anomalies),
|
||
}
|
||
|
||
|
||
def _row_from_summary(*, scene: SceneSpec, variant: str, summary: dict[str, Any],
|
||
wall_s: float) -> dict[str, Any]:
|
||
validity = _validity(summary, scene.truth_channels)
|
||
cleared_count = len(summary.get("cleared_channels", []))
|
||
virtual_s = float(summary.get("virtual_time_s") or 0.0)
|
||
avg_s = summary.get("avg_locate_clear_time_s")
|
||
totals = summary.get("step_totals", {}) or {}
|
||
return {
|
||
"scene_id": scene.scene_id,
|
||
"variant": variant,
|
||
"problem": scene.problem,
|
||
"truth_count": len(scene.truth_channels),
|
||
"cleared_count": cleared_count,
|
||
"virtual_time_s": virtual_s,
|
||
"avg_locate_clear_time_s": (float(avg_s) if avg_s is not None else None),
|
||
"accounted_virtual_time_s": summary.get("accounted_virtual_time_s"),
|
||
"program_run_time_s": round(wall_s, 6),
|
||
"real_elapsed_total_s": totals.get("real_elapsed_total_s"),
|
||
"total_requests": summary.get("total_requests"),
|
||
"measures": summary.get("measures"),
|
||
"clears": summary.get("clears"),
|
||
"clear_success": summary.get("clear_success"),
|
||
"completed": bool(summary.get("completed")),
|
||
"stop_reason": summary.get("stop_reason"),
|
||
"exit_reason": summary.get("exit_reason"),
|
||
**validity,
|
||
}
|
||
|
||
|
||
def _error_row(scene: SceneSpec, variant: str, exc: BaseException) -> dict[str, Any]:
|
||
return {
|
||
"scene_id": scene.scene_id,
|
||
"variant": variant,
|
||
"problem": scene.problem,
|
||
"truth_count": len(scene.truth_channels),
|
||
"cleared_count": 0,
|
||
"virtual_time_s": None,
|
||
"avg_locate_clear_time_s": None,
|
||
"accounted_virtual_time_s": None,
|
||
"program_run_time_s": None,
|
||
"real_elapsed_total_s": None,
|
||
"total_requests": None,
|
||
"measures": None,
|
||
"clears": None,
|
||
"clear_success": None,
|
||
"completed": False,
|
||
"stop_reason": None,
|
||
"exit_reason": None,
|
||
"valid": False,
|
||
"cleared_matches_truth": False,
|
||
"no_source_declared_absent": False,
|
||
"truth_channels": sorted(scene.truth_channels),
|
||
"declared_absent_truth": [],
|
||
"anomaly_count": 0,
|
||
"error": f"{type(exc).__name__}: {exc}",
|
||
}
|
||
|
||
|
||
def run_scene_variant(scene: SceneSpec, variant: str, *, config: ProblemConfig,
|
||
hrl_model: str | None, run_dir: Path,
|
||
hrl_device: str = "auto") -> dict[str, Any]:
|
||
"""运行一个 scene × variant。返回可序列化结果行与 summary。"""
|
||
run_dir.mkdir(parents=True, exist_ok=True)
|
||
sources = [_source_from_dict(raw) for raw in scene.sources]
|
||
simulator = LocalSimulator(
|
||
sources,
|
||
robot_id="hrl-benchmark",
|
||
spec=referee_spec(config),
|
||
error_seed=scene.error_seed,
|
||
stress_mode=False,
|
||
)
|
||
# 真值只用于 LocalSimulator；策略端仅持有 client/loop，不读取 simulator.sources。
|
||
log_path = run_dir / "session.jsonl"
|
||
log = SessionLog(path=log_path)
|
||
client = SimulatorClient("hrl-benchmark", config=config,
|
||
transport=LocalSimTransport(simulator), log=log)
|
||
loop = _build_loop(variant, client, config, scene.problem,
|
||
hrl_model=hrl_model, hrl_device=hrl_device)
|
||
started = time.monotonic()
|
||
try:
|
||
summary = loop.run()
|
||
finally:
|
||
wall_s = time.monotonic() - started
|
||
row = _row_from_summary(scene=scene, variant=variant, summary=summary,
|
||
wall_s=wall_s)
|
||
# 完整 summary 可能很大；仍写入 runs 目录，供复算与异常审计。
|
||
(run_dir / "summary.json").write_text(
|
||
json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
return {"row": row, "summary": summary}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# aggregate
|
||
# ---------------------------------------------------------------------------
|
||
def _aggregate(rows: list[dict[str, Any]]) -> dict[str, Any]:
|
||
grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||
for row in rows:
|
||
grouped[row["variant"]].append(row)
|
||
result: dict[str, Any] = {"variants": {}, "all_valid": True}
|
||
for variant, items in sorted(grouped.items()):
|
||
virtual = [float(item["virtual_time_s"]) for item in items
|
||
if item.get("virtual_time_s") is not None]
|
||
per_source = [float(item["avg_locate_clear_time_s"])
|
||
for item in items
|
||
if item.get("avg_locate_clear_time_s") is not None]
|
||
cleared = [int(item["cleared_count"] or 0) for item in items]
|
||
program = [float(item["program_run_time_s"]) for item in items
|
||
if item.get("program_run_time_s") is not None]
|
||
success = [1.0 if item.get("valid") else 0.0 for item in items]
|
||
result["variants"][variant] = {
|
||
"episodes": len(items),
|
||
"valid_episodes": int(sum(success)),
|
||
"valid_rate": (sum(success) / len(items)) if items else 0.0,
|
||
"success_rate": (sum(success) / len(items)) if items else 0.0,
|
||
"cleared_count_mean": statistics.fmean(cleared) if cleared else 0.0,
|
||
"cleared_count_std": (statistics.pstdev(cleared)
|
||
if len(cleared) > 1 else 0.0),
|
||
"virtual_time_mean_s": statistics.fmean(virtual) if virtual else 0.0,
|
||
"virtual_time_std_s": (statistics.pstdev(virtual)
|
||
if len(virtual) > 1 else 0.0),
|
||
"virtual_time_min_s": min(virtual) if virtual else 0.0,
|
||
"virtual_time_max_s": max(virtual) if virtual else 0.0,
|
||
"avg_locate_clear_time_mean_s": (statistics.fmean(per_source)
|
||
if per_source else None),
|
||
"program_run_time_mean_s": (statistics.fmean(program)
|
||
if program else None),
|
||
"total_requests_mean": statistics.fmean(
|
||
[int(item.get("total_requests") or 0) for item in items]),
|
||
}
|
||
if result["variants"][variant]["valid_rate"] < 1.0:
|
||
result["all_valid"] = False
|
||
return result
|
||
|
||
|
||
def main(argv=None) -> int:
|
||
parser = argparse.ArgumentParser(description=__doc__)
|
||
parser.add_argument("--variants", default="hrl_heuristic,bayes,directional",
|
||
help="逗号分隔：hrl_heuristic,bayes,directional,hrl")
|
||
parser.add_argument("--manifest", default="output/hrl/manifests/dev.json")
|
||
parser.add_argument("--output", default="output/hrl/p1-benchmark")
|
||
parser.add_argument("--problem", type=int, choices=(3, 4), default=3)
|
||
parser.add_argument("--scenes", type=int, default=6,
|
||
help="manifest 不存在时生成的场景数")
|
||
parser.add_argument("--sources", type=int, default=None,
|
||
help="每局源数；默认按题面 10~16 随机")
|
||
parser.add_argument("--directional", type=int, default=3,
|
||
help="问题4定向源数量（生成 manifest 时使用）")
|
||
parser.add_argument("--start-seed", type=int, default=91_000)
|
||
parser.add_argument("--hrl-model", default=None)
|
||
parser.add_argument("--hrl-device", choices=("auto", "cpu", "cuda", "mps"),
|
||
default="auto")
|
||
parser.add_argument("--force", action="store_true",
|
||
help="允许覆盖已有输出目录")
|
||
parser.add_argument("--log", action="store_true",
|
||
help="为兼容 CLI 保留；当前始终写 session.jsonl")
|
||
args = parser.parse_args(argv)
|
||
|
||
config = ProblemConfig()
|
||
output = Path(args.output)
|
||
if output.exists() and any(output.iterdir()) and not args.force:
|
||
print(f"输出目录已存在且非空：{output}；如需覆盖请加 --force。",
|
||
file=sys.stderr)
|
||
return 2
|
||
output.mkdir(parents=True, exist_ok=True)
|
||
manifest_path = Path(args.manifest)
|
||
manifest = load_or_generate_manifest(
|
||
manifest_path, problem=args.problem, scenes=args.scenes,
|
||
sources=args.sources, directional=args.directional,
|
||
start_seed=args.start_seed, config=config)
|
||
scenes = scenes_from_manifest(manifest)
|
||
variants = [item.strip() for item in args.variants.split(",") if item.strip()]
|
||
if not variants:
|
||
parser.error("--variants 不能为空")
|
||
if "hrl" in variants and not args.hrl_model:
|
||
# manifest 中也允许写明模型路径；此处先支持 CLI 显式路径。
|
||
print("variant=hrl 需要 --hrl-model（P1 冒烟模型路径）。", file=sys.stderr)
|
||
return 2
|
||
known = {"hrl_heuristic", "bayes", "directional", "hrl"}
|
||
unknown = sorted(set(variants) - known)
|
||
if unknown:
|
||
parser.error(f"未知 variant: {unknown}")
|
||
|
||
rows: list[dict[str, Any]] = []
|
||
runs_root = output / "runs"
|
||
wall_started = time.monotonic()
|
||
for scene in scenes:
|
||
for variant in variants:
|
||
run_dir = runs_root / scene.scene_id / variant
|
||
print(json.dumps({"event": "run_start", "scene": scene.scene_id,
|
||
"variant": variant}, ensure_ascii=False))
|
||
try:
|
||
result = run_scene_variant(
|
||
scene, variant, config=config, hrl_model=args.hrl_model,
|
||
run_dir=run_dir, hrl_device=args.hrl_device)
|
||
row = result["row"]
|
||
except Exception as exc: # 单场异常也进入聚合，all_valid=false
|
||
row = _error_row(scene, variant, exc)
|
||
rows.append(row)
|
||
print(json.dumps({"event": "run_done", **{
|
||
key: row.get(key) for key in
|
||
("scene_id", "variant", "cleared_count", "virtual_time_s",
|
||
"valid", "program_run_time_s", "error")}}, ensure_ascii=False))
|
||
aggregate = _aggregate(rows)
|
||
aggregate["wall_seconds"] = round(time.monotonic() - wall_started, 3)
|
||
aggregate["manifest"] = str(manifest_path)
|
||
aggregate["output"] = str(output)
|
||
aggregate["variants_requested"] = variants
|
||
aggregate["scenes"] = len(scenes)
|
||
(output / "results.json").write_text(
|
||
json.dumps(rows, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
(output / "manifest.json").write_text(
|
||
json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
(output / "aggregate.json").write_text(
|
||
json.dumps(aggregate, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
fieldnames = list(rows[0].keys())
|
||
for row in rows:
|
||
for key in row:
|
||
if key not in fieldnames:
|
||
fieldnames.append(key)
|
||
with (output / "aggregate.csv").open("w", newline="", encoding="utf-8") as handle:
|
||
writer = csv.DictWriter(handle, fieldnames=fieldnames)
|
||
writer.writeheader()
|
||
writer.writerows(rows)
|
||
print(json.dumps(aggregate, ensure_ascii=False, indent=2))
|
||
return 0 if aggregate.get("all_valid") else 1
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|