强化学习第一版
This commit is contained in:
@@ -0,0 +1,238 @@
|
||||
"""用自建模拟器的既有专家逻辑做行为克隆，初始化 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())
|
||||
Reference in New Issue
Block a user