"""用自建模拟器的既有专家逻辑做行为克隆，初始化 PPO 策略网络。 思路：自建 referee 的观测模型与真值逻辑既快又可复现。先在它上面用 手写专家（Q3: worst-posterior 选点器；Q4: 遮挡梯子选点器）生成 (观测, 专家动作) 数据集，训练一个与 PPO policy_net 同构的 MLP 分类器， 再把权重拷入 PPO 模型作为初始化；随后 scripts/train_rl.py --init 继续 PPO 微调。这样 RL 从强先验出发，而不是从随机策略开始。 产物： output/rl/bc/q3_expert.npz output/rl/bc/q3_bc.pt output/rl/q3_ppo_bc_init.zip output/rl/bc/q4_expert.npz output/rl/bc/q4_bc.pt output/rl/q4_ppo_bc_init.zip """ from __future__ import annotations import argparse import json import math import time from pathlib import Path import numpy as np import torch from stable_baselines3 import PPO from mcm26b.config import ProblemConfig from mcm26b.planner import safe_disk_radius, select_measurement_point from mcm26b.rl._constants import NUM_ACTIONS, OBS_DIM from mcm26b.rl.env import LocalizationEnv from mcm26b.rl.features import candidate_actions HIDDEN = (128, 128) def expert_action(env: LocalizationEnv) -> int: """当前状态的专家动作编号（与训练环境候选同源，就近映射）。""" config = env.config channel = env._channel region = env._belief.candidate_set(channel) position = env._client.state.position blocked = env._blocked candidates = candidate_actions(region, position, config, blocked) if env.problem == 3: expert_point = select_measurement_point( region, position, config).position else: center, radius = region.mec() safe = safe_disk_radius(region, config) if not blocked: expert_point = select_measurement_point( region, position, config).position elif len(blocked) == 1 and safe >= 0: first = blocked[0] dx, dy = center[0] - first[0], center[1] - first[1] norm = math.hypot(dx, dy) if norm <= 1e-9: expert_point = (center[0] + safe, center[1]) else: scale = min(1.0, safe / norm) expert_point = (center[0] + dx * scale, center[1] + dy * scale) else: first = blocked[0] base = math.atan2(first[1] - center[1], first[0] - center[0]) angle = base + math.radians(45.0 * (len(blocked) - 1)) expert_point = (center[0] + radius * math.cos(angle), center[1] + radius * math.sin(angle)) return min(range(len(candidates)), key=lambda k: math.dist(candidates[k], expert_point)) class BCMlp(torch.nn.Module): def __init__(self): super().__init__() self.net = torch.nn.Sequential( torch.nn.Linear(OBS_DIM, HIDDEN[0]), torch.nn.ReLU(), torch.nn.Linear(HIDDEN[0], HIDDEN[1]), torch.nn.ReLU(), torch.nn.Linear(HIDDEN[1], NUM_ACTIONS)) def forward(self, x): return self.net(x) def collect(problem: int, episodes: int, seed: int, config: ProblemConfig): env = LocalizationEnv( config, problem=problem, directional_probability=1.0 if problem == 4 else 0.0, max_episode_steps=16) observations, actions = [], [] for episode in range(episodes): obs, info = env.reset(seed=seed + episode) for _step in range(env.max_episode_steps + 2): action = expert_action(env) observations.append(obs) actions.append(action) obs, _reward, terminated, truncated, _info = env.step(action) if terminated or truncated: break return np.asarray(observations, dtype=np.float32), \ np.asarray(actions, dtype=np.int64) def train_bc(x, y, *, epochs: int, batch_size: int, lr: float, device: str, valid_frac: float = 0.1): model = BCMlp().to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=lr) count = len(x) split = max(1, int(count * valid_frac)) perm = torch.randperm(count) x_t, y_t = torch.as_tensor(x, device=device), torch.as_tensor(y, device=device) x_v, y_v = x_t[perm[:split]], y_t[perm[:split]] x_tr, y_tr = x_t[perm[split:]], y_t[perm[split:]] loss_fn = torch.nn.CrossEntropyLoss() for epoch in range(epochs): order = torch.randperm(len(x_tr), device=device) total = correct = 0.0 for start in range(0, len(order), batch_size): idx = order[start:start + batch_size] logits = model(x_tr[idx]) loss = loss_fn(logits, y_tr[idx]) optimizer.zero_grad() loss.backward() optimizer.step() total += float(loss.item()) * len(idx) correct += int((logits.argmax(1) == y_tr[idx]).sum().item()) train_acc = correct / len(x_tr) with torch.no_grad(): valid_acc = float((model(x_v).argmax(1) == y_v).float().mean().item()) if epoch == 0 or (epoch + 1) % 5 == 0: print(json.dumps({"event": "bc_epoch", "epoch": epoch + 1, "loss": round(total / len(x_tr), 4), "train_acc": round(train_acc, 4), "valid_acc": round(valid_acc, 4)}, ensure_ascii=False)) return model, float(valid_acc) def copy_into_ppo(problem: int, model: BCMlp, device: str) -> Path | None: """把 BC 网络权重按 SB3 的 policy_net/action_net 拆分拷入 PPO。""" ppo_path = Path(f"output/rl/q{problem}_ppo.zip") if not ppo_path.is_file(): return None ppo = PPO.load(str(ppo_path), device=device) state = {key.removeprefix("net."): value for key, value in model.state_dict().items()} # SB3 MlpPolicy：policy_net 只含前两个 Linear（Tanh 激活）， # 输出层在 action_net。BC 的 Sequential 键为 0/2/4。 policy_state = {key: value for key, value in state.items() if key in ("0.weight", "0.bias", "2.weight", "2.bias")} action_state = { "weight": state["4.weight"], "bias": state["4.bias"]} ppo.policy.mlp_extractor.policy_net.load_state_dict(policy_state) ppo.policy.action_net.load_state_dict(action_state) init_path = Path(f"output/rl/q{problem}_ppo_bc_init.zip") ppo.save(str(init_path)) return init_path def main(argv=None) -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--problem", type=int, choices=(3, 4), required=True) parser.add_argument("--episodes", type=int, default=3000) parser.add_argument("--epochs", type=int, default=20) parser.add_argument("--batch-size", type=int, default=256) parser.add_argument("--lr", type=float, default=3e-4) parser.add_argument("--seed", type=int, default=20260911) parser.add_argument("--device", default="auto") parser.add_argument("--copy-only", action="store_true", help="跳过采集/训练，只把已有 q_bc.pt 拷入 PPO 初始化") args = parser.parse_args(argv) if args.device == "auto": device = "cuda" if torch.cuda.is_available() else "cpu" else: device = args.device print(json.dumps({"event": "bc_start", "problem": args.problem, "episodes": args.episodes, "device": device, "cuda": torch.cuda.is_available()}, ensure_ascii=False)) if args.copy_only: bc_dir = Path("output/rl/bc") checkpoint = torch.load(bc_dir / f"q{args.problem}_bc.pt", map_location="cpu") model = BCMlp() model.load_state_dict(checkpoint["state_dict"]) init_path = copy_into_ppo(args.problem, model, device) dataset = np.load(bc_dir / f"q{args.problem}_expert.npz", allow_pickle=False) report = { "problem": args.problem, "samples": int(len(dataset["obs"])), "valid_accuracy": checkpoint.get("valid_acc"), "dataset": str(bc_dir / f"q{args.problem}_expert.npz"), "bc_model": str(bc_dir / f"q{args.problem}_bc.pt"), "ppo_init": str(init_path) if init_path is not None else None, } (bc_dir / f"q{args.problem}_bc.json").write_text( json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8") print(json.dumps({"event": "bc_copy", "problem": args.problem, "ppo_init": str(init_path), **report}, ensure_ascii=False, indent=2)) return 0 if init_path is not None else 2 config = ProblemConfig() started = time.monotonic() x, y = collect(args.problem, args.episodes, args.seed, config) collect_s = time.monotonic() - started bc_dir = Path("output/rl/bc") bc_dir.mkdir(parents=True, exist_ok=True) np.savez_compressed(bc_dir / f"q{args.problem}_expert.npz", obs=x, action=y) model, valid_acc = train_bc(x, y, epochs=args.epochs, batch_size=args.batch_size, lr=args.lr, device=device) torch.save({"state_dict": model.state_dict(), "obs_dim": OBS_DIM, "actions": NUM_ACTIONS, "valid_acc": valid_acc}, bc_dir / f"q{args.problem}_bc.pt") # 拷入现有 PPO 的 policy_net/action_net 作为初始化（critic 保持原权重）。 init_path = copy_into_ppo(args.problem, model, device) report = { "problem": args.problem, "episodes": args.episodes, "samples": int(len(x)), "valid_accuracy": valid_acc, "collect_seconds": round(collect_s, 1), "train_seconds": round(time.monotonic() - started - collect_s, 1), "dataset": str(bc_dir / f"q{args.problem}_expert.npz"), "bc_model": str(bc_dir / f"q{args.problem}_bc.pt"), "ppo_init": str(init_path) if init_path is not None else None, } (bc_dir / f"q{args.problem}_bc.json").write_text( json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8") print(json.dumps(report, ensure_ascii=False, indent=2)) return 0 if __name__ == "__main__": raise SystemExit(main())