239 lines
10 KiB
Python
239 lines
10 KiB
Python
"""用自建模拟器的既有专家逻辑做行为克隆，初始化 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<problem>_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())
|