13

scripts/pretrain_rl_expert.py

@@ -73,8 +73,8 @@ 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(OBS_DIM, HIDDEN[0]), torch.nn.Tanh(),
+            torch.nn.Linear(HIDDEN[0], HIDDEN[1]), torch.nn.Tanh(),
             torch.nn.Linear(HIDDEN[1], NUM_ACTIONS))
 
     def forward(self, x):
@@ -86,30 +86,33 @@ def collect(problem: int, episodes: int, seed: int, config: ProblemConfig):
         config, problem=problem,
         directional_probability=1.0 if problem == 4 else 0.0,
         max_episode_steps=16)
-    observations, actions = [], []
+    observations, actions, groups = [], [], []
     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)
+            groups.append(episode)
             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)
+        np.asarray(actions, dtype=np.int64), np.asarray(groups, dtype=np.int64)
 
 
 def train_bc(x, y, *, epochs: int, batch_size: int, lr: float,
-             device: str, valid_frac: float = 0.1):
+             device: str, groups, 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)
+    episodes = np.unique(groups)
+    if len(episodes) < 2:
+        raise ValueError('行为克隆至少需要两个独立episode才能划分验证集')
+    valid_groups = np.random.default_rng(0).permutation(episodes)[:max(1, int(len(episodes)*valid_frac))]
+    valid = torch.as_tensor(np.isin(groups, valid_groups), device=device)
     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:]]
+    x_v, y_v = x_t[valid], y_t[valid]
+    x_tr, y_tr = x_t[~valid], y_t[~valid]
     loss_fn = torch.nn.CrossEntropyLoss()
     for epoch in range(epochs):
         order = torch.randperm(len(x_tr), device=device)
@@ -137,10 +140,10 @@ def train_bc(x, y, *, epochs: int, batch_size: int, lr: float,
 
 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)
+    # 当前奖励/特征已修订，重新初始化critic，不能沿用旧任务的价值预测。
+    env = LocalizationEnv(problem=problem)
+    ppo = PPO('MlpPolicy', env, policy_kwargs={'net_arch': list(HIDDEN)},
+              n_steps=512, batch_size=128, gamma=1.0, device=device)
     state = {key.removeprefix("net."): value
              for key, value in model.state_dict().items()}
     # SB3 MlpPolicy：policy_net 只含前两个 Linear（Tanh 激活），
@@ -151,8 +154,17 @@ def copy_into_ppo(problem: int, model: BCMlp, device: str) -> Path | None:
         "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)
+    # 数值检查整个前向映射，而不只是线性层形状。
+    probe = torch.randn((8, OBS_DIM), device=device)
+    with torch.no_grad():
+        copied = ppo.policy.action_net(ppo.policy.mlp_extractor.policy_net(probe))
+        if not torch.allclose(model(probe), copied, atol=1e-6):
+            raise RuntimeError('行为克隆与PPO激活函数/前向映射不一致')
     init_path = Path(f"output/rl/q{problem}_ppo_bc_init.zip")
     ppo.save(str(init_path))
+    env.close()
+    init_path.with_suffix('.json').write_text(json.dumps(
+        dict(algorithm='ppo', feature_version=2, problem=problem)), encoding='utf-8')
     return init_path
 
 
@@ -182,6 +194,8 @@ def main(argv=None) -> int:
         bc_dir = Path("output/rl/bc")
         checkpoint = torch.load(bc_dir / f"q{args.problem}_bc.pt",
                                 map_location="cpu")
+        if checkpoint.get('activation') != 'Tanh' or checkpoint.get('feature_version') != 2:
+            parser.error('旧ReLU/旧特征BC检查点不能直接拷贝，请重新采集并训练')
         model = BCMlp()
         model.load_state_dict(checkpoint["state_dict"])
         init_path = copy_into_ppo(args.problem, model, device)
@@ -204,17 +218,19 @@ def main(argv=None) -> int:
 
     config = ProblemConfig()
     started = time.monotonic()
-    x, y = collect(args.problem, args.episodes, args.seed, config)
+    torch.manual_seed(args.seed)
+    x, y, groups = 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)
+    np.savez_compressed(bc_dir / f"q{args.problem}_expert.npz", obs=x, action=y, episode=groups)
 
     model, valid_acc = train_bc(x, y, epochs=args.epochs,
                                 batch_size=args.batch_size, lr=args.lr,
-                                device=device)
+                                device=device, groups=groups)
     torch.save({"state_dict": model.state_dict(), "obs_dim": OBS_DIM,
-                "actions": NUM_ACTIONS, "valid_acc": valid_acc},
+                "actions": NUM_ACTIONS, "valid_acc": valid_acc,
+                "activation": "Tanh", "feature_version": 2},
                bc_dir / f"q{args.problem}_bc.pt")
 
     # 拷入现有 PPO 的 policy_net/action_net 作为初始化（critic 保持原权重）。