holy fuking shit

scripts/pretrain_rl_expert.py

@@ -88,7 +88,7 @@ def collect(problem: int, episodes: int, seed: int, config: ProblemConfig):
         max_episode_steps=16)
     observations, actions, groups = [], [], []
     for episode in range(episodes):
-        obs, info = env.reset(seed=seed + episode)
+        obs, _info = env.reset(seed=seed + episode)
         for _step in range(env.max_episode_steps + 2):
             action = expert_action(env)
             observations.append(obs)
@@ -108,7 +108,8 @@ def train_bc(x, y, *, epochs: int, batch_size: int, lr: float,
     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_groups = np.random.default_rng(0).permutation(episodes)
+    valid_groups = valid_groups[: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[valid], y_t[valid]