195 lines
7.9 KiB
Python
195 lines
7.9 KiB
Python
"""问题3/4 强化学习选点策略训练（PPO + CUDA）。
|
||
|
||
训练环境 mcm26b.rl.env.LocalizationEnv 与运行闭环共享 features 观测/候选
|
||
布局。模型输出到 output/rl/q3_ppo.zip（问题3）或 q4_ppo.zip（问题4），
|
||
由 `mcm26b run --strategy rl` / `--q4-strategy rl` 加载。
|
||
|
||
示例：
|
||
uv sync --group rl --group dev
|
||
uv run python scripts/train_rl.py --problem 3 --timesteps 30000
|
||
uv run python scripts/train_rl.py --problem 4 --timesteps 30000
|
||
uv run mcm26b run --robot-id 2026001 --self-test --strategy rl
|
||
uv run mcm26b run --robot-id 2026001 --self-test --problem 4 --q4-strategy rl
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import time
|
||
from dataclasses import asdict
|
||
import hashlib
|
||
from pathlib import Path
|
||
|
||
from stable_baselines3 import PPO
|
||
from sb3_contrib import MaskablePPO
|
||
from stable_baselines3.common.vec_env import DummyVecEnv
|
||
|
||
from mcm26b.config import ProblemConfig
|
||
from mcm26b.rl.env import LocalizationEnv
|
||
|
||
|
||
def _make_env(problem: int, seed: int, config: ProblemConfig,
|
||
directional_probability: float):
|
||
def _factory():
|
||
env = LocalizationEnv(config, problem=problem,
|
||
directional_probability=directional_probability,
|
||
real_budget_s=3600.0)
|
||
env.reset(seed=seed)
|
||
return env
|
||
return _factory
|
||
|
||
|
||
def evaluate(model, problem: int, *, episodes: int = 24,
|
||
directional_probability: float = 0.5,
|
||
config: ProblemConfig = ProblemConfig()) -> dict:
|
||
"""独立场景评估：确定性策略的成功率、步数与平均回报。"""
|
||
env = LocalizationEnv(config, problem=problem,
|
||
directional_probability=directional_probability)
|
||
total_reward = 0.0
|
||
successes = 0
|
||
steps = []
|
||
costs = []
|
||
for episode in range(episodes):
|
||
obs, _info = env.reset(seed=1_000_000 + episode)
|
||
episode_reward = 0.0
|
||
for step in range(env.max_episode_steps + 2):
|
||
kwargs = {'action_masks': env.action_masks()} if isinstance(model, MaskablePPO) else {}
|
||
action, _state = model.predict(obs, deterministic=True, **kwargs)
|
||
obs, reward, terminated, truncated, info = env.step(int(action))
|
||
episode_reward += float(reward)
|
||
if terminated or truncated:
|
||
break
|
||
total_reward += episode_reward
|
||
cleared = bool(info.get("cleared"))
|
||
successes += int(cleared)
|
||
steps.append(step + 1)
|
||
costs.append(info['virtual_cost_s'])
|
||
env.close()
|
||
return {
|
||
"episodes": episodes,
|
||
"success_rate": successes / episodes,
|
||
"mean_steps": sum(steps) / episodes,
|
||
"mean_episode_reward": total_reward / episodes,
|
||
"successes": successes,
|
||
"mean_virtual_cost_s": sum(costs) / episodes,
|
||
"max_virtual_cost_s": max(costs),
|
||
}
|
||
|
||
|
||
def main(argv=None) -> int:
|
||
parser = argparse.ArgumentParser(description=__doc__)
|
||
parser.add_argument("--problem", type=int, choices=(3, 4), default=3,
|
||
help="训练问题：3=全向源选点，4=含定向遮挡梯子")
|
||
parser.add_argument("--timesteps", type=int, default=30000,
|
||
help="PPO 总环境步数")
|
||
parser.add_argument("--seed", type=int, default=20260911)
|
||
parser.add_argument("--envs", type=int, default=2, help="并行环境份数（CPU 侧）")
|
||
parser.add_argument("--directional-probability", type=float, default=None,
|
||
help="问题4训练场景中定向源比例（默认 1.0，先学遮挡分支）")
|
||
parser.add_argument("--device", default="auto",
|
||
help="auto=优先 CUDA；或 cuda/cpu")
|
||
parser.add_argument("--model", default=None,
|
||
help="模型输出路径（默认 output/rl/q<problem>_ppo.zip）")
|
||
parser.add_argument("--init", default=None,
|
||
help="初始模型路径（如 output/rl/q3_ppo_bc_init.zip；"
|
||
"加载后继续 PPO 微调）")
|
||
parser.add_argument("--eval-episodes", type=int, default=24)
|
||
parser.add_argument("--verbose", type=int, default=1)
|
||
parser.add_argument('--algorithm', choices=('ppo', 'maskable_ppo'), default='maskable_ppo')
|
||
parser.add_argument('--torch-threads', type=int, default=1)
|
||
args = parser.parse_args(argv)
|
||
directional_probability = (args.directional_probability
|
||
if args.directional_probability is not None
|
||
else (1.0 if args.problem == 4 else 0.0))
|
||
if not 0.0 <= directional_probability <= 1.0:
|
||
parser.error("--directional-probability 须在 [0,1]")
|
||
|
||
import torch
|
||
torch.set_num_threads(args.torch_threads)
|
||
if args.device == "auto":
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
else:
|
||
device = args.device
|
||
if device == "cuda" and not torch.cuda.is_available():
|
||
parser.error("--device cuda 但当前环境检测不到 CUDA GPU")
|
||
print(json.dumps({
|
||
"event": "train_start", "problem": args.problem,
|
||
"timesteps": args.timesteps, "device": device,
|
||
"directional_probability": directional_probability,
|
||
"torch": torch.__version__,
|
||
"gpu": torch.cuda.get_device_name(0) if device == "cuda" else None,
|
||
"cuda_available": torch.cuda.is_available(),
|
||
}, ensure_ascii=False))
|
||
|
||
config = ProblemConfig()
|
||
model_cls = MaskablePPO if args.algorithm == 'maskable_ppo' else PPO
|
||
vec_env = DummyVecEnv(
|
||
[_make_env(args.problem, args.seed + rank, config,
|
||
directional_probability)
|
||
for rank in range(max(1, args.envs))])
|
||
if args.init:
|
||
metadata = json.loads(Path(args.init).with_suffix('.json').read_text(encoding='utf-8'))
|
||
if metadata.get('algorithm', 'ppo') != args.algorithm or metadata.get('feature_version') != 2:
|
||
parser.error('初始化模型的算法/特征版本与当前训练不符，请重训或提供v2初始化模型')
|
||
model = model_cls.load(args.init, env=vec_env, device=device)
|
||
else:
|
||
model = model_cls(
|
||
"MlpPolicy", vec_env,
|
||
policy_kwargs={"net_arch": [128, 128]},
|
||
n_steps=512,
|
||
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=device,
|
||
verbose=args.verbose,
|
||
)
|
||
|
||
started = time.monotonic()
|
||
model.learn(total_timesteps=args.timesteps,
|
||
progress_bar=args.verbose > 0)
|
||
elapsed_s = time.monotonic() - started
|
||
|
||
model_path = Path(args.model) if args.model \
|
||
else Path("output/rl") / f"q{args.problem}_ppo.zip"
|
||
model_path.parent.mkdir(parents=True, exist_ok=True)
|
||
model.save(str(model_path))
|
||
|
||
evaluation = evaluate(model, args.problem, episodes=args.eval_episodes,
|
||
directional_probability=directional_probability,
|
||
config=config)
|
||
report = {
|
||
"problem": args.problem,
|
||
"timesteps": args.timesteps,
|
||
"seed": args.seed,
|
||
"device": device,
|
||
"directional_probability": directional_probability,
|
||
"wall_seconds": round(elapsed_s, 1),
|
||
"model": str(model_path),
|
||
"init": args.init,
|
||
"algorithm": args.algorithm,
|
||
"feature_version": 2,
|
||
"reward": 'negative_virtual_cost_div100_success20_failure100',
|
||
"gamma": 1.0,
|
||
"config": asdict(config),
|
||
"model_sha256": hashlib.sha256(model_path.read_bytes()).hexdigest(),
|
||
**evaluation,
|
||
}
|
||
report_path = model_path.with_suffix(".json")
|
||
report_path.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())
|