13

scripts/train_rl.py

@@ -17,9 +17,12 @@ 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
@@ -46,11 +49,13 @@ def evaluate(model, problem: int, *, episodes: int = 24,
     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):
-            action, _state = model.predict(obs, deterministic=True)
+            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:
@@ -59,12 +64,16 @@ def evaluate(model, problem: int, *, episodes: int = 24,
         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),
     }
 
 
@@ -87,6 +96,8 @@ def main(argv=None) -> int:
                              "加载后继续 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
@@ -95,6 +106,7 @@ def main(argv=None) -> int:
         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:
@@ -111,20 +123,24 @@ def main(argv=None) -> int:
     }, 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:
-        model = PPO.load(args.init, env=vec_env, device=device)
+        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 = PPO(
+        model = model_cls(
             "MlpPolicy", vec_env,
             policy_kwargs={"net_arch": [128, 128]},
             n_steps=512,
             batch_size=128,
             n_epochs=6,
-            gamma=0.99,
+            gamma=1.0,
             gae_lambda=0.95,
             learning_rate=3e-4,
             clip_range=0.2,
@@ -158,12 +174,19 @@ def main(argv=None) -> int:
         "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