BC训练

Env/bc_env.py

@@ -88,10 +88,48 @@ class BCScenarioEnv(MultiAgentScenarioEnv):
         self._spawn_controlled_agents()
         self._update_background_vehicles()
         obs = self._get_all_obs()
-        rewards = {aid: 0.0 for aid in self.controlled_agents}
+
+        # Reward shaping for evaluation/rollout monitoring (BC training itself doesn't use env reward).
+        speed_coef = float(self.config.get("reward_speed_coef", 0.05))
+        collision_distance = float(self.config.get("collision_distance", 6.0))
+        collision_penalty = float(self.config.get("collision_penalty", 100.0))
+
+        # Pre-collect all active vehicles (includes background vehicles).
+        active_agents = list(self.engine.agent_manager.active_agents.items())
+
+        rewards = {}
+        infos = {}
+        for aid, vehicle in self.controlled_agents.items():
+            # Speed reward
+            speed = getattr(vehicle, "speed", None)
+            if speed is None:
+                speed = float(np.linalg.norm(vehicle.velocity))
+            r_speed = speed_coef * float(speed)
+
+            # Near-collision penalty (distance-based, simulator-agnostic)
+            min_dist = float("inf")
+            for other_id, other_vehicle in active_agents:
+                if other_id == aid:
+                    continue
+                try:
+                    dist = float(np.linalg.norm(vehicle.position - other_vehicle.position))
+                except Exception:
+                    continue
+                if dist < min_dist:
+                    min_dist = dist
+
+            near_collision = bool(min_dist < collision_distance)
+            r_collision = -collision_penalty if near_collision else 0.0
+
+            rewards[aid] = float(r_speed + r_collision)
+            infos[aid] = {
+                "near_collision": near_collision,
+                "min_dist": (min_dist if np.isfinite(min_dist) else None),
+                "r_speed": float(r_speed),
+                "r_collision": float(r_collision),
+            }
         dones = {aid: False for aid in self.controlled_agents}
         dones["__all__"] = self.episode_step >= self.config["horizon"]
-        infos = {aid: {} for aid in self.controlled_agents}
         return obs, rewards, dones, infos
 
     def _get_all_obs(self):

Env/inverse_dynamics.py

@@ -2,7 +2,7 @@ import numpy as np
 import math
 
 class InverseDynamics:
-    def __init__(self, max_steering=0.7, max_acc=15.0, length=4.5):
+    def __init__(self, max_steering=0.7, max_acc=8.0, length=4.5):
         """
         :param max_steering: Max steering angle in radians (approx 40 degrees)
         :param max_acc: Max acceleration in m/s^2

README.md

@@ -93,6 +93,10 @@ python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir
 
 ### 2. 行为克隆 (BC)
 - **训练**：`python train_bc.py`（模型保存到 `models/bc/`，日志到 `logs/bc/`）
+```
+# 注意替换文件名
+python train_bc.py --expert_data_path ./data/training/expert_data_0_50.pkl --epochs 100
+```
 - **可视化**：`python scripts/visualize.py policy --policy_type bc --model_path models/bc/policy_best.pt`
 
 ### 3. 多智能体对抗模仿学习 (MAGAIL)

dataset/loader.py

@@ -10,8 +10,15 @@ import torch
 from torch.utils.data import Dataset
 
 
-def load_expert_pkl(expert_data_path):
-    """从目录或单个 pkl 加载专家 (obs, acts)，返回 concat 后的 obs_data, act_data。"""
+def load_expert_pkl(expert_data_path, *, filter_terminal_last_step: bool = False):
+    """从目录或单个 pkl 加载专家 (obs, acts)，返回 concat 后的 obs_data, act_data。
+
+    Args:
+        expert_data_path: Directory containing pkl files or a single pkl file.
+        filter_terminal_last_step: If True, drop the last (obs, act) pair of each trajectory.
+            This approximates II's \"train only on non-terminal steps\" when the dataset doesn't
+            explicitly store dones.
+    """
     if os.path.isdir(expert_data_path):
         pkl_files = glob.glob(os.path.join(expert_data_path, "*.pkl"))
         if not pkl_files:
@@ -30,12 +37,27 @@ def load_expert_pkl(expert_data_path):
             if isinstance(data, list):
                 for traj in data:
                     if "obs" in traj and "acts" in traj:
-                        obs_data.append(traj["obs"])
-                        act_data.append(traj["acts"])
+                        obs = traj["obs"]
+                        acts = traj["acts"]
+                        if filter_terminal_last_step and len(obs) > 0 and len(acts) > 0:
+                            # Drop last step of each trajectory
+                            obs = obs[:-1]
+                            acts = acts[:-1]
+                        if len(obs) == 0 or len(acts) == 0:
+                            continue
+                        obs_data.append(obs)
+                        act_data.append(acts)
             elif isinstance(data, dict):
                 if "observations" in data and "actions" in data:
-                    obs_data.append(data["observations"])
-                    act_data.append(data["actions"])
+                    obs = data["observations"]
+                    acts = data["actions"]
+                    if filter_terminal_last_step and len(obs) > 0 and len(acts) > 0:
+                        obs = obs[:-1]
+                        acts = acts[:-1]
+                    if len(obs) == 0 or len(acts) == 0:
+                        continue
+                    obs_data.append(obs)
+                    act_data.append(acts)
             else:
                 print(f"Skipping {pkl_file}: Unknown data format {type(data)}")
         except Exception as e:
@@ -79,7 +101,7 @@ def get_expert_scenario_ids(expert_data_path, max_ids=10):
 
 
 class MAGAILExpertDataset(Dataset):
-    def __init__(self, data_dir, transform=None):
+    def __init__(self, data_dir, transform=None, *, filter_terminal_last_step: bool = False):
         """
         Args:
             data_dir (str): Directory containing .pkl files from generate_expert_data.py
@@ -110,7 +132,10 @@ class MAGAILExpertDataset(Dataset):
             acts = traj["acts"]
 
             # obs: (T, 45), acts: (T, 2)
-            for i in range(len(obs)):
+            max_i = len(obs)
+            if filter_terminal_last_step and max_i > 0:
+                max_i -= 1
+            for i in range(max_i):
                 self.flat_data.append((obs[i], acts[i]))
 
         print(f"Total samples: {len(self.flat_data)}")

train_bc.py

@@ -35,6 +35,8 @@ def evaluate_policy(policy, args, device):
         scenario_ids = [0, 1, 2]
 
     total_rewards = []
+    total_steps = []
+    collision_episodes = 0
     horizon = 200
 
     for idx, scenario_id in enumerate(scenario_ids):
@@ -67,6 +69,7 @@ def evaluate_policy(policy, args, device):
 
         episode_reward = 0.0
         step_count = 0
+        had_near_collision = False
         dones = {"__all__": False}
         while not dones["__all__"] and step_count < horizon:
             step_count += 1
@@ -80,10 +83,18 @@ def evaluate_policy(policy, args, device):
                 actions, _ = policy.sample(obs_tensor)
                 actions = actions.cpu().numpy()
             action_dict = {aid: act for aid, act in zip(agent_ids, actions)}
-            obs_dict, rewards, dones, _ = env.step(action_dict)
+            obs_dict, rewards, dones, infos = env.step(action_dict)
             episode_reward += sum(rewards.values())
+            if infos:
+                for _aid, info in infos.items():
+                    if isinstance(info, dict) and info.get("near_collision", False):
+                        had_near_collision = True
+                        break
 
         total_rewards.append(episode_reward)
+        total_steps.append(step_count)
+        if had_near_collision:
+            collision_episodes += 1
         print(
             f"  Eval Episode {idx} (scenario {scenario_id}): Total Reward {episode_reward:.2f}, steps {step_count}, "
             f"agents (current): {n_controlled}, total in scenario: {n_total_in_scenario}"
@@ -92,10 +103,15 @@ def evaluate_policy(policy, args, device):
 
     if not total_rewards:
         print("  No valid eval episodes (all skipped or failed).")
-        return 0.0
+        return 0.0, 0.0, 0.0
     avg_reward = float(np.mean(total_rewards))
-    print(f"  Average Evaluation Reward: {avg_reward:.2f}")
-    return avg_reward
+    avg_steps = float(np.mean(total_steps)) if total_steps else 0.0
+    collision_rate = float(collision_episodes / max(1, len(total_rewards)))
+    print(
+        f"  Average Evaluation Reward: {avg_reward:.2f} | Mean Episode Length: {avg_steps:.1f} | "
+        f"Collision Rate (near): {collision_rate:.3f}"
+    )
+    return avg_reward, collision_rate, avg_steps
 
 
 def main(args):
@@ -108,7 +124,10 @@ def main(args):
     print(f"TensorBoard logging to: {log_dir}")
     os.makedirs(args.save_dir, exist_ok=True)
 
-    obs_data, act_data = load_expert_pkl(args.expert_data_path)
+    obs_data, act_data = load_expert_pkl(
+        args.expert_data_path,
+        filter_terminal_last_step=args.filter_terminal_last_step,
+    )
     obs_tensor = torch.FloatTensor(obs_data)
     act_tensor = torch.FloatTensor(act_data)
     dataset = TensorDataset(obs_tensor, act_tensor)
@@ -147,9 +166,15 @@ def main(args):
             best_val_loss = avg_val_loss
             torch.save(policy.state_dict(), os.path.join(args.save_dir, "policy_best.pt"))
 
+        # Periodic checkpointing (II-style)
+        if args.checkpoint_freq > 0 and (epoch + 1) % args.checkpoint_freq == 0:
+            torch.save(policy.state_dict(), os.path.join(args.save_dir, f"policy_epoch{epoch+1}.pt"))
+
         if (epoch + 1) % args.eval_freq == 0:
-            eval_reward = evaluate_policy(policy, args, device)
+            eval_reward, eval_collision_rate, eval_mean_steps = evaluate_policy(policy, args, device)
             writer.add_scalar("Reward/eval", eval_reward, epoch)
+            writer.add_scalar("Eval/collision_rate_near", eval_collision_rate, epoch)
+            writer.add_scalar("Eval/mean_episode_length", eval_mean_steps, epoch)
 
     torch.save(policy.state_dict(), os.path.join(args.save_dir, "policy_final.pt"))
     writer.close()
@@ -164,5 +189,11 @@ if __name__ == "__main__":
     parser.add_argument("--batch_size", type=int, default=64)
     parser.add_argument("--lr", type=float, default=3e-4)
     parser.add_argument("--eval_freq", type=int, default=10)
+    parser.add_argument("--checkpoint_freq", type=int, default=50, help="Save policy_epochN.pt every N epochs. Set <=0 to disable.")
+    parser.add_argument(
+        "--filter_terminal_last_step",
+        action="store_true",
+        help="Drop the last (obs, act) pair of each trajectory to approximate training on non-terminal steps (II-style).",
+    )
     args = parser.parse_args()
     main(args)