268 lines
11 KiB
Python
268 lines
11 KiB
Python
"""低层 PPO 训练：8 方向到达任务（BC 直线专家热启动 + PPO 微调）。
|
||
|
||
产物：
|
||
output/hrl/hrl_low.zip / hrl_low.json PPO 模型（部署用）
|
||
output/hrl/bc_low_expert.npz / bc_low.pt 行为克隆数据与网络
|
||
output/hrl/hrl_low_curve_env*.monitor.csv 每个环境独立学习曲线
|
||
|
||
评估口径：到达成功率 + 路径效率（直线距离/实际路径长，1 为完美直线）。
|
||
这是独立导航实验：目前低层 PPO 仅供轨迹诊断，不影响正式 HTTP 任务。
|
||
|
||
示例：
|
||
uv run python scripts/train_hrl_low.py --timesteps 60000
|
||
uv run python scripts/train_hrl_low.py --skip-bc # 纯 PPO 对照
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import hashlib
|
||
import json
|
||
import math
|
||
import time
|
||
from pathlib import Path
|
||
|
||
from mcm26b.config import ProblemConfig
|
||
from mcm26b.hrl.constants import (FEATURE_VERSION, LOW_MAX_STEPS,
|
||
LOW_NUM_ACTIONS, LOW_OBS_DIM)
|
||
from mcm26b.hrl.low_level_env import LowLevelEnv
|
||
|
||
HIDDEN = (64, 64)
|
||
|
||
|
||
def _bc_module(obs_dim: int, n_actions: int, hidden: tuple[int, ...]):
|
||
import torch
|
||
|
||
class BCMlp(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
layers = []
|
||
prev = obs_dim
|
||
for width in hidden:
|
||
layers += [torch.nn.Linear(prev, width), torch.nn.Tanh()]
|
||
prev = width
|
||
layers.append(torch.nn.Linear(prev, n_actions))
|
||
self.net = torch.nn.Sequential(*layers)
|
||
|
||
def forward(self, x):
|
||
return self.net(x)
|
||
|
||
return BCMlp
|
||
|
||
|
||
def collect(episodes: int, seed: int, config: ProblemConfig):
|
||
import numpy as np
|
||
env = LowLevelEnv(config=config, seed=seed)
|
||
observations, actions, episode_ids = [], [], []
|
||
for episode in range(episodes):
|
||
obs, _info = env.reset(seed=seed + episode)
|
||
for _step in range(LOW_MAX_STEPS + 2):
|
||
observations.append(obs)
|
||
actions.append(env.expert_action())
|
||
episode_ids.append(episode)
|
||
obs, _r, terminated, truncated, _info = env.step(env.expert_action())
|
||
if terminated or truncated:
|
||
break
|
||
env.close()
|
||
return (np.asarray(observations, dtype=np.float32),
|
||
np.asarray(actions, dtype=np.int64),
|
||
np.asarray(episode_ids, dtype=np.int64))
|
||
|
||
|
||
def bc_validation_mask(episode_ids, *, seed: int):
|
||
"""按整局留出，避免同一轨迹相邻状态跨训练/验证集。"""
|
||
import numpy as np
|
||
unique = np.unique(episode_ids)
|
||
if len(unique) < 2:
|
||
raise ValueError("BC 至少需要 2 局才能按整局划分训练/验证集")
|
||
rng = np.random.default_rng(seed)
|
||
validation_ids = rng.permutation(unique)[:max(1, len(unique) // 10)]
|
||
return np.isin(episode_ids, validation_ids)
|
||
|
||
|
||
def train_bc(x, y, episode_ids, *, epochs: int, batch_size: int, lr: float,
|
||
device: str, seed: int):
|
||
import numpy as np
|
||
import torch
|
||
BCMlp = _bc_module(LOW_OBS_DIM, LOW_NUM_ACTIONS, HIDDEN)
|
||
model = BCMlp().to(device)
|
||
optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
|
||
valid = torch.as_tensor(
|
||
bc_validation_mask(episode_ids, seed=seed), device=device)
|
||
x_t = torch.as_tensor(x, device=device)
|
||
y_t = torch.as_tensor(y, device=device)
|
||
x_v, y_v, x_tr, y_tr = x_t[valid], y_t[valid], x_t[~valid], y_t[~valid]
|
||
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())
|
||
if epoch == 0 or (epoch + 1) % 5 == 0:
|
||
with torch.no_grad():
|
||
valid_acc = float((model(x_v).argmax(1) == y_v).float().mean().item())
|
||
print(json.dumps({"event": "bc_epoch", "epoch": epoch + 1,
|
||
"loss": round(total / len(x_tr), 4),
|
||
"train_acc": round(correct / len(x_tr), 4),
|
||
"valid_acc": round(valid_acc, 4)},
|
||
ensure_ascii=False))
|
||
with torch.no_grad():
|
||
valid_acc = float((model(x_v).argmax(1) == y_v).float().mean().item())
|
||
return model, valid_acc
|
||
|
||
|
||
def copy_into_ppo(bc_model, device: str, env, *, seed: int,
|
||
n_steps: int = 512) -> object:
|
||
"""把 BC 权重拷入 MaskablePPO（policy_net 两层 + action_net），前向校验。"""
|
||
import torch
|
||
from sb3_contrib import MaskablePPO
|
||
ppo = MaskablePPO("MlpPolicy", env, policy_kwargs={"net_arch": list(HIDDEN)},
|
||
n_steps=n_steps, batch_size=128, n_epochs=6,
|
||
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=seed, device=device)
|
||
state = {key.removeprefix("net."): value
|
||
for key, value in bc_model.state_dict().items()}
|
||
policy_state = {key: value for key, value in state.items()
|
||
if key in ("0.weight", "0.bias", "2.weight", "2.bias")}
|
||
ppo.policy.mlp_extractor.policy_net.load_state_dict(policy_state)
|
||
ppo.policy.action_net.load_state_dict(
|
||
{"weight": state["4.weight"], "bias": state["4.bias"]})
|
||
probe = torch.randn((8, LOW_OBS_DIM), device=device)
|
||
with torch.no_grad():
|
||
copied = ppo.policy.action_net(ppo.policy.mlp_extractor.policy_net(probe))
|
||
if not torch.allclose(bc_model(probe), copied, atol=1e-6):
|
||
raise RuntimeError("行为克隆与 PPO 前向映射不一致")
|
||
return ppo
|
||
|
||
|
||
def evaluate(model, *, episodes: int, seed: int,
|
||
config: ProblemConfig) -> dict:
|
||
"""确定性策略评估：成功率与路径效率（对照直线）。"""
|
||
env = LowLevelEnv(config=config, seed=seed)
|
||
successes = 0
|
||
efficiencies = []
|
||
for episode in range(episodes):
|
||
obs, info = env.reset(seed=seed + episode)
|
||
start, goal = tuple(info["position"]), tuple(info["goal"])
|
||
straight = math.dist(start, goal)
|
||
steps = 0
|
||
for _ in range(LOW_MAX_STEPS + 2):
|
||
action, _state = model.predict(obs, deterministic=True)
|
||
obs, _r, terminated, truncated, info = env.step(int(action))
|
||
steps += 1
|
||
if terminated or truncated:
|
||
break
|
||
successes += int(terminated)
|
||
efficiencies.append(min(1.0, straight /
|
||
max(info["path_length_m"], 1e-9)))
|
||
env.close()
|
||
return {"episodes": episodes,
|
||
"success_rate": successes / episodes,
|
||
"mean_path_efficiency": sum(efficiencies) / len(efficiencies)}
|
||
|
||
|
||
def main(argv=None) -> int:
|
||
parser = argparse.ArgumentParser(description=__doc__)
|
||
parser.add_argument("--timesteps", type=int, default=60_000)
|
||
parser.add_argument("--seed", type=int, default=20260911)
|
||
parser.add_argument("--bc-episodes", type=int, default=600)
|
||
parser.add_argument("--bc-epochs", type=int, default=12)
|
||
parser.add_argument("--envs", type=int, default=4)
|
||
parser.add_argument("--rollout-steps", type=int, default=512,
|
||
help="每环境每次 PPO rollout 长度；短 smoke 可设 8")
|
||
parser.add_argument("--device", default="cpu")
|
||
parser.add_argument("--model", default="output/hrl/hrl_low.zip")
|
||
parser.add_argument("--eval-episodes", type=int, default=40)
|
||
parser.add_argument("--skip-bc", action="store_true",
|
||
help="跳过行为克隆（纯 PPO 对照实验用）")
|
||
parser.add_argument("--verbose", type=int, default=1)
|
||
args = parser.parse_args(argv)
|
||
if min(args.envs, args.timesteps, args.eval_episodes,
|
||
args.rollout_steps) <= 0:
|
||
parser.error("envs/timesteps/eval-episodes/rollout-steps 必须为正数")
|
||
if args.envs * args.rollout_steps < 2:
|
||
parser.error("PPO rollout 至少需要 2 个样本")
|
||
if not args.skip_bc and (args.bc_episodes < 2 or args.bc_epochs < 1):
|
||
parser.error("BC 至少需要 2 局和 1 个 epoch")
|
||
|
||
import torch
|
||
torch.set_num_threads(1)
|
||
from sb3_contrib import MaskablePPO
|
||
from stable_baselines3.common.monitor import Monitor
|
||
from stable_baselines3.common.vec_env import DummyVecEnv
|
||
from stable_baselines3.common.utils import set_random_seed
|
||
set_random_seed(args.seed)
|
||
|
||
model_path = Path(args.model)
|
||
model_path.parent.mkdir(parents=True, exist_ok=True)
|
||
curve_paths = [model_path.with_name(
|
||
f"{model_path.stem}_curve_env{rank}.monitor.csv")
|
||
for rank in range(args.envs)]
|
||
|
||
def _factory(rank: int):
|
||
def _make():
|
||
env = Monitor(LowLevelEnv(seed=args.seed + rank),
|
||
filename=str(curve_paths[rank]))
|
||
return env
|
||
return _make
|
||
|
||
vec_env = DummyVecEnv([_factory(rank) for rank in range(args.envs)])
|
||
started = time.monotonic()
|
||
bc_valid_acc = None
|
||
if args.skip_bc:
|
||
model = MaskablePPO("MlpPolicy", vec_env,
|
||
policy_kwargs={"net_arch": list(HIDDEN)},
|
||
n_steps=args.rollout_steps, batch_size=128, n_epochs=6,
|
||
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)
|
||
else:
|
||
x, y, episode_ids = collect(args.bc_episodes, args.seed, ProblemConfig())
|
||
bc_model, bc_valid_acc = train_bc(
|
||
x, y, episode_ids, epochs=args.bc_epochs, batch_size=256,
|
||
lr=3e-4, device=args.device, seed=args.seed)
|
||
model = copy_into_ppo(bc_model, args.device, vec_env, seed=args.seed,
|
||
n_steps=args.rollout_steps)
|
||
bc_dir = model_path.parent / "bc_low.pt"
|
||
torch.save(bc_model.state_dict(), str(bc_dir))
|
||
print(json.dumps({"event": "bc_done", "samples": int(len(x))},
|
||
ensure_ascii=False))
|
||
model.learn(total_timesteps=args.timesteps,
|
||
progress_bar=args.verbose > 0)
|
||
elapsed_s = time.monotonic() - started
|
||
|
||
model.save(str(model_path))
|
||
eval_seed = args.seed + 1_000_000
|
||
evaluation = evaluate(model, episodes=args.eval_episodes, seed=eval_seed,
|
||
config=ProblemConfig())
|
||
report = {
|
||
"level": "low", "timesteps": args.timesteps, "seed": args.seed,
|
||
"device": args.device, "wall_seconds": round(elapsed_s, 1),
|
||
"model": str(model_path), "algorithm": "maskable_ppo",
|
||
"feature_version": FEATURE_VERSION, "bc_skipped": args.skip_bc,
|
||
"bc_valid_acc": bc_valid_acc, "evaluation_seed": eval_seed,
|
||
"monitor_files": [str(path) for path in curve_paths],
|
||
"training_role": "standalone_navigation_only",
|
||
"obs_dim": LOW_OBS_DIM, "act_dim": LOW_NUM_ACTIONS,
|
||
"model_sha256": hashlib.sha256(model_path.read_bytes()).hexdigest(),
|
||
**evaluation,
|
||
}
|
||
model_path.with_suffix(".json").write_text(
|
||
json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
print(json.dumps(report, ensure_ascii=False, indent=2))
|
||
vec_env.close()
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|