强化学习第一版
This commit is contained in:
@@ -0,0 +1,171 @@
|
||||
"""问题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 pathlib import Path
|
||||
|
||||
from stable_baselines3 import PPO
|
||||
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 = []
|
||||
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):
|
||||
action, _state = model.predict(obs, deterministic=True)
|
||||
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)
|
||||
return {
|
||||
"episodes": episodes,
|
||||
"success_rate": successes / episodes,
|
||||
"mean_steps": sum(steps) / episodes,
|
||||
"mean_episode_reward": total_reward / episodes,
|
||||
"successes": successes,
|
||||
}
|
||||
|
||||
|
||||
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)
|
||||
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
|
||||
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()
|
||||
vec_env = DummyVecEnv(
|
||||
[_make_env(args.problem, args.seed + rank, config,
|
||||
directional_probability)
|
||||
for rank in range(max(1, args.envs))])
|
||||
if args.init:
|
||||
model = PPO.load(args.init, env=vec_env, device=device)
|
||||
else:
|
||||
model = PPO(
|
||||
"MlpPolicy", vec_env,
|
||||
policy_kwargs={"net_arch": [128, 128]},
|
||||
n_steps=512,
|
||||
batch_size=128,
|
||||
n_epochs=6,
|
||||
gamma=0.99,
|
||||
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,
|
||||
**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))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user