HBBC部署到代码中

.gitignore

@@ -1,3 +1,57 @@
 # 日志文件
 Env/logs/
 *.log
+
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# 虚拟环境
+venv/
+env/
+ENV/
+.venv
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# 数据和模型文件
+data/
+runs/
+*.pkl
+*.h5
+*.ckpt
+*.pth
+*.pt
+checkpoints/
+models/
+
+# 第三方库（如果已安装）
+metadrive/
+scenarionet/
+
+# 系统文件
+.DS_Store
+Thumbs.db

Algorithm/bc.py

@@ -0,0 +1,52 @@
+"""
+Behavior Cloning (BC) 算法：仅包含损失与单 epoch 训练/评估逻辑。
+数据加载、环境评估、日志与保存由训练脚本 (train_bc.py) 负责。
+"""
+import torch
+
+
+def bc_loss(policy, states, actions):
+    """
+    BC 损失：负对数似然 -E[log pi(a|s)]。
+    states: (B, state_dim), actions: (B, action_dim), 均在 policy 所在 device 上。
+    """
+    log_pi = policy.evaluate_log_pi(states, actions)
+    return -log_pi.mean()
+
+
+def train_bc_epoch(policy, train_loader, optimizer, device):
+    """
+    训练一个 epoch，返回平均 train loss。
+    policy 与 optimizer 由调用方管理，本函数只做前向、损失、反向与 step。
+    """
+    policy.train()
+    total_loss = 0.0
+    n_batches = 0
+    for states, actions in train_loader:
+        states = states.to(device)
+        actions = actions.to(device)
+        loss = bc_loss(policy, states, actions)
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        total_loss += loss.item()
+        n_batches += 1
+    return total_loss / n_batches if n_batches else 0.0
+
+
+def eval_bc_epoch(policy, val_loader, device):
+    """
+    在验证集上评估一个 epoch，返回平均 val loss（无梯度）。
+    """
+    policy.eval()
+    total_loss = 0.0
+    n_batches = 0
+    with torch.no_grad():
+        for states, actions in val_loader:
+            states = states.to(device)
+            actions = actions.to(device)
+            log_pi = policy.evaluate_log_pi(states, actions)
+            loss = -log_pi.mean().item()
+            total_loss += loss
+            n_batches += 1
+    return total_loss / n_batches if n_batches else 0.0

Env/bc_ego_replay_env.py

@@ -0,0 +1,194 @@
+"""
+Single-agent BC evaluation environment: only ego (SDC) is controlled by the policy;
+other vehicles are replayed from expert trajectories (same as data collection).
+"""
+import numpy as np
+from Env.expert_replay_env import ExpertReplayEnv
+from Env.hbbc_background_policy import HBBCBackgroundController
+
+
+class BCEgoReplayEnv(ExpertReplayEnv):
+    """
+    For single-agent BC evaluation: controlled_agents exposes only SDC (default_agent).
+    Other vehicles are still spawned and replayed by expert; internally we keep them
+    in _replay_agents so step() can update them.
+    """
+
+    def reset(self, seed=None):
+        obs = super().reset(seed=seed)
+        self.enable_hbbc_background = bool(self.config.get("enable_hbbc_background", False))
+        self.hbbc_controller = None
+        self._hbbc_runtime_logged = False
+        if self.enable_hbbc_background:
+            self.hbbc_controller = HBBCBackgroundController(
+                model_path=self.config.get("hbbc_model_path", "models/hbbc/hbbc.pt"),
+                device=self.config.get("hbbc_inference_device", "cpu"),
+                latent_mode=self.config.get("hbbc_latent_mode", "per_vehicle_fixed"),
+                latent_json_path=self.config.get("hbbc_latent_json_path"),
+                seed=int(self.config.get("seed", 0)),
+                dt=float(self.config.get("hbbc_dt", 0.1)),
+            )
+            self.hbbc_controller.reset_episode()
+        # Expose only SDC as the controlled agent for the evaluator
+        self._replay_agents = dict(self.controlled_agents)
+        if self.replay_sdc and self.sdc_vehicle is not None:
+            self.controlled_agents = {self.sdc_agent_id: self.sdc_vehicle}
+            self.controlled_agent_ids = [self.sdc_agent_id]
+        else:
+            self.controlled_agents = {}
+            self.controlled_agent_ids = []
+        return self._get_all_obs()
+
+    def _get_all_obs(self):
+        """Return only ego (SDC) observation so evaluator has a single agent."""
+        if not self.controlled_agents or self.sdc_vehicle is None:
+            return {}
+        obs = self._obs_for_vehicle(self.sdc_vehicle, exclude_agent_id=self.sdc_agent_id)
+        return {self.sdc_agent_id: obs}
+
+    def step(self, action_dict=None):
+        self.round += 1
+        expert_actions = {}
+        agents_to_remove = []
+
+        # SDC: use policy action if provided, else expert replay
+        if self.replay_sdc and self.sdc_vehicle is not None and self.sdc_track is not None:
+            policy_action = None
+            if action_dict and self.sdc_agent_id in action_dict:
+                policy_action = np.asarray(action_dict[self.sdc_agent_id], dtype=np.float64)
+            next_step = self.round
+            curr_step = self.round - 1
+            if next_step < len(self.sdc_track["state"]["position"]) and self.sdc_track["state"]["valid"][next_step]:
+                curr_state = {
+                    "position": self.sdc_track["state"]["position"][curr_step],
+                    "heading": self.sdc_track["state"]["heading"][curr_step],
+                    "velocity": self.sdc_track["state"]["velocity"][curr_step],
+                }
+                if policy_action is not None:
+                    next_state = self.inverse_dynamics.apply_action(curr_state, policy_action, dt=0.1)
+                    expert_actions[self.sdc_agent_id] = policy_action
+                else:
+                    next_state = {
+                        "position": self.sdc_track["state"]["position"][next_step],
+                        "heading": self.sdc_track["state"]["heading"][next_step],
+                        "velocity": self.sdc_track["state"]["velocity"][next_step],
+                    }
+                    action, _ = self.inverse_dynamics.compute_action(curr_state, next_state, dt=0.1)
+                    expert_actions[self.sdc_agent_id] = action
+                self.sdc_vehicle.set_position(next_state["position"])
+                self.sdc_vehicle.set_heading_theta(next_state["heading"])
+                self.sdc_vehicle.set_velocity(next_state["velocity"])
+                self.sdc_vehicle.last_expert_action = expert_actions[self.sdc_agent_id]
+
+        # Replay other vehicles: restore full controlled_agents for internal logic
+        self.controlled_agents = dict(self._replay_agents)
+        self.controlled_agent_ids = list(self.controlled_agents.keys())
+        hbbc_batch = []
+        hbbc_curr_states = {}
+        for agent_id, vehicle in self.controlled_agents.items():
+            track = vehicle.expert_track
+            next_step = self.round
+            if next_step >= len(track["state"]["position"]):
+                agents_to_remove.append(agent_id)
+                continue
+            if not track["state"]["valid"][next_step]:
+                agents_to_remove.append(agent_id)
+                continue
+            if self.enable_hbbc_background and self.hbbc_controller is not None:
+                # HBBC autonomous rollout: use vehicle's own previous-step state
+                curr_state = {
+                    "position": np.asarray(vehicle.position, dtype=np.float64),
+                    "heading": float(vehicle.heading_theta),
+                    "velocity": np.asarray(vehicle.velocity, dtype=np.float64),
+                }
+                object_id = str(getattr(vehicle, "original_id", agent_id))
+                hbbc_batch.append((agent_id, vehicle, object_id, agent_id))
+                hbbc_curr_states[agent_id] = curr_state
+            else:
+                curr_step = self.round - 1
+                curr_state = {
+                    "position": track["state"]["position"][curr_step],
+                    "heading": track["state"]["heading"][curr_step],
+                    "velocity": track["state"]["velocity"][curr_step],
+                }
+                next_state = {
+                    "position": track["state"]["position"][next_step],
+                    "heading": track["state"]["heading"][next_step],
+                    "velocity": track["state"]["velocity"][next_step],
+                }
+                action, _ = self.inverse_dynamics.compute_action(curr_state, next_state, dt=0.1)
+                expert_actions[agent_id] = action
+                vehicle.set_position(next_state["position"])
+                vehicle.set_heading_theta(next_state["heading"])
+                vehicle.set_velocity(next_state["velocity"])
+                vehicle.last_expert_action = action
+
+        if hbbc_batch and self.hbbc_controller is not None:
+            hbbc_actions = self.hbbc_controller.infer_actions(hbbc_batch)
+            if not self._hbbc_runtime_logged:
+                print(f"[HBBC] background policy active, current dynamic agents: {len(hbbc_batch)}")
+                self._hbbc_runtime_logged = True
+            for agent_id, _, _, _ in hbbc_batch:
+                curr_state = hbbc_curr_states[agent_id]
+                action = hbbc_actions[agent_id]
+                next_state = self.inverse_dynamics.apply_action(curr_state, action, dt=0.1)
+                expert_actions[agent_id] = action
+                vehicle = self.controlled_agents[agent_id]
+                vehicle.set_position(next_state["position"])
+                vehicle.set_heading_theta(next_state["heading"])
+                vehicle.set_velocity(next_state["velocity"])
+                try:
+                    vehicle.last_current_action.append(action)
+                except Exception:
+                    pass
+                vehicle.last_expert_action = action
+        for agent_id in agents_to_remove:
+            vehicle = self.controlled_agents[agent_id]
+            self.controlled_agents.pop(agent_id)
+            self.controlled_agent_ids.remove(agent_id)
+            self.engine.agent_manager.active_agents.pop(agent_id, None)
+            self.engine.clear_objects([vehicle.id])
+            if self.hbbc_controller is not None:
+                self.hbbc_controller.remove_vehicle(agent_id)
+        self.engine.taskMgr.step()
+        self._spawn_controlled_agents()
+        self._update_background_vehicles()
+        self._replay_agents = dict(self.controlled_agents)
+        # Expose only SDC again
+        if self.replay_sdc and self.sdc_vehicle is not None:
+            self.controlled_agents = {self.sdc_agent_id: self.sdc_vehicle}
+            self.controlled_agent_ids = [self.sdc_agent_id]
+        else:
+            self.controlled_agents = {}
+            self.controlled_agent_ids = []
+
+        obs = self._get_all_obs()
+        rewards = {}
+        infos = {aid: {"expert_action": expert_actions.get(aid, np.zeros(2))} for aid in self.controlled_agents}
+        if self.sdc_agent_id in self.controlled_agents and self.sdc_vehicle is not None:
+            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))
+            speed = float(np.linalg.norm(self.sdc_vehicle.velocity))
+            r_speed = speed_coef * speed
+            min_dist = float("inf")
+            for other_id, other_vehicle in self.engine.agent_manager.active_agents.items():
+                if other_id == self.sdc_agent_id:
+                    continue
+                try:
+                    d = float(np.linalg.norm(self.sdc_vehicle.position - other_vehicle.position))
+                    min_dist = min(min_dist, d)
+                except Exception:
+                    continue
+            near_collision = min_dist < collision_distance
+            r_collision = -collision_penalty if near_collision else 0.0
+            rewards[self.sdc_agent_id] = r_speed + r_collision
+            infos[self.sdc_agent_id].update(
+                near_collision=near_collision,
+                min_dist=min_dist if np.isfinite(min_dist) else None,
+                r_speed=r_speed,
+                r_collision=r_collision,
+            )
+        dones = {aid: False for aid in self.controlled_agents}
+        dones["__all__"] = self.round >= self.config["horizon"] or (len(self._replay_agents) == 0 and self.round > 190)
+        return obs, rewards, dones, infos

Env/bc_env.py

@@ -0,0 +1,246 @@
+from Env.scenario_env import MultiAgentScenarioEnv
+from Env.hbbc_background_policy import HBBCBackgroundController
+from Env.utils import filter_traffic_tracks_to_birth_lists
+from metadrive.component.vehicle.vehicle_type import DefaultVehicle
+import numpy as np
+
+
+class BCScenarioEnv(MultiAgentScenarioEnv):
+    """
+    Environment for Behavior Cloning Evaluation.
+    Uses the same 45-dim observation as ExpertReplayEnv:
+    - Ego State (5): x, y, vx, vy, heading
+    - Neighbors (40): 10 nearest * (rel_x, rel_y, vx, vy)
+
+    Spawns background (static) vehicles so that observation distribution matches expert data collection:
+    expert data is generated with ExpertReplayEnv which includes bg_* in active_agents, so the policy
+    was trained on obs that can include those neighbors. Demo should use the same scene for consistency.
+    """
+    def _init_hbbc_background(self):
+        self.enable_hbbc_background = bool(self.config.get("enable_hbbc_background", False))
+        self.hbbc_dynamic_agents = {}
+        self._spawned_dynamic_bg_ids = set()
+        self.hbbc_controller = None
+        if not self.enable_hbbc_background:
+            return
+        self.hbbc_controller = HBBCBackgroundController(
+            model_path=self.config.get("hbbc_model_path", "models/hbbc/hbbc.pt"),
+            device=self.config.get("hbbc_inference_device", "cpu"),
+            latent_mode=self.config.get("hbbc_latent_mode", "per_vehicle_fixed"),
+            latent_json_path=self.config.get("hbbc_latent_json_path"),
+            seed=int(self.config.get("seed", 0)),
+            dt=float(self.config.get("hbbc_dt", 0.1)),
+        )
+        self.hbbc_controller.reset_episode()
+
+    def _move_excess_controlled_to_hbbc_background(self):
+        if not self.enable_hbbc_background:
+            return
+        keep_n = int(self.config.get("num_controlled_agents", 0))
+        keep_n = max(0, keep_n)
+        ordered_ids = list(self.controlled_agents.keys())
+        keep_ids = set(ordered_ids[:keep_n])
+        move_ids = [aid for aid in ordered_ids if aid not in keep_ids]
+        for aid in move_ids:
+            self.hbbc_dynamic_agents[aid] = self.controlled_agents[aid]
+            self.controlled_agents.pop(aid, None)
+            if aid in self.controlled_agent_ids:
+                self.controlled_agent_ids.remove(aid)
+        self._spawned_dynamic_bg_ids.update(move_ids)
+
+    def _apply_hbbc_before_step(self):
+        if not self.enable_hbbc_background or not self.hbbc_dynamic_agents:
+            return
+        batch = []
+        for aid, vehicle in self.hbbc_dynamic_agents.items():
+            object_id = getattr(vehicle, "original_id", None) or aid.replace("controlled_", "", 1)
+            batch.append((aid, vehicle, str(object_id) if object_id is not None else None, aid))
+        actions = self.hbbc_controller.infer_actions(batch)
+        for aid, vehicle in self.hbbc_dynamic_agents.items():
+            action = actions.get(aid, np.zeros(2, dtype=np.float32))
+            vehicle.before_step(action)
+
+    def _apply_hbbc_after_step(self):
+        if not self.enable_hbbc_background:
+            return
+        for vehicle in self.hbbc_dynamic_agents.values():
+            vehicle.after_step()
+
+    def reset(self, seed=None):
+        self._init_hbbc_background()
+        # Clear background vehicles from previous episode so engine.reset() passes _object_clean_check
+        if getattr(self, "engine", None) is not None:
+            ids_bg = [
+                oid for oid, obj in self.engine.get_objects().items()
+                if (getattr(obj, "name", None) or getattr(obj, "id", None) or "").startswith(("bg_", "controlled_"))
+            ]
+            if ids_bg:
+                self.engine.clear_objects(ids_bg, force_destroy=True)
+            for aid in list(self.engine.agent_manager.active_agents.keys()):
+                if aid.startswith("bg_") or aid.startswith("controlled_"):
+                    self.engine.agent_manager.active_agents.pop(aid, None)
+        obs = super().reset(seed=seed)
+        self._move_excess_controlled_to_hbbc_background()
+        self._spawn_background_vehicles()
+        return self._get_all_obs()
+
+    def _build_birth_lists_from_traffic(self):
+        """Same lane/static filter as expert data; return background_vehicles so we spawn them (match training obs)."""
+        car_birth_info_list, background_vehicles, obj_to_clean, stats = filter_traffic_tracks_to_birth_lists(
+            self.engine.traffic_manager.current_traffic_data,
+            self.engine.traffic_manager.sdc_scenario_id,
+            self.engine.map_manager,
+            return_stats=True,
+        )
+        if stats["n_controlled"] == 0 and stats["n_total"] > 0:
+            print(
+                "[BCScenarioEnv] 0 controlled agents: total_vehicles={}, off_lane={}, static={}, no_valid={}.".format(
+                    stats["n_total"],
+                    stats["n_off_lane"],
+                    stats["n_static"],
+                    stats["n_no_valid"],
+                )
+            )
+        return car_birth_info_list, background_vehicles, obj_to_clean
+
+    def _spawn_background_vehicles(self):
+        """Spawn all static background vehicles once at reset (no show_time filter; same as ExpertReplayEnv)."""
+        for sid, car in self.background_vehicles.items():
+            bg_id = f"bg_{car['id']}"
+            if bg_id in self.engine.agent_manager.active_agents:
+                continue
+            vehicle_config = {}
+            if "length" in car and "width" in car:
+                vehicle_config = {"length": car["length"], "width": car["width"]}
+            v = self.engine.spawn_object(
+                DefaultVehicle,
+                name=bg_id,
+                vehicle_config=vehicle_config,
+                position=car["begin"],
+                heading=car["heading"],
+            )
+            v.set_velocity([0, 0])
+            self.engine.agent_manager.active_agents[bg_id] = v
+            v.valid_mask = car.get("valid")
+            v.start_t = car.get("show_time")
+
+    def _update_background_vehicles(self):
+        # Static vehicles are spawned once at init and never removed.
+        pass
+
+    def step(self, action_dict):
+        if action_dict is None:
+            action_dict = {}
+        self.round += 1
+        for agent_id, action in action_dict.items():
+            if agent_id in self.controlled_agents:
+                self.controlled_agents[agent_id].before_step(action)
+        self._apply_hbbc_before_step()
+        self.engine.step()
+        self.engine.after_step()
+        for agent_id in action_dict:
+            if agent_id in self.controlled_agents:
+                self.controlled_agents[agent_id].after_step()
+        self._spawn_controlled_agents()
+        self._move_excess_controlled_to_hbbc_background()
+        self._apply_hbbc_after_step()
+        self._update_background_vehicles()
+        obs = self._get_all_obs()
+
+        # 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"]
+        return obs, rewards, dones, infos
+
+    def _get_all_obs(self):
+        # Implement custom observation: 30m range, 10 nearest vehicles
+        obs_dict = {}
+        
+        for agent_id, vehicle in self.controlled_agents.items():
+            # 1. Ego State
+            ego_state = [
+                vehicle.position[0], vehicle.position[1],
+                vehicle.velocity[0], vehicle.velocity[1],
+                vehicle.heading_theta
+            ]
+            
+            # 2. Neighbors
+            neighbors = []
+            # Iterate through all vehicles in the engine
+            candidates = []
+            # Use engine.agent_manager.active_agents to find neighbors
+            # Note: This includes background vehicles if they are in active_agents
+            for other_id, other_vehicle in self.engine.agent_manager.active_agents.items():
+                if other_id == agent_id:
+                    continue
+                
+                # Check if vehicle is valid/active
+                # (MetaDrive manages active_agents, so they should be active)
+                
+                dist = np.linalg.norm(vehicle.position - other_vehicle.position)
+                if dist < 30.0:
+                    candidates.append((dist, other_vehicle))
+            
+            # Sort by distance
+            candidates.sort(key=lambda x: x[0])
+            
+            # Take top 10
+            top_10 = candidates[:10]
+            
+            neighbor_feats = []
+            for _, neighbor in top_10:
+                neighbor_feats.extend([
+                    neighbor.position[0] - vehicle.position[0], # Relative pos
+                    neighbor.position[1] - vehicle.position[1],
+                    neighbor.velocity[0], # Absolute vel
+                    neighbor.velocity[1]
+                ])
+                
+            # Pad if < 10
+            missing = 10 - len(top_10)
+            if missing > 0:
+                neighbor_feats.extend([0.0] * (4 * missing))
+                
+            # Flatten
+            obs = np.array(ego_state + neighbor_feats, dtype=np.float32)
+            obs_dict[agent_id] = obs
+            
+        return obs_dict

Env/expert_replay_env.py

@@ -35,132 +35,44 @@ class ExpertReplayEnv(MultiAgentScenarioEnv):
         if self.engine is None:
             raise ValueError("Broken MetaDrive instance.")
             
-        self.background_vehicles = {} # Vehicles that exist but are static/background
-        
-        # Helper function to check if a position is on a valid lane
-        def is_on_lane(pos, map_manager, threshold=2.0):
-            # Check if point is close to any lane in the road network
-            # This can be expensive if checked for every point, so we check sample points
-            # or rely on lane index if available.
-            # Waymo tracks don't have lane index, just positions.
-            # We can use map.road_network.get_closest_lane_index(pos)
-            if map_manager is None or map_manager.current_map is None:
-                return True # If no map, assume valid
-                
-            try:
-                # Use a larger search radius to catch slightly offset lanes
-                lane, lane_index = map_manager.current_map.road_network.get_closest_lane_index(pos, return_lane=True)
-                if lane is None:
-                    return False
-                
-                # Check lateral distance
-                long, lat = lane.local_coordinates(pos)
-                width = lane.width
-                # Allow being slightly off-lane (e.g. changing lanes)
-                # But parking lots are usually far from defined lanes in Waymo converted maps
-                if abs(lat) <= (width / 2 + threshold):
-                    return True
-                return False
-            except:
-                return False
-
-        # --- MODIFIED SECTION START ---
-        # Capture expert tracks before they are cleaned
+        self.background_vehicles = {}
         self.expert_tracks = {}
-        # Capture SDC track for ego replay (MetaDrive default agent)
         self.sdc_track = None
         self.sdc_vehicle = None
+
+        # 在加载新场景前，必须清除上一轮通过 spawn_object 生成的物体，否则 engine.reset() 内 _object_clean_check 会报错
+        # 从 engine 当前对象中按名称筛选（与 manager 无关的对象需在此清理），并强制销毁
+        ids_to_clear = []
+        for oid, obj in self.engine.get_objects().items():
+            name = getattr(obj, "name", None) or getattr(obj, "id", None)
+            if name and (str(name).startswith("controlled_") or str(name).startswith("bg_")):
+                ids_to_clear.append(oid)
+        if ids_to_clear:
+            self.engine.clear_objects(ids_to_clear, force_destroy=True)
+        self.controlled_agents.clear()
+        self.controlled_agent_ids.clear()
+        for aid in list(self.engine.agent_manager.active_agents.keys()):
+            if aid.startswith("bg_") or aid.startswith("controlled_"):
+                self.engine.agent_manager.active_agents.pop(aid, None)
+
         if self.replay_sdc and hasattr(self.engine, "traffic_manager"):
             sdc_sid = self.engine.traffic_manager.sdc_scenario_id
             self.sdc_track = self.engine.traffic_manager.current_traffic_data.get(sdc_sid, None)
-        _obj_to_clean_this_frame = []
-        self.car_birth_info_list = []
 
-        # Pre-filter: Check tracks against map AND check for static vehicles
-        
-        for scenario_id, track in self.engine.traffic_manager.current_traffic_data.items():
-            if scenario_id == self.engine.traffic_manager.sdc_scenario_id:
-                continue
-            else:
-                if track["type"] == MetaDriveType.VEHICLE:
-                    _obj_to_clean_this_frame.append(scenario_id)
-                    
-                    valid = track['state']['valid']
-                    if not valid.any():
-                        continue
-                        
-                    first_show = np.argmax(valid)
-                    last_show = len(valid) - 1 - np.argmax(valid[::-1])
-                    mid_show = (first_show + last_show) // 2
-                    
-                    # 1. Lane check (existing logic)
-                    points_to_check = [first_show, mid_show, last_show]
-                    on_road_count = 0
-                    is_valid_track = True
-                    start_pos = track['state']['position'][first_show]
-                    if not is_on_lane(start_pos, self.engine.map_manager, threshold=5.0): # 5m tolerance
-                        mid_pos = track['state']['position'][mid_show]
-                        if not is_on_lane(mid_pos, self.engine.map_manager, threshold=5.0):
-                            is_valid_track = False
-                    
-                    # 2. Static check
-                    # Calculate total displacement and max speed
-                    positions = track['state']['position'][valid.astype(bool)]
-                    velocities = track['state']['velocity'][valid.astype(bool)]
-                    
-                    total_displacement = 0
-                    max_speed = 0
-                    if len(positions) > 1:
-                        total_displacement = np.linalg.norm(positions[-1] - positions[0])
-                        max_speed = np.max(np.linalg.norm(velocities, axis=1))
-                    
-                    is_static = False
-                    if total_displacement < 5.0 and max_speed < 1.0: # Relaxed threshold: <5m move and <1m/s
-                        is_static = True
-                        
-                    # Decision logic:
-                    # - If off-road AND static: Skip completely (don't even spawn as background)
-                    # - If off-road but moving: Maybe keep? Or skip? Usually off-road moving is weird, skip.
-                    # - If on-road but static: Spawn as BACKGROUND (visible but not controlled agent)
-                    # - If on-road and moving: Spawn as CONTROLLED agent
-                    
-                    if not is_valid_track:
-                        # Skip off-road vehicles entirely (both static and moving off-road)
-                        continue
-                        
-                    if is_static:
-                        # Add to background list, but NOT to car_birth_info_list (which is for controlled agents)
-                        # We need a way to spawn them. Let's add a separate list.
-                         self.background_vehicles[scenario_id] = {
-                            'id': track['metadata']['object_id'],
-                            'show_time': first_show,
-                            'begin': (track['state']['position'][first_show, 0], track['state']['position'][first_show, 1]),
-                            'heading': track['state']['heading'][first_show],
-                            'end': (track['state']['position'][last_show, 0], track['state']['position'][last_show, 1]),
-                            'scenario_id': scenario_id,
-                            'length': track['state']['length'][first_show],
-                            'width': track['state']['width'][first_show],
-                            'valid': valid # Need validity to know when to show/hide
-                        }
-                         continue # Do not add to controlled list
-
-                    # Store the full track for replay (only for controlled agents)
-                    self.expert_tracks[scenario_id] = track
-                    
-                    self.car_birth_info_list.append({
-                        'id': track['metadata']['object_id'],
-                        'show_time': first_show,
-                        'begin': (track['state']['position'][first_show, 0], track['state']['position'][first_show, 1]),
-                        'heading': track['state']['heading'][first_show],
-                        'end': (track['state']['position'][last_show, 0], track['state']['position'][last_show, 1]),
-                        'scenario_id': scenario_id, # Keep track of original ID to lookup tracks
-                        'length': track['state']['length'][first_show],
-                        'width': track['state']['width'][first_show]
-                    })
-
-        for scenario_id in _obj_to_clean_this_frame:
+        from Env.utils import filter_traffic_tracks_to_birth_lists
+        traffic_data = self.engine.traffic_manager.current_traffic_data
+        car_birth_info_list, self.background_vehicles, obj_to_clean = filter_traffic_tracks_to_birth_lists(
+            traffic_data,
+            self.engine.traffic_manager.sdc_scenario_id,
+            self.engine.map_manager,
+        )
+        for entry in car_birth_info_list:
+            sid = entry["scenario_id"]
+            if sid in traffic_data:
+                self.expert_tracks[sid] = traffic_data[sid]
+        self.car_birth_info_list = car_birth_info_list
+        for scenario_id in obj_to_clean:
             self.engine.traffic_manager.current_traffic_data.pop(scenario_id)
-        # --- MODIFIED SECTION END ---
 
         self.engine.reset()
         self.reset_sensors()
@@ -185,7 +97,7 @@ class ExpertReplayEnv(MultiAgentScenarioEnv):
         # We covered most of it.
         
         self._spawn_controlled_agents()
-        self._spawn_background_vehicles() # Initial spawn for background
+        self._spawn_all_background_vehicles_at_init()
 
         # Ensure SDC/ego is moved to the correct initial expert state.
         if self.replay_sdc:
@@ -202,119 +114,33 @@ class ExpertReplayEnv(MultiAgentScenarioEnv):
         
         return self._get_all_obs()
 
-    def _spawn_background_vehicles(self):
-        # Spawn static/background vehicles
-        # Since they are static, we might just spawn them once if their show_time is 0
-        # But Waymo tracks have valid bits, they might appear/disappear.
-        # For optimization, if they are truly static (never move), we just spawn them when show_time matches.
-        
-        # We need to track spawned background vehicles to remove them if they become invalid?
-        # Since we defined them as "static", they probably stay put. 
-        # But validity might change (e.g. late spawn).
-        
-        # For simplicity in this step, let's just iterate and spawn if time matches
+    def _spawn_all_background_vehicles_at_init(self):
+        """Spawn all static background vehicles once at reset (no show_time filter; no removal by valid)."""
         for sid, car in self.background_vehicles.items():
-            if car['show_time'] == self.round:
-                # Spawn as a Traffic Vehicle (not PolicyVehicle), or just a static object?
-                # Using DefaultVehicle is fine, but don't add to controlled_agents
-                
-                # Check duplication
-                bg_id = f"bg_{car['id']}"
-                # if bg_id in self.engine.obj_to_id: # obj_to_id might not be available in all versions
-                if bg_id in self.engine.agent_manager.active_agents:
-                    continue
-                    
-                vehicle_config = {}
-                if 'length' in car and 'width' in car:
-                     vehicle_config = {
-                         "length": car['length'],
-                         "width": car['width']
-                     }
-
-                v = self.engine.spawn_object(
-                    DefaultVehicle,
-                    name=bg_id,
-                    vehicle_config=vehicle_config,
-                    position=car['begin'],
-                    heading=car['heading']
-                )
-                
-                # Set color to grey/dark to indicate background
-                v.set_velocity([0, 0])
-                # Maybe set color? MetaDrive vehicles random color.
-                # v.set_color(...) if supported
-                
-                # Register as an active object but NOT controlled agent
-                # The engine manages it. 
-                # CRITICAL: We need it in self.engine.agent_manager.active_agents for Observation?
-                # If we want it to be seen by Lidar/Observation, it needs to be an "agent" or "traffic".
-                # DefaultVehicle spawned this way is just an object.
-                # We should add it to traffic manager? Or just leave it as object?
-                # MultiAgentScenarioEnv._get_all_obs iterates self.engine.agent_manager.active_agents
-                
-                # If we want it in observation, we must add it to active_agents OR iterate over all objects.
-                # Adding to active_agents is easier for compatibility.
-                self.engine.agent_manager.active_agents[bg_id] = v
-                
-                # Store valid mask to remove it later if needed?
-                v.valid_mask = car['valid']
-                v.start_t = car['show_time']
+            bg_id = f"bg_{car['id']}"
+            if bg_id in self.engine.agent_manager.active_agents:
+                continue
+            vehicle_config = {}
+            if 'length' in car and 'width' in car:
+                vehicle_config = {
+                    "length": car['length'],
+                    "width": car['width']
+                }
+            v = self.engine.spawn_object(
+                DefaultVehicle,
+                name=bg_id,
+                vehicle_config=vehicle_config,
+                position=car['begin'],
+                heading=car['heading']
+            )
+            v.set_velocity([0, 0])
+            self.engine.agent_manager.active_agents[bg_id] = v
+            v.valid_mask = car.get('valid')
+            v.start_t = car['show_time']
 
     def _update_background_vehicles(self):
-        # Remove background vehicles if they become invalid
-        # Or spawn new ones
-        self._spawn_background_vehicles()
-        
-        # Check validity for existing
-        to_remove = []
-        for aid, v in self.engine.agent_manager.active_agents.items():
-            if aid.startswith("bg_"):
-                # Check validity
-                if hasattr(v, 'valid_mask'):
-                    curr_step = self.round
-                    if curr_step >= len(v.valid_mask) or not v.valid_mask[curr_step]:
-                        to_remove.append(aid)
-        
-        for aid in to_remove:
-             self.engine.agent_manager.active_agents.pop(aid, None)
-             # if aid in self.engine.obj_to_id:
-             #    self.engine.clear_objects([self.engine.obj_to_id[aid]])
-             # Instead, we should find the object by ID and clear it.
-             # Since we don't track obj directly, we can't easily clear it without obj ref.
-             # Wait, active_agents stores the vehicle object.
-             # So we can just clear that object.
-             pass 
-             
-        # Re-iterate to clear objects properly
-        for aid in to_remove:
-             # We need to find the vehicle object to clear it.
-             # But we popped it from active_agents.
-             # Wait, we should get it before pop.
-             pass
-             
-    def _update_background_vehicles(self):
-        # Remove background vehicles if they become invalid
-        # Or spawn new ones
-        self._spawn_background_vehicles()
-        
-        # Check validity for existing
-        to_remove = []
-        objects_to_clear = []
-        
-        for aid, v in self.engine.agent_manager.active_agents.items():
-            if aid.startswith("bg_"):
-                # Check validity
-                if hasattr(v, 'valid_mask'):
-                    curr_step = self.round
-                    if curr_step >= len(v.valid_mask) or not v.valid_mask[curr_step]:
-                        to_remove.append(aid)
-                        objects_to_clear.append(v)
-        
-        for aid in to_remove:
-             self.engine.agent_manager.active_agents.pop(aid, None)
-             
-        if objects_to_clear:
-            self.engine.clear_objects(objects_to_clear)
+        # Static vehicles are spawned once at init and never removed (no spawn/remove by show_time or valid).
+        pass
 
     def _spawn_controlled_agents(self):
         for car in self.car_birth_info_list:
@@ -468,60 +294,52 @@ class ExpertReplayEnv(MultiAgentScenarioEnv):
         # Get observations
         obs = self._get_all_obs()
 
-        rewards = {aid: 0.0 for aid in self.controlled_agents}
-        dones = {aid: False for aid in self.controlled_agents}
-        dones["__all__"] = (self.round >= self.config["horizon"]) or (len(self.controlled_agents) == 0 and self.round > 190) # Waymo scenarios are usually ~198 steps (20s @ 10Hz) or 90 steps (9s)
-        
-        infos = {aid: {"expert_action": expert_actions.get(aid, np.zeros(2))} for aid in self.controlled_agents}
+        # Build rewards/dones/infos: include controlled_agents and optionally SDC for data collection
+        all_agent_ids = list(self.controlled_agents.keys())
+        if self.replay_sdc and self.sdc_vehicle is not None and self.sdc_agent_id not in all_agent_ids:
+            all_agent_ids = all_agent_ids + [self.sdc_agent_id]
+        rewards = {aid: 0.0 for aid in all_agent_ids}
+        dones = {aid: False for aid in all_agent_ids}
+        dones["__all__"] = (self.round >= self.config["horizon"]) or (len(self.controlled_agents) == 0 and self.round > 190)  # Waymo scenarios are usually ~198 steps (20s @ 10Hz) or 90 steps (9s)
+        infos = {aid: {"expert_action": expert_actions.get(aid, np.zeros(2))} for aid in all_agent_ids}
 
         return obs, rewards, dones, infos
 
+    def _obs_for_vehicle(self, vehicle, exclude_agent_id=None):
+        """Compute 45-dim obs (ego 5 + 10 neighbors x 4) for a vehicle. exclude_agent_id: do not count as neighbor."""
+        ego_state = [
+            vehicle.position[0], vehicle.position[1],
+            vehicle.velocity[0], vehicle.velocity[1],
+            vehicle.heading_theta
+        ]
+        candidates = []
+        for other_id, other_vehicle in self.engine.agent_manager.active_agents.items():
+            if other_id == exclude_agent_id:
+                continue
+            dist = np.linalg.norm(vehicle.position - other_vehicle.position)
+            if dist < 30.0:
+                candidates.append((dist, other_vehicle))
+        candidates.sort(key=lambda x: x[0])
+        top_10 = candidates[:10]
+        neighbor_feats = []
+        for _, neighbor in top_10:
+            neighbor_feats.extend([
+                neighbor.position[0] - vehicle.position[0],
+                neighbor.position[1] - vehicle.position[1],
+                neighbor.velocity[0],
+                neighbor.velocity[1]
+            ])
+        missing = 10 - len(top_10)
+        if missing > 0:
+            neighbor_feats.extend([0.0] * (4 * missing))
+        return np.array(ego_state + neighbor_feats, dtype=np.float32)
+
     def _get_all_obs(self):
         # Implement custom observation: 30m range, 10 nearest vehicles
         obs_dict = {}
-        
         for agent_id, vehicle in self.controlled_agents.items():
-            # 1. Ego State
-            ego_state = [
-                vehicle.position[0], vehicle.position[1],
-                vehicle.velocity[0], vehicle.velocity[1],
-                vehicle.heading_theta
-            ]
-            
-            # 2. Neighbors
-            neighbors = []
-            # Iterate through all vehicles in the engine
-            candidates = []
-            for other_id, other_vehicle in self.engine.agent_manager.active_agents.items():
-                if other_id == agent_id:
-                    continue
-                
-                dist = np.linalg.norm(vehicle.position - other_vehicle.position)
-                if dist < 30.0:
-                    candidates.append((dist, other_vehicle))
-            
-            # Sort by distance
-            candidates.sort(key=lambda x: x[0])
-            
-            # Take top 10
-            top_10 = candidates[:10]
-            
-            neighbor_feats = []
-            for _, neighbor in top_10:
-                neighbor_feats.extend([
-                    neighbor.position[0] - vehicle.position[0], # Relative pos
-                    neighbor.position[1] - vehicle.position[1],
-                    neighbor.velocity[0], # Absolute vel? or Relative? Usually relative in MultiAgent
-                    neighbor.velocity[1]
-                ])
-                
-            # Pad if < 10
-            missing = 10 - len(top_10)
-            if missing > 0:
-                neighbor_feats.extend([0.0] * (4 * missing))
-                
-            # Flatten
-            obs = np.array(ego_state + neighbor_feats, dtype=np.float32)
-            obs_dict[agent_id] = obs
-            
+            obs_dict[agent_id] = self._obs_for_vehicle(vehicle, exclude_agent_id=agent_id)
+        # Include SDC/ego obs for expert data collection (e.g. single-agent)
+        if self.replay_sdc and self.sdc_vehicle is not None:
+            obs_dict[self.sdc_agent_id] = self._obs_for_vehicle(self.sdc_vehicle, exclude_agent_id=self.sdc_agent_id)
         return obs_dict

Env/hbbc_actor_critic.py

@@ -0,0 +1,69 @@
+import torch
+import torch.nn as nn
+
+
+def _get_activation(name: str):
+    name = (name or "elu").lower()
+    mapping = {
+        "elu": nn.ELU,
+        "relu": nn.ReLU,
+        "tanh": nn.Tanh,
+        "leakyrelu": nn.LeakyReLU,
+    }
+    if name not in mapping:
+        raise ValueError(f"Unsupported activation: {name}")
+    return mapping[name]()
+
+
+class ActorCritic(nn.Module):
+    """Minimal HBBC ActorCritic for inference-only deployment."""
+
+    def __init__(
+        self,
+        num_actor_obs=18,
+        num_critic_obs=18,
+        num_actions=2,
+        latent_c_dim=4,
+        latent_eps_dim=6,
+        use_style_latent=True,
+        actor_hidden_dims=None,
+        activation="elu",
+    ):
+        super().__init__()
+        _ = num_critic_obs  # kept for checkpoint compatibility
+        if actor_hidden_dims is None:
+            actor_hidden_dims = [512, 256, 128]
+
+        act_fn = _get_activation(activation)
+        self.latent_c_dim = int(latent_c_dim)
+        self.latent_eps_dim = int(latent_eps_dim)
+        self.use_style_latent = bool(use_style_latent)
+
+        layers = [nn.Linear(num_actor_obs, actor_hidden_dims[0]), act_fn]
+        for i in range(len(actor_hidden_dims) - 1):
+            layers.append(nn.Linear(actor_hidden_dims[i], actor_hidden_dims[i + 1]))
+            layers.append(_get_activation(activation))
+        self.actor_trunk = nn.Sequential(*layers)
+        self.actor_head = nn.Linear(actor_hidden_dims[-1], num_actions)
+
+        if self.use_style_latent:
+            self.style_trunk = nn.Sequential(
+                nn.Linear(self.latent_eps_dim, 512),
+                _get_activation(activation),
+                nn.Linear(512, 256),
+                _get_activation(activation),
+                nn.Linear(256, 128),
+                _get_activation(activation),
+            )
+            self.style_head = nn.Linear(128, self.latent_eps_dim)
+            self.style_activation = torch.tanh
+
+    def act_inference(self, observations: torch.Tensor) -> torch.Tensor:
+        if self.use_style_latent:
+            obs = observations[..., :-(self.latent_c_dim + self.latent_eps_dim)]
+            eps = observations[..., -self.latent_c_dim - self.latent_eps_dim:-self.latent_c_dim]
+            c = observations[..., -self.latent_c_dim:]
+            eps = self.style_activation(self.style_head(self.style_trunk(eps)))
+            observations = torch.cat([obs, eps, c], dim=-1)
+        embedding = self.actor_trunk(observations)
+        return self.actor_head(embedding)

Env/hbbc_background_policy.py

@@ -0,0 +1,274 @@
+import json
+import os
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+
+from Env.hbbc_actor_critic import ActorCritic
+
+
+def _wrap_to_pi(angle: float) -> float:
+    return (angle + np.pi) % (2 * np.pi) - np.pi
+
+
+def _normalize_eps(eps: np.ndarray) -> np.ndarray:
+    eps = np.asarray(eps, dtype=np.float32).reshape(-1)
+    if eps.shape[0] != 6:
+        raise ValueError(f"latent_eps must be 6-dim, got {eps.shape[0]}")
+    norm = float(np.linalg.norm(eps))
+    if norm < 1e-8:
+        eps = np.array([1.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float32)
+    else:
+        eps = eps / norm
+    return np.clip(eps, -1.0, 1.0)
+
+
+def _normalize_c(latent_c: np.ndarray) -> np.ndarray:
+    c = np.asarray(latent_c, dtype=np.float32).reshape(-1)
+    if c.shape[0] != 4:
+        raise ValueError(f"latent_c must be 4-dim, got {c.shape[0]}")
+    idx = int(np.argmax(c))
+    one_hot = np.zeros(4, dtype=np.float32)
+    one_hot[idx] = 1.0
+    return one_hot
+
+
+def _sample_latent(rng: np.random.RandomState) -> Tuple[np.ndarray, np.ndarray]:
+    eps = _normalize_eps(rng.randn(6).astype(np.float32))
+    mode = int(rng.randint(0, 4))
+    c = np.zeros(4, dtype=np.float32)
+    c[mode] = 1.0
+    return eps, c
+
+
+@dataclass
+class VehicleStateCache:
+    last_heading_theta: Optional[float] = None
+    last_action: Tuple[float, float] = (0.0, 0.0)
+    last_speed_km_h: Optional[float] = None
+
+
+class HBBCModelWrapper:
+    _cache: Dict[Tuple[str, str], "HBBCModelWrapper"] = {}
+
+    def __init__(self, model_path: str, device: str = "cpu"):
+        self.model_path = os.path.abspath(model_path)
+        self.device = torch.device(device)
+        self.model = self._load_model()
+
+    @classmethod
+    def get(cls, model_path: str, device: str = "cpu") -> "HBBCModelWrapper":
+        key = (os.path.abspath(model_path), str(torch.device(device)))
+        if key not in cls._cache:
+            cls._cache[key] = HBBCModelWrapper(model_path=key[0], device=key[1])
+        return cls._cache[key]
+
+    def _load_model(self) -> ActorCritic:
+        model = ActorCritic(
+            num_actor_obs=18,
+            num_critic_obs=18,
+            num_actions=2,
+            latent_c_dim=4,
+            latent_eps_dim=6,
+            use_style_latent=True,
+        ).to(self.device)
+        try:
+            ckpt = torch.load(self.model_path, map_location=self.device, weights_only=True)
+        except Exception:
+            ckpt = torch.load(self.model_path, map_location=self.device, weights_only=False)
+        state_dict = ckpt["actor_critic"] if isinstance(ckpt, dict) and "actor_critic" in ckpt else ckpt
+        missing, unexpected = model.load_state_dict(state_dict, strict=False)
+        if missing:
+            raise RuntimeError(
+                f"HBBC checkpoint missing required keys for {self.model_path}: {missing}"
+            )
+        if unexpected:
+            print(f"[HBBC] ignore extra checkpoint keys: {unexpected[:8]}{'...' if len(unexpected) > 8 else ''}")
+        model.eval()
+        return model
+
+    def act_batch(self, obs_batch: np.ndarray) -> np.ndarray:
+        obs_batch = np.asarray(obs_batch, dtype=np.float32)
+        with torch.no_grad():
+            obs_t = torch.from_numpy(obs_batch).to(self.device)
+            actions = self.model.act_inference(obs_t).cpu().numpy()
+        return np.clip(actions, -1.0, 1.0)
+
+
+class HBBCLatentManager:
+    def __init__(self, mode: str = "per_vehicle_fixed", seed: int = 0, latent_json_path: Optional[str] = None):
+        self.mode = mode
+        self.rng = np.random.RandomState(seed)
+        self.latent_json_path = latent_json_path
+        self.manual_object_latent: Dict[str, Dict[str, np.ndarray]] = {}
+        self.manual_agent_latent: Dict[str, Dict[str, np.ndarray]] = {}
+        self.manual_global_latent: Optional[Tuple[np.ndarray, np.ndarray]] = None
+        self.vehicle_latent: Dict[str, Tuple[np.ndarray, np.ndarray]] = {}
+        self._episode_latent: Optional[Tuple[np.ndarray, np.ndarray]] = None
+        self._load_manual_latent_json()
+
+    def reset_episode(self):
+        self.vehicle_latent.clear()
+        self._episode_latent = None
+        if self.mode == "per_episode_reset":
+            self._episode_latent = _sample_latent(self.rng)
+
+    def _load_manual_latent_json(self):
+        if not self.latent_json_path:
+            return
+        path = os.path.abspath(self.latent_json_path)
+        if not os.path.exists(path):
+            print(f"[HBBC] latent json not found: {path}, fallback to random sampling.")
+            return
+        try:
+            with open(path, "r", encoding="utf-8") as f:
+                data = json.load(f)
+        except Exception as e:
+            print(f"[HBBC] failed to load latent json ({path}): {e}. fallback to random sampling.")
+            return
+
+        object_section = data.get("object_id", {})
+        agent_section = data.get("agent_id", {})
+        global_section = data.get("global")
+
+        if global_section is not None:
+            parsed = self._parse_one_latent(global_section, "global")
+            if parsed is not None:
+                self.manual_global_latent = (parsed["latent_eps"], parsed["latent_c"])
+
+        for key, value in object_section.items():
+            parsed = self._parse_one_latent(value, f"object_id:{key}")
+            if parsed is not None:
+                self.manual_object_latent[str(key)] = parsed
+        for key, value in agent_section.items():
+            parsed = self._parse_one_latent(value, f"agent_id:{key}")
+            if parsed is not None:
+                self.manual_agent_latent[str(key)] = parsed
+
+    @staticmethod
+    def _parse_one_latent(value: dict, name: str) -> Optional[Dict[str, np.ndarray]]:
+        if not isinstance(value, dict):
+            print(f"[HBBC] invalid latent entry ({name}): expect dict.")
+            return None
+        try:
+            eps = _normalize_eps(value["latent_eps"])
+            c = _normalize_c(value["latent_c"])
+            return {"latent_eps": eps, "latent_c": c}
+        except Exception as e:
+            print(f"[HBBC] invalid latent entry ({name}): {e}")
+            return None
+
+    def _lookup_manual(self, object_id: Optional[str], agent_id: Optional[str]) -> Optional[Tuple[np.ndarray, np.ndarray]]:
+        if object_id is not None and object_id in self.manual_object_latent:
+            e = self.manual_object_latent[object_id]["latent_eps"]
+            c = self.manual_object_latent[object_id]["latent_c"]
+            return e, c
+        if agent_id is not None and agent_id in self.manual_agent_latent:
+            e = self.manual_agent_latent[agent_id]["latent_eps"]
+            c = self.manual_agent_latent[agent_id]["latent_c"]
+            return e, c
+        if self.manual_global_latent is not None:
+            return self.manual_global_latent
+        return None
+
+    def get_latent(self, vehicle_key: str, object_id: Optional[str], agent_id: Optional[str]) -> Tuple[np.ndarray, np.ndarray]:
+        manual = self._lookup_manual(object_id=object_id, agent_id=agent_id)
+        if manual is not None:
+            return manual
+        if self.mode == "per_episode_reset":
+            if self._episode_latent is None:
+                self._episode_latent = _sample_latent(self.rng)
+            return self._episode_latent
+        if vehicle_key not in self.vehicle_latent:
+            self.vehicle_latent[vehicle_key] = _sample_latent(self.rng)
+        return self.vehicle_latent[vehicle_key]
+
+
+class HBBCBackgroundController:
+    def __init__(
+        self,
+        model_path: str,
+        device: str = "cpu",
+        latent_mode: str = "per_vehicle_fixed",
+        latent_json_path: Optional[str] = None,
+        seed: int = 0,
+        dt: float = 0.1,
+    ):
+        self.model = HBBCModelWrapper.get(model_path=model_path, device=device)
+        self.latent_mgr = HBBCLatentManager(mode=latent_mode, seed=seed, latent_json_path=latent_json_path)
+        self.dt = float(dt)
+        self.vehicle_state: Dict[str, VehicleStateCache] = {}
+
+    def reset_episode(self):
+        self.latent_mgr.reset_episode()
+        self.vehicle_state.clear()
+
+    def remove_vehicle(self, vehicle_key: str):
+        self.vehicle_state.pop(vehicle_key, None)
+        self.latent_mgr.vehicle_latent.pop(vehicle_key, None)
+
+    def _build_base_state(self, vehicle, vehicle_key: str) -> np.ndarray:
+        state = self.vehicle_state.get(vehicle_key)
+        if state is None:
+            state = VehicleStateCache()
+            self.vehicle_state[vehicle_key] = state
+
+        speed_km_h = float(getattr(vehicle, "speed_km_h", 0.0))
+        max_speed_km_h = float(getattr(vehicle, "max_speed_km_h", 120.0))
+        veh_vel = np.clip((speed_km_h + 1.0) / (max_speed_km_h + 1.0), 0.0, 1.0)
+
+        heading_theta = float(getattr(vehicle, "heading_theta", 0.0))
+        if state.last_heading_theta is None:
+            yaw_rate = 0.0
+        else:
+            yaw_rate = _wrap_to_pi(heading_theta - state.last_heading_theta) / self.dt
+        yaw_rate = float(np.clip(yaw_rate, -5.0, 5.0))
+
+        current_action = getattr(vehicle, "current_action", None)
+        if current_action is None:
+            last_action_0, last_action_1 = state.last_action
+        else:
+            try:
+                last_action_0, last_action_1 = float(current_action[0]), float(current_action[1])
+            except Exception:
+                last_action_0, last_action_1 = state.last_action
+
+        state.last_heading_theta = heading_theta
+        state.last_speed_km_h = speed_km_h
+        state.last_action = (last_action_0, last_action_1)
+
+        obs = np.array(
+            [
+                0.0,
+                0.0,
+                0.0,
+                veh_vel,
+                0.0,
+                yaw_rate * 0.5,
+                last_action_0,
+                last_action_1,
+            ],
+            dtype=np.float32,
+        )
+        return obs
+
+    def build_obs(self, vehicle, vehicle_key: str, object_id: Optional[str], agent_id: Optional[str]) -> np.ndarray:
+        base = self._build_base_state(vehicle, vehicle_key=vehicle_key)
+        eps, c = self.latent_mgr.get_latent(vehicle_key=vehicle_key, object_id=object_id, agent_id=agent_id)
+        return np.concatenate([base, eps, c], axis=-1).astype(np.float32)
+
+    def infer_actions(self, batch: List[Tuple[str, object, Optional[str], Optional[str]]]) -> Dict[str, np.ndarray]:
+        if not batch:
+            return {}
+        obs_list = []
+        vehicle_ids = []
+        for vehicle_key, vehicle, object_id, agent_id in batch:
+            obs_list.append(self.build_obs(vehicle, vehicle_key=vehicle_key, object_id=object_id, agent_id=agent_id))
+            vehicle_ids.append(vehicle_key)
+        actions = self.model.act_batch(np.stack(obs_list, axis=0))
+        out = {}
+        for idx, key in enumerate(vehicle_ids):
+            out[key] = actions[idx].astype(np.float32)
+        return out

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
@@ -63,3 +63,31 @@ class InverseDynamics:
         norm_steering = np.clip(steering / self.max_steering, -1.0, 1.0)
 
         return np.array([norm_steering, norm_acc]), {'raw_acc': acc, 'raw_steering': steering}
+
+    def apply_action(self, current_state, action, dt=0.1):
+        """
+        Forward dynamics: given current_state and action [steering, acc] in [-1, 1], return next_state.
+        State format: dict with position (x,y), heading, velocity (vx, vy).
+        """
+        steering_norm, acc_norm = float(action[0]), float(action[1])
+        acc = acc_norm * self.max_acc
+        steering = steering_norm * self.max_steering
+        pos = np.array(current_state['position'][:2], dtype=np.float64)
+        heading = float(current_state['heading'])
+        vel = np.array(current_state['velocity'], dtype=np.float64)
+        v = np.linalg.norm(vel)
+        if v < 0.1:
+            v = 0.1
+        theta_dot = v * np.tan(steering) / self.wheelbase
+        v_next = v + acc * dt
+        v_next = max(0.0, v_next)
+        heading_next = heading + theta_dot * dt
+        heading_next = np.arctan2(np.sin(heading_next), np.cos(heading_next))
+        vx_next = v_next * np.cos(heading_next)
+        vy_next = v_next * np.sin(heading_next)
+        pos_next = pos + dt * np.array([vx_next, vy_next])
+        return {
+            'position': pos_next,
+            'heading': heading_next,
+            'velocity': np.array([vx_next, vy_next]),
+        }

Env/scenario_env.py

@@ -53,6 +53,13 @@ class MultiAgentScenarioEnv(ScenarioEnv):
             data_directory=None,
             num_controlled_agents=3,
             horizon=1000,
+            # HBBC background vehicle controls (optional)
+            enable_hbbc_background=False,
+            hbbc_model_path="models/hbbc/hbbc.pt",
+            hbbc_inference_device="cpu",
+            hbbc_latent_mode="per_vehicle_fixed",
+            hbbc_latent_json_path=None,
+            hbbc_dt=0.1,
         ))
         return config
 
@@ -64,6 +71,11 @@ class MultiAgentScenarioEnv(ScenarioEnv):
         self.round = 0
         super().__init__(config)
 
+    @property
+    def num_controlled_in_scenario(self) -> int:
+        """整个场景中受控车轨迹总数（car_birth_info_list 长度），会在不同 show_time 陆续 spawn。"""
+        return len(getattr(self, "car_birth_info_list", []))
+
     def reset(self, seed: Union[None, int] = None):
         self.round = 0
         if self.logger is None:
@@ -76,30 +88,23 @@ class MultiAgentScenarioEnv(ScenarioEnv):
         if self.engine is None:
             raise ValueError("Broken MetaDrive instance.")
 
-        # 记录专家数据中每辆车的位置，接着全部清除，只保留位置等信息，用于后续生成
-        _obj_to_clean_this_frame = []
-        self.car_birth_info_list = []
-        for scenario_id, track in self.engine.traffic_manager.current_traffic_data.items():
-            if scenario_id == self.engine.traffic_manager.sdc_scenario_id:
-                continue
-            else:
-                if track["type"] == MetaDriveType.VEHICLE:
-                    _obj_to_clean_this_frame.append(scenario_id)
-                    valid = track['state']['valid']
-                    first_show = np.argmax(valid) if valid.any() else -1
-                    last_show = len(valid) - 1 - np.argmax(valid[::-1]) if valid.any() else -1
-                    # id，出现时间，出生点坐标，出生朝向，目的地
-                    self.car_birth_info_list.append({
-                        'id': track['metadata']['object_id'],
-                        'show_time': first_show,
-                        'begin': (track['state']['position'][first_show, 0], track['state']['position'][first_show, 1]),
-                        'heading': track['state']['heading'][first_show],
-                        'end': (track['state']['position'][last_show, 0], track['state']['position'][last_show, 1])
-                    })
-
-        for scenario_id in _obj_to_clean_this_frame:
+        # 注意：_build_birth_lists_from_traffic() 在 engine.reset() 之前执行，读的是当前 engine 的
+        # current_traffic_data 与 map_manager.current_map。若复用同一 env 连续 reset(0)、reset(1)，
+        # MetaDrive 可能已按 seed 更新了 traffic 为 scenario 1，但 map 仍为 scenario 0（在 engine.reset() 才切图），
+        # 导致 is_on_lane( scenario_1 车位, scenario_0 地图 ) 全为 False → 全部 off_lane → 0 受控车。
+        # 因此多场景时应“每个 scenario 单独建 env”（start_scenario_index=i, num_scenarios=1）再 reset(seed=i)。
+        self.background_vehicles = getattr(self, "background_vehicles", {})
+        self.car_birth_info_list, self.background_vehicles, _obj_to_clean = self._build_birth_lists_from_traffic()
+        for scenario_id in _obj_to_clean:
             self.engine.traffic_manager.current_traffic_data.pop(scenario_id)
 
+        # Clear vehicles we spawned via engine.spawn_object() so _object_clean_check() passes
+        ids_to_clear = [v.id for v in self.controlled_agents.values()]
+        if ids_to_clear:
+            self.engine.clear_objects(ids_to_clear)
+        self.controlled_agents.clear()
+        self.controlled_agent_ids.clear()
+
         self.engine.reset()
         self.reset_sensors()
         self.engine.taskMgr.step()
@@ -114,14 +119,32 @@ class MultiAgentScenarioEnv(ScenarioEnv):
         self.episode_rewards = defaultdict(float)
         self.episode_lengths = defaultdict(int)
 
-        self.controlled_agents.clear()
-        self.controlled_agent_ids.clear()
-
         super().reset(seed)  # 初始化场景
         self._spawn_controlled_agents()
 
         return self._get_all_obs()
 
+    def _build_birth_lists_from_traffic(self):
+        """Build car_birth_info_list and obj_to_clean from current_traffic_data. Override for filtered (lane/static) selection."""
+        _obj_to_clean_this_frame = []
+        car_birth_info_list = []
+        for scenario_id, track in self.engine.traffic_manager.current_traffic_data.items():
+            if scenario_id == self.engine.traffic_manager.sdc_scenario_id:
+                continue
+            if track["type"] == MetaDriveType.VEHICLE:
+                _obj_to_clean_this_frame.append(scenario_id)
+                valid = track["state"]["valid"]
+                first_show = int(np.argmax(valid)) if valid.any() else -1
+                last_show = len(valid) - 1 - int(np.argmax(valid[::-1])) if valid.any() else -1
+                car_birth_info_list.append({
+                    "id": track["metadata"]["object_id"],
+                    "show_time": first_show,
+                    "begin": (track["state"]["position"][first_show, 0], track["state"]["position"][first_show, 1]),
+                    "heading": track["state"]["heading"][first_show],
+                    "end": (track["state"]["position"][last_show, 0], track["state"]["position"][last_show, 1]),
+                })
+        return car_birth_info_list, {}, _obj_to_clean_this_frame
+
     def _spawn_controlled_agents(self):
         # ego_vehicle = self.engine.agent_manager.active_agents.get("default_agent")
         # ego_position = ego_vehicle.position if ego_vehicle else np.array([0, 0])
@@ -190,6 +213,7 @@ class MultiAgentScenarioEnv(ScenarioEnv):
                 self.controlled_agents[agent_id].before_step(action)
 
         self.engine.step()
+        self.engine.after_step()
 
         for agent_id in action_dict:
             if agent_id in self.controlled_agents:

Env/utils.py

@@ -2,6 +2,227 @@ import numpy as np
 import torch
 import random
 
+from metadrive.type import MetaDriveType
+
+
+def _static_obbs_overlap(car_a, car_b):
+    """
+    Check if two static vehicle OBBs overlap (2D SAT).
+    car_a, car_b: dicts with "begin" (x, y), "heading" (rad), "length", "width".
+    begin is center; half-extents are length/2, width/2.
+    """
+    def _get_corners(car):
+        cx, cy = car["begin"][0], car["begin"][1]
+        h = float(car["heading"])
+        L2 = float(car["length"]) / 2.0
+        W2 = float(car["width"]) / 2.0
+        ux, uy = np.cos(h), np.sin(h)
+        vx, vy = -np.sin(h), np.cos(h)
+        return np.array([
+            [cx + L2 * ux + W2 * vx, cy + L2 * uy + W2 * vy],
+            [cx + L2 * ux - W2 * vx, cy + L2 * uy - W2 * vy],
+            [cx - L2 * ux - W2 * vx, cy - L2 * uy - W2 * vy],
+            [cx - L2 * ux + W2 * vx, cy - L2 * uy + W2 * vy],
+        ])
+
+    def _get_axes(car):
+        h = float(car["heading"])
+        return [
+            np.array([np.cos(h), np.sin(h)]),
+            np.array([-np.sin(h), np.cos(h)]),
+        ]
+
+    corners_a = _get_corners(car_a)
+    corners_b = _get_corners(car_b)
+    axes = _get_axes(car_a) + _get_axes(car_b)
+
+    for axis in axes:
+        proj_a = corners_a @ axis
+        proj_b = corners_b @ axis
+        min_a, max_a = proj_a.min(), proj_a.max()
+        min_b, max_b = proj_b.min(), proj_b.max()
+        if max_a < min_b or max_b < min_a:
+            return False
+    return True
+
+
+def _deduplicate_background_by_collision(background_vehicles):
+    """
+    Merge static tracks that collide (same physical vehicle). Build collision graph,
+    find connected components, keep one representative per component (min show_time, then scenario_id).
+    """
+    if not background_vehicles:
+        return background_vehicles
+    items = list(background_vehicles.items())
+    n = len(items)
+    # Build adjacency by index
+    parent = list(range(n))
+
+    def find(i):
+        if parent[i] != i:
+            parent[i] = find(parent[i])
+        return parent[i]
+
+    def union(i, j):
+        pi, pj = find(i), find(j)
+        if pi != pj:
+            parent[pi] = pj
+
+    for i in range(n):
+        for j in range(i + 1, n):
+            if _static_obbs_overlap(items[i][1], items[j][1]):
+                union(i, j)
+
+    # Representative per component: index with min (show_time, scenario_id)
+    comp_rep = {}
+    for i in range(n):
+        r = find(i)
+        sid, car = items[i][0], items[i][1]
+        key = (car.get("show_time", 0), sid)
+        if r not in comp_rep or key < comp_rep[r][0]:
+            comp_rep[r] = (key, sid, car)
+
+    return {sid: car for (_, sid, car) in comp_rep.values()}
+
+
+def is_on_lane(pos, map_manager, threshold=2.0):
+    """Check if a position is on a valid lane (within lateral tolerance)."""
+    if map_manager is None or map_manager.current_map is None:
+        return True
+    try:
+        lane, _ = map_manager.current_map.road_network.get_closest_lane_index(pos, return_lane=True)
+        if lane is None:
+            return False
+        long, lat = lane.local_coordinates(pos)
+        width = lane.width
+        if abs(lat) <= (width / 2 + threshold):
+            return True
+        return False
+    except Exception:
+        return False
+
+
+def filter_traffic_tracks_to_birth_lists(
+    current_traffic_data,
+    sdc_scenario_id,
+    map_manager,
+    *,
+    lane_threshold=5.0,
+    static_displacement_threshold=5.0,
+    static_speed_threshold=1.0,
+    return_stats=False,
+    deduplicate_static_by_collision=True,
+):
+    """
+    Filter traffic tracks into controlled (car_birth_info_list) and background lists.
+
+    - controlled (car_birth_info_list): 非 SDC、类型 VEHICLE、至少一帧 valid、在车道内、且非静态
+      （位移/速度超过阈值）。用于策略控制或专家回放，spawn 时机为 show_time == round。
+    - background (background_vehicles): 同上但在车道内且判定为静态（位移 < 5m、速度 < 1 m/s）。
+      仅作场景占位与观测邻居，spawn 时机为 show_time == round，按 valid 在 step 中移除。
+
+    Returns (car_birth_info_list, background_vehicles, obj_to_clean) or, if return_stats=True,
+    (car_birth_info_list, background_vehicles, obj_to_clean, stats_dict).
+    stats_dict: n_total, n_no_valid, n_off_lane, n_static, n_controlled.
+    """
+    car_birth_info_list = []
+    background_vehicles = {}
+    obj_to_clean = []
+    n_total = 0
+    n_no_valid = 0
+    n_off_lane = 0
+    n_static = 0
+
+    for scenario_id, track in current_traffic_data.items():
+        if scenario_id == sdc_scenario_id:
+            continue
+        if track["type"] != MetaDriveType.VEHICLE:
+            continue
+
+        n_total += 1
+        obj_to_clean.append(scenario_id)
+        valid = track["state"]["valid"]
+        if not valid.any():
+            n_no_valid += 1
+            continue
+
+        first_show = int(np.argmax(valid))
+        last_show = len(valid) - 1 - int(np.argmax(valid[::-1]))
+        mid_show = (first_show + last_show) // 2
+
+        start_pos = track["state"]["position"][first_show]
+        is_valid_track = True
+        if not is_on_lane(start_pos, map_manager, threshold=lane_threshold):
+            mid_pos = track["state"]["position"][mid_show]
+            if not is_on_lane(mid_pos, map_manager, threshold=lane_threshold):
+                is_valid_track = False
+
+        if not is_valid_track:
+            n_off_lane += 1
+            continue
+
+        positions = track["state"]["position"][valid.astype(bool)]
+        velocities = track["state"]["velocity"][valid.astype(bool)]
+        total_displacement = 0.0
+        max_speed = 0.0
+        if len(positions) > 1:
+            total_displacement = float(np.linalg.norm(positions[-1] - positions[0]))
+            max_speed = float(np.max(np.linalg.norm(velocities, axis=1)))
+        is_static = total_displacement < static_displacement_threshold and max_speed < static_speed_threshold
+
+        if is_static:
+            n_static += 1
+            background_vehicles[scenario_id] = {
+                "id": track["metadata"]["object_id"],
+                "show_time": first_show,
+                "begin": (
+                    float(track["state"]["position"][first_show, 0]),
+                    float(track["state"]["position"][first_show, 1]),
+                ),
+                "heading": float(track["state"]["heading"][first_show]),
+                "end": (
+                    float(track["state"]["position"][last_show, 0]),
+                    float(track["state"]["position"][last_show, 1]),
+                ),
+                "scenario_id": scenario_id,
+                "length": track["state"]["length"][first_show],
+                "width": track["state"]["width"][first_show],
+                "valid": valid,
+            }
+            continue
+
+        car_birth_info_list.append({
+            "id": track["metadata"]["object_id"],
+            "show_time": first_show,
+            "begin": (
+                float(track["state"]["position"][first_show, 0]),
+                float(track["state"]["position"][first_show, 1]),
+            ),
+            "heading": float(track["state"]["heading"][first_show]),
+            "end": (
+                float(track["state"]["position"][last_show, 0]),
+                float(track["state"]["position"][last_show, 1]),
+            ),
+            "scenario_id": scenario_id,
+            "length": track["state"]["length"][first_show],
+            "width": track["state"]["width"][first_show],
+        })
+
+    if deduplicate_static_by_collision and background_vehicles:
+        background_vehicles = _deduplicate_background_by_collision(background_vehicles)
+
+    if return_stats:
+        stats = {
+            "n_total": n_total,
+            "n_no_valid": n_no_valid,
+            "n_off_lane": n_off_lane,
+            "n_static": n_static,
+            "n_controlled": len(car_birth_info_list),
+        }
+        return car_birth_info_list, background_vehicles, obj_to_clean, stats
+    return car_birth_info_list, background_vehicles, obj_to_clean
+
+
 def set_seed(seed):
     if seed == -1:
         seed = np.random.randint(0, 10000)

README.md

@@ -1,98 +1,148 @@
 # MAGAIL4AutoDrive
 
-> 基于多智能体生成对抗模仿学习(MAGAIL)的自动驾驶训练系统 | MetaDrive + Waymo Open Motion Dataset
+基于 **MetaDrive** 仿真器和 **Waymo Open Motion Dataset** 的自动驾驶多智能体模仿学习（MAGAIL）与行为克隆（BC）训练系统。
 
-本项目利用 Waymo 真实驾驶数据，通过 MetaDrive 仿真环境构建专家回放系统，提取车辆状态与动作，用于训练多智能体模仿学习算法 (MAGAIL)。
+本项目旨在从真实的 Waymo 驾驶数据中提取专家轨迹，并通过模仿学习（Imitation Learning）训练能够适应复杂交互场景的自动驾驶策略。
 
-## 📁 核心模块
+## 目录结构
 
-*   **`Env/expert_replay_env.py`**: 专家回放环境。核心类 `ExpertReplayEnv`，负责读取 Waymo 轨迹，计算逆动力学动作，并过滤非道路/静态车辆。
-*   **`Env/inverse_dynamics.py`**: 逆动力学模块。根据车辆位置和航向计算油门、刹车和转向动作。
-*   **`scripts/generate_expert_data.py`**: 数据收集脚本。批量运行场景并保存训练数据。
-*   **`scripts/visualize_replay.py`**: 可视化脚本。用于观察回放效果和数据质量。
-
-***
-
-## 🚀 1. 数据收集
-
-### 生成专家数据
-使用 `generate_expert_data.py` 脚本从 Waymo 数据集中批量提取 (State, Action) 对。
-
-```bash
-# 设置 Python 路径
-export PYTHONPATH=$PYTHONPATH:.:./metadrive
-
-# 运行生成脚本
-# --data_dir: Waymo 数据路径 (建议使用 exp_filtered)
-# --output_dir: 结果保存路径
-# --num_scenarios: 要处理的场景数量
-python scripts/generate_expert_data.py \
-    --data_dir data/exp_filtered \
-    --output_dir data/training_data \
-    --num_scenarios 100 \
-    --start_index 0
+```text
+MAGAIL4AutoDrive/
+├── Algorithm/                 # 强化学习与模仿学习算法实现
+│   ├── policy.py              # 基础策略网络 (MLP 等)
+│   ├── ppo.py                 # PPO 算法实现
+│   ├── magail.py              # MAGAIL 算法核心逻辑
+│   ├── disc.py                # 判别器 (Discriminator) 网络
+│   └── ...
+├── Env/                       # 仿真环境封装 (MetaDrive Wrapper)
+│   ├── bc_env.py              # BCScenarioEnv，45 维观测（BC/MAGAIL 共用）
+│   ├── bc_ego_replay_env.py   # BCEgoReplayEnv，单智能体 BC 评估（仅 ego 受控）
+│   ├── scenario_env.py       # 多智能体基础场景环境
+│   ├── expert_replay_env.py  # 专家轨迹回放环境（数据生成与回放）
+│   ├── inverse_dynamics.py   # 逆动力学模块 (轨迹 -> 动作)
+│   ├── simple_idm_policy.py  # ConstantVelocityPolicy 占位策略
+│   └── ...
+├── dataset/                   # 数据集加载器
+│   ├── loader.py              # 主流水线：load_expert_pkl、MAGAILExpertDataset
+│   └── expert_dataset.py      # 可选 107 维/5 维管线
+├── scripts/                   # 工具脚本（数据、回放、可视化、分析）
+│   ├── generate_expert_data.py    # 从 Waymo 生成专家 (obs, act) pkl
+│   ├── visualize.py              # 可视化统一入口（replay / policy / trajectory）
+│   ├── analyze_expert_data.py    # 数据分布分析
+│   ├── launch_tensorboard.py     # 启动 TensorBoard
+│   ├── README.md                  # 脚本用法说明
+│   └── ...
+├── data/                      # 数据目录（相对路径）
+│   ├── exp_filtered/          # Waymo 场景数据
+│   ├── training_data/        # 专家 pkl 输出（generate_expert_data）
+│   └── trajectories/         # 其他轨迹 pkl（如 expert_dataset 输出）
+├── models/                    # 模型保存目录（相对路径）
+│   ├── bc/                    # BC 模型 (.pt)
+│   └── magail/                # MAGAIL 模型 (*_actor.pth, *_critic.pth)
+├── logs/                      # 训练日志 (TensorBoard)
+│   ├── bc/
+│   └── magail/
+├── train_bc.py                # [根目录] BC 训练
+├── train_magail.py            # [根目录] MAGAIL 训练
+└── README.md
 ```
 
-**生成的 `.pkl` 文件结构**：
-包含一个列表，每个元素是一条车辆轨迹（Trajectory Dictionary）：
-*   `obs`: `(T, 45)` - 观测矩阵。包含 Ego 状态 (5维) + 10辆邻居车相对信息 (40维)。
-*   `acts`: `(T, 2)` - 动作矩阵。`[Steering, Accel]`，归一化到 `[-1, 1]`。
-*   `agent_id`: 车辆 ID。
-*   `scenario_id`: 所属场景 ID。
+## 路径约定（相对项目根）
 
-**内置过滤器**：
-脚本会自动过滤掉以下无效车辆：
-1.  **非道路车辆**：始终在停车场或路外行驶的车辆。
-2.  **静态车辆**：全称移动距离小于 5米 且速度从未超过 1m/s 的车辆（作为背景流存在，不收集数据）。
+- **数据**：Waymo 场景 `data/exp_filtered`；专家 pkl `data/training_data`；其他轨迹 `data/trajectories`
+- **模型**：BC `models/bc/`，MAGAIL `models/magail/`
+- **日志**：TensorBoard 写入 `logs/bc/`、`logs/magail/`
 
----
+所有默认路径均为相对项目根，便于在不同设备上复用。
 
-## 🔍 2. 数据可视化与验证
+## 数据处理流程
 
-### 回放可视化
-使用 `visualize_replay.py` 直观地观察回放效果，确认车辆行为是否自然，以及过滤逻辑是否生效。
+从 Waymo Motion 原始数据到本项目训练用专家 pkl，依次为：
+
+**1) 下载 Waymo Motion（TFRecord）**  
+安装 `gsutil` 并登录 Google 账号后，例如只下载 training_20s：
 
 ```bash
-# 运行可视化
-# --horizon: 回放的最大步数 (Waymo 场景通常为 90 或 198 步)
-python scripts/visualize_replay.py \
-    --data_dir data/exp_filtered \
-    --start_index 0 \
-    --num_scenarios 1 \
-    --horizon 200
+gsutil -m cp -r "gs://waymo_open_dataset_motion_v_1_2_0/uncompressed/scenario/training_20s" ./waymo/
 ```
 
-**观察要点**：
-*   **受控车辆 (Controlled Agents)**：控制台会显示数量（如 `Controlled agents: 2`）。这些是真正产生数据的车辆。
-*   **背景车辆**：如果在渲染图中看到其他车（通常是路边停放的），但受控数量很少，说明静态过滤生效了。
-
-### 数据分析
-使用 `analyze_expert_data.py` 查看生成数据的统计分布。
+**2) ScenarioNet Convert（TFRecord → ScenarioNet 场景库）**  
+需安装 ScenarioNet、MetaDrive 及 TensorFlow 2.11、protobuf 3.20；转换时不用 GPU。
 
 ```bash
-python scripts/analyze_expert_data.py --data_path data/training_data/expert_data_0_100.pkl
+python -m scenarionet.convert_waymo -d data/exp_converted --raw_data_path ./waymo/training_20s --num_workers 64
 ```
 
----
+**3) ScenarioNet Filter（按需筛选场景）**  
+从 convert 得到的场景库中筛掉含红绿灯、天桥等场景，输出到如 `data/exp_filtered`。具体命令以 ScenarioNet 文档为准（Operations → Filter）。
 
-## 🧠 3. 模型训练 (Next Steps)
+**4) 本项目：生成专家 pkl**  
+使用筛选后的场景目录，生成训练用 pkl 到 `data/training_data`：
 
-有了 `data/training_data/` 下的专家数据后，您可以开始训练 MAGAIL 模型。
+- **多智能体**（所有受控车轨迹，输出 `expert_data_{start_index}_{num_scenarios}.pkl`）：
+```bash
+python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100 --start_index 0
+```
 
-### 训练流程
-1.  **加载数据**：使用 `dataset/expert_dataset.py` 中的 `ExpertDataset` 类加载 `.pkl` 数据。
-2.  **初始化 MAGAIL**：
-    *   **Generator (Policy)**: 接收观测 `(B, 45)`，输出动作 `(B, 2)`。
-    *   **Discriminator**: 接收状态-动作对 `(s, a)`，判断是专家还是生成器。
-3.  **交互采样**：
-    *   在 `MultiAgentScenarioEnv`（非回放模式）中运行 Policy。
-    *   收集 Policy 生成的轨迹。
-4.  **对抗更新**：
-    *   利用专家数据和 Policy 数据训练 Discriminator。
-    *   利用 Discriminator 的输出作为 Reward (GAIL Reward) 训练 Policy (PPO/TRPO)。
+- **单智能体**（仅 ego 车轨迹，输出 `expert_data_ego_{start_index}_{num_scenarios}.pkl`，用于单智能体 BC）：
+```bash
+python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100 --start_index 0 --ego_only
+```
 
-### 推荐配置
-*   **Observation**: 45维 (Ego + 10 Neighbors)
-*   **Action**: 2维 Continuous (Steering, Accel)
-*   **Horizon**: 200 steps
-*   **Batch Size**: 1024+ (多智能体环境下数据量很大)
+## 核心工作流
+
+### 1. 数据准备
+使用 `scripts/generate_expert_data.py` 将 Waymo 数据转换为训练用 `.pkl`，输出到 `data/training_data/`。
+
+- **多智能体**：
+```bash
+python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100 --start_index 0
+```
+
+- **单智能体（仅 ego）**：
+```bash
+python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100 --start_index 0 --ego_only
+```
+
+### 2. 行为克隆 (BC)
+BC 支持两种模式：**多智能体**（默认，所有受控车共用同一策略）与 **单智能体**（仅 ego 车，评估时其他车按专家轨迹回放）。
+
+- **多智能体训练**（模型保存到 `models/bc/`，日志到 `logs/bc/`）：
+```bash
+python train_bc.py --expert_data_path data/training_data/expert_data_0_50.pkl --epochs 100
+```
+
+- **单智能体训练**（使用 ego-only 数据，评估时仅 ego 受策略控制，其他车专家回放）：
+```bash
+python train_bc.py --expert_data_path data/training_data/expert_data_ego_0_50.pkl --epochs 100 --single_agent
+```
+
+- **可视化**：`python scripts/visualize.py policy --policy_type bc --model_path models/bc/policy_best.pt`  
+  仅自车用策略、其他车回放（单智能体可视化）：加 `--ego_only`，例如  
+  `python scripts/visualize.py policy --policy_type bc --model_path models/bc/policy_best.pt --ego_only --num_scenarios 1`
+
+### 3. 多智能体对抗模仿学习 (MAGAIL)
+- **训练**：`python train_magail.py`（模型保存到 `models/magail/`，日志到 `logs/magail/`）
+- **可视化**：`python scripts/visualize.py policy --policy_type magail --model_path models/magail/model_50_actor.pth`
+
+### 4. 可视化统一入口
+可视化统一使用 `scripts/visualize.py`，子命令：`replay`（场景回放）、`policy`（BC/MAGAIL 策略）、`trajectory`（专家轨迹 2D 动画）。详见 [scripts/README.md](scripts/README.md)。
+
+## 文件与模块职责
+
+### 根目录脚本
+- **train_bc.py**：BC 训练，从 `dataset.loader` 加载专家 pkl，模型与日志写入 `models/bc/`、`logs/bc/`
+- **train_magail.py**：MAGAIL 训练，环境使用 `BCScenarioEnv`（45 维），从 `dataset.loader` 加载专家数据，模型与日志写入 `models/magail/`、`logs/magail/`
+
+### Env 模块
+- **Env/bc_env.py**：`BCScenarioEnv`，45 维观测（Ego 5 维 + 10 邻居×4 维），BC 与 MAGAIL 训练/评估共用
+- **Env/bc_ego_replay_env.py**：`BCEgoReplayEnv`，单智能体 BC 评估环境，仅 ego 受策略控制，其他车按专家轨迹回放
+- **Env/scenario_env.py**：`MultiAgentScenarioEnv` 基类，Waymo 场景加载与步进
+- **Env/expert_replay_env.py**：专家轨迹回放与逆动力学动作，供 `generate_expert_data.py` 与回放可视化
+- **Env/inverse_dynamics.py**：轨迹 → 油门/转向动作
+
+### Algorithm 模块
+- **Algorithm/policy.py**：`StateIndependentPolicy`，BC 使用的 MLP 策略
+
+### scripts 目录
+工具脚本用途与用法见 [scripts/README.md](scripts/README.md)。

algorithms/__init__.py

@@ -0,0 +1,2 @@
+"""Compatibility package for legacy HBBC checkpoints."""
+

algorithms/utils.py

@@ -0,0 +1,18 @@
+import numpy as np
+
+
+class RunningMeanStd(object):
+    def __init__(self, epsilon=1e-4, shape=()):
+        self.mean = np.zeros(shape, np.float64)
+        self.var = np.ones(shape, np.float64)
+        self.count = epsilon
+
+
+class Normalizer(RunningMeanStd):
+    def __init__(self, input_dim, epsilon=1e-4, clip_obs=10.0):
+        super().__init__(shape=input_dim)
+        self.epsilon = epsilon
+        self.clip_obs = clip_obs
+
+    def normalize(self, input):
+        return np.clip((input - self.mean) / np.sqrt(self.var + self.epsilon), -self.clip_obs, self.clip_obs)

dataset/expert_dataset.py

@@ -244,6 +244,7 @@ class ExpertTrajectoryDataset(Dataset):
                 print(f"  观测维度: {obs_dim} (应为107)")
         
         if save_path:
+            os.makedirs(os.path.dirname(save_path), exist_ok=True)
             with open(save_path, "wb") as f:
                 pickle.dump({
                     "trajectories": all_trajectories,
@@ -282,7 +283,7 @@ if __name__ == "__main__":
         trajectories, observations = ExpertTrajectoryDataset.collect_with_full_obs(
             env_config,
             num_scenarios=10,
-            save_path="./expert_trajectories_full.pkl"
+            save_path="data/trajectories/expert_trajectories_full.pkl"
         )
         
         if len(trajectories) > 0:

dataset/loader.py

@@ -0,0 +1,164 @@
+"""
+统一数据加载：BC/MAGAIL 训练用专家 pkl 的加载函数与 Dataset。
+主训练流水线使用本模块；dataset/expert_dataset.py 为可选 107 维/5 维管线。
+"""
+import os
+import glob
+import pickle
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+
+def load_expert_pkl(expert_data_path, *, filter_terminal_last_step: bool = False, agent_id_filter=None):
+    """从目录或单个 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.
+        agent_id_filter: If not None, only load trajectories with traj[\"agent_id\"] == agent_id_filter
+            (e.g. \"default_agent\" for single-agent/ego-only).
+    """
+    if os.path.isdir(expert_data_path):
+        pkl_files = glob.glob(os.path.join(expert_data_path, "*.pkl"))
+        if not pkl_files:
+            raise FileNotFoundError(f"No .pkl files in {expert_data_path}")
+        print(f"Found {len(pkl_files)} pickle files in {expert_data_path}")
+    elif os.path.exists(expert_data_path):
+        pkl_files = [expert_data_path]
+    else:
+        raise FileNotFoundError(f"Expert data path not found: {expert_data_path}")
+
+    obs_data, act_data = [], []
+    for pkl_file in pkl_files:
+        try:
+            with open(pkl_file, "rb") as f:
+                data = pickle.load(f)
+            if isinstance(data, list):
+                for traj in data:
+                    if agent_id_filter is not None and traj.get("agent_id") != agent_id_filter:
+                        continue
+                    if "obs" in traj and "acts" in traj:
+                        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["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:
+            print(f"Error loading {pkl_file}: {e}")
+
+    if len(obs_data) == 0:
+        raise ValueError("No valid data loaded from provided path.")
+    obs_data = np.concatenate(obs_data, axis=0)
+    act_data = np.concatenate(act_data, axis=0)
+    print(f"Total loaded samples: {len(obs_data)}")
+    return obs_data, act_data
+
+
+def get_expert_scenario_ids(expert_data_path, max_ids=10):
+    """
+    从专家 pkl 中收集出现过的 scenario_id（这些场景在采集时曾有受控车）。
+    用于 eval 时只在这些场景上评估，保证 eval 有受控车。
+    返回排序后的 list，最多 max_ids 个。
+    """
+    if os.path.isdir(expert_data_path):
+        pkl_files = glob.glob(os.path.join(expert_data_path, "*.pkl"))
+    elif os.path.exists(expert_data_path):
+        pkl_files = [expert_data_path]
+    else:
+        return []
+
+    seen = set()
+    for pkl_file in pkl_files:
+        try:
+            with open(pkl_file, "rb") as f:
+                data = pickle.load(f)
+            if isinstance(data, list):
+                for traj in data:
+                    if "scenario_id" in traj:
+                        seen.add(traj["scenario_id"])
+            # dict 格式通常没有 per-trajectory scenario_id，跳过
+        except Exception:
+            continue
+    out = sorted(seen)[:max_ids]
+    return out
+
+
+class MAGAILExpertDataset(Dataset):
+    def __init__(self, data_dir, transform=None, *, filter_terminal_last_step: bool = False, agent_id_filter=None):
+        """
+        Args:
+            data_dir (str): Directory containing .pkl files from generate_expert_data.py
+            transform (callable, optional): Optional transform to be applied on a sample.
+            agent_id_filter: If not None, only load trajectories with traj[\"agent_id\"] == agent_id_filter.
+        """
+        self.data_dir = data_dir
+        self.transform = transform
+        self.trajectories = []
+        self.flat_data = []  # (obs, act) pairs
+
+        # Load all .pkl files
+        pkl_files = glob.glob(os.path.join(data_dir, "*.pkl"))
+        print(f"Loading data from {len(pkl_files)} files in {data_dir}...")
+
+        for pkl_file in pkl_files:
+            try:
+                with open(pkl_file, "rb") as f:
+                    data = pickle.load(f)
+                    # data is a list of dicts: {'obs': (T, 45), 'acts': (T, 2), ...}
+                    if agent_id_filter is not None:
+                        data = [t for t in data if t.get("agent_id") == agent_id_filter]
+                    self.trajectories.extend(data)
+            except Exception as e:
+                print(f"Error loading {pkl_file}: {e}")
+
+        # Flatten for training Discriminator/BC
+        print(f"Processing {len(self.trajectories)} trajectories...")
+        for traj in self.trajectories:
+            obs = traj["obs"]
+            acts = traj["acts"]
+
+            # obs: (T, 45), acts: (T, 2)
+            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)}")
+
+    def __len__(self):
+        return len(self.flat_data)
+
+    def __getitem__(self, idx):
+        obs, act = self.flat_data[idx]
+
+        obs = torch.from_numpy(obs).float()
+        act = torch.from_numpy(act).float()
+
+        sample = {"state": obs, "action": act}
+
+        if self.transform:
+            sample = self.transform(sample)
+
+        return sample

dataset/magail_dataset.py

@@ -1,61 +0,0 @@
-import torch
-from torch.utils.data import Dataset
-import pickle
-import numpy as np
-import os
-import glob
-
-class MAGAILExpertDataset(Dataset):
-    def __init__(self, data_dir, transform=None):
-        """
-        Args:
-            data_dir (str): Directory containing .pkl files from generate_expert_data.py
-            transform (callable, optional): Optional transform to be applied on a sample.
-        """
-        self.data_dir = data_dir
-        self.transform = transform
-        self.trajectories = []
-        self.flat_data = [] # (obs, act) pairs
-        
-        # Load all .pkl files
-        pkl_files = glob.glob(os.path.join(data_dir, "*.pkl"))
-        print(f"Loading data from {len(pkl_files)} files in {data_dir}...")
-        
-        for pkl_file in pkl_files:
-            try:
-                with open(pkl_file, 'rb') as f:
-                    data = pickle.load(f)
-                    # data is a list of dicts: {'obs': (T, 45), 'acts': (T, 2), ...}
-                    self.trajectories.extend(data)
-            except Exception as e:
-                print(f"Error loading {pkl_file}: {e}")
-                
-        # Flatten for training Discriminator/BC
-        print(f"Processing {len(self.trajectories)} trajectories...")
-        for traj in self.trajectories:
-            obs = traj['obs']
-            acts = traj['acts']
-            
-            # obs: (T, 45), acts: (T, 2)
-            # We pair them up
-            for i in range(len(obs)):
-                self.flat_data.append((obs[i], acts[i]))
-                
-        print(f"Total samples: {len(self.flat_data)}")
-        
-    def __len__(self):
-        return len(self.flat_data)
-
-    def __getitem__(self, idx):
-        obs, act = self.flat_data[idx]
-        
-        # Convert to tensor
-        obs = torch.from_numpy(obs).float()
-        act = torch.from_numpy(act).float()
-        
-        sample = {'state': obs, 'action': act}
-        
-        if self.transform:
-            sample = self.transform(sample)
-            
-        return sample

docs/HBBC_Deploy_guied.md

@@ -0,0 +1,439 @@
+# HBBC 策略部署指南
+
+本文档说明如何将 `weights/hbbc.pt` 部署到 MetaDrive 项目中的**背景车辆**上，作为车辆控制策略使用。
+
+---
+
+## 0. 本仓库适配说明（MAGAIL4AutoDrive）
+
+本仓库已落地一套可直接使用的 HBBC 背景车接入实现，核心代码：
+
+- `Env/hbbc_actor_critic.py`：HBBC 所需 `ActorCritic` 最小推理网络
+- `Env/hbbc_background_policy.py`：模型加载、18 维观测构建、latent 管理（含 JSON 覆盖）
+- `Env/bc_env.py`：`BCScenarioEnv` 动态背景车 HBBC 接入（静态背景车保持不变）
+- `Env/bc_ego_replay_env.py`：`BCEgoReplayEnv` 动态背景车 HBBC 接入（ego-only 评估兼容）
+
+与原文档示例不同点：
+
+1. 当前仓库 `BaseVehicle` 没有 `pos_buffer/rot_buffer/action_buffer`，因此 8 维 `base_state` 使用当前可得车辆状态重建；
+2. 仅动态背景车使用 HBBC，静态背景车仍作为占位/邻居车辆；
+3. 支持通过 JSON 手动指定场景中某些车辆的 latent（`object_id` / `agent_id` 双 key）。
+
+---
+
+## 1. 概述
+
+### 1.1 HBBC 是什么
+
+**HBBC**（Hierarchical Behavior-Based Controller）是一个低层驾驶策略网络，输入车辆状态和行为条件，输出连续控制动作 `[steering, acceleration]`，可直接用于 MetaDrive 的车辆控制。
+
+### 1.2 依赖
+
+- **PyTorch**
+- **NumPy**
+- **MetaDrive**（需包含 `BaseVehicle`、`BasePolicy` 等基础组件）
+
+---
+
+## 2. 模型加载
+
+### 2.1 模型架构
+
+HBBC 对应 `ActorCritic` 网络，需按以下参数实例化：
+
+```python
+import torch
+from algorithms.modules import ActorCritic  # 或复制 actor_critic.py 到目标项目
+
+hbbc = ActorCritic(
+    num_actor_obs=18,
+    num_critic_obs=18,
+    num_actions=2,
+    latent_c_dim=4,      # 行为模式数
+    latent_eps_dim=6,    # 风格向量维度
+    use_style_latent=True,
+).to(device)
+
+# 加载权重
+checkpoint = torch.load("path/to/hbbc.pt", map_location=device, weights_only=False)
+hbbc.load_state_dict(checkpoint['actor_critic'])
+hbbc.eval()
+```
+
+### 2.2 推理接口
+
+```python
+with torch.no_grad():
+    actions = hbbc.act_inference(obs_tensor)  # obs_tensor: (batch, 18), 输出: (batch, 2)
+```
+
+---
+
+## 3. 输入规格（18 维）
+
+HBBC 的输入为 `hbbc_obs`，维度 18，由三部分拼接：
+
+```
+hbbc_obs = [base_state(8) | latent_eps(6) | latent_c(4)]
+```
+
+### 3.1 base_state（8 维）
+
+从车辆对象构建，需按**精确顺序**拼接。实现如下（需配合 `relative_pos_local`、`rot_matrix_inv`、`clip` 等工具函数）：
+
+```python
+import numpy as np
+
+def build_hbbc_base_state(vehicle):
+    """
+    从 MetaDrive 车辆对象构建 HBBC 的 8 维 base_state。
+    要求 vehicle 具有: position, pos_buffer, rot_buffer, heading_buffer,
+    speed_km_h, max_speed_km_h, eps_step, acceleration, yaw_rate, action_buffer
+    """
+    from metadrive.utils.math import clip  # 或 np.clip
+    
+    veh_pos = list(vehicle.position) + [0]
+    init_veh_rot = np.array([vehicle.rot_buffer[0][0], vehicle.rot_buffer[0][1], vehicle.rot_buffer[0][2]])
+    init_veh_pos = list(vehicle.pos_buffer[0]) + [0]
+    init_veh_heading = vehicle.heading_buffer[0]
+
+    # 局部位置（本实现中置 0）
+    veh_pos_local = relative_pos_local(init_veh_pos, veh_pos, init_veh_rot)[:2]
+    veh_pos_local[0] /= 10
+    veh_pos_local[1] /= 2
+
+    # 局部航向（本实现中置 0）
+    veh_heading = vehicle.heading
+    cross = np.cross(init_veh_heading, veh_heading)
+    dot = np.dot(init_veh_heading, veh_heading)
+    veh_heading_local = np.arctan2(cross, dot)
+
+    veh_vel = clip((vehicle.speed_km_h + 1) / (vehicle.max_speed_km_h + 1), 0.0, 1.0)
+    veh_acc = vehicle.acceleration / 5 if vehicle.eps_step > 1 else 0
+    yaw_rate = vehicle.yaw_rate
+    last_action_0 = vehicle.action_buffer[-1][0]
+    last_action_1 = vehicle.action_buffer[-1][1]
+
+    # 8 维，顺序固定
+    obs = np.concatenate((
+        veh_pos_local * 0,           # 2 维，置 0
+        [veh_heading_local * 0],     # 1 维，置 0
+        [veh_vel],                   # 1 维
+        [veh_acc * 0],               # 1 维，置 0
+        [yaw_rate * 0.5],            # 1 维
+        [last_action_0], [last_action_1]  # 2 维
+    )).astype(np.float32)
+    return obs
+```
+
+### 3.2 latent_eps（6 维）
+
+风格向量，需 **L2 归一化** 且在 `[-1, 1]` 内：
+
+```python
+# 随机采样（每个 episode 或每辆车可固定/随机）
+latent_eps = np.random.randn(6).astype(np.float32)
+latent_eps = latent_eps / (np.linalg.norm(latent_eps) + 1e-8)
+latent_eps = np.clip(latent_eps, -1.0, 1.0)
+```
+
+### 3.3 latent_c（4 维）
+
+行为模式 one-hot，4 选 1：
+
+```python
+# 随机选一个模式 (0~3)
+mode = np.random.randint(0, 4)
+latent_c = np.zeros(4, dtype=np.float32)
+latent_c[mode] = 1.0
+```
+
+### 3.4 完整观测拼接
+
+```python
+def build_hbbc_obs(vehicle, latent_eps, latent_c):
+    base = build_hbbc_base_state(vehicle)
+    return np.concatenate([base, latent_eps, latent_c], axis=-1)  # shape: (18,)
+```
+
+---
+
+## 4. 必需工具函数
+
+若目标项目无以下函数，需自行实现或从 styledrive 的 `envs/utils.py` 拷贝：
+
+```python
+def rot_matrix(t):
+    """t: [roll, pitch, yaw], 返回 3x3 旋转矩阵"""
+    roll, pitch, yaw = t[0], t[1], t[2]
+    sr, cr = np.sin(roll), np.cos(roll)
+    sp, cp = np.sin(pitch), np.cos(pitch)
+    sy, cy = np.sin(yaw), np.cos(yaw)
+    r_roll = np.array([[1, 0, 0], [0, cr, -sr], [0, sr, cr]])
+    r_pitch = np.array([[cp, 0, sp], [0, 1, 0], [-sp, 0, cp]])
+    r_yaw = np.array([[cy, -sy, 0], [sy, cy, 0], [0, 0, 1]])
+    return np.dot(np.dot(r_yaw, r_pitch), r_roll)
+
+def rot_matrix_inv(t):
+    return rot_matrix(t).T
+
+def relative_pos_local(coord, coord_t, veh_rot):
+    """将 coord_t 从世界坐标变换到以 coord 为原点、veh_rot 为姿态的局部坐标"""
+    r_pos_global = np.array(coord_t) - np.array(coord)
+    rot_mat_inv = rot_matrix_inv(veh_rot)
+    return rot_mat_inv @ r_pos_global
+```
+
+`clip` 可用 `np.clip` 或 `metadrive.utils.math.clip`。
+
+---
+
+## 5. 车辆属性要求
+
+使用 HBBC 的车辆需继承或兼容 MetaDrive 的 `BaseVehicle`，并具备：
+
+| 属性 | 说明 |
+|------|------|
+| `position` | 当前位置 (x, y) 或 (x, y, z) |
+| `heading` | 航向单位向量 |
+| `heading_theta` | 航向角（弧度） |
+| `pos_buffer` | `deque`，至少 1 个元素，`pos_buffer[0]` 为 episode 起始位姿 |
+| `rot_buffer` | `deque`，`(roll, pitch, yaw)`，`rot_buffer[0]` 为起始姿态 |
+| `heading_buffer` | `deque`，`heading_buffer[0]` 为起始航向 |
+| `action_buffer` | `deque`，`action_buffer[-1]` 为上一时刻动作 `(steering, acc)` |
+| `speed_km_h` | 当前速度 km/h |
+| `max_speed_km_h` | 最大速度 km/h |
+| `acceleration` | 当前加速度 |
+| `yaw_rate` | 偏航角速度 (rad/s) |
+| `eps_step` | 本 episode 的步数 |
+| `last_heading_theta` | 上一帧航向角（用于 yaw_rate） |
+
+`BaseVehicle` 在 `before_step` 中会更新 `pos_buffer`、`rot_buffer`、`heading_buffer`、`action_buffer`，只要在配置中设置 `veh_obs_len >= 1`（建议 3–10）即可。
+
+---
+
+## 6. 输出动作格式
+
+HBBC 输出 2 维连续动作，与 MetaDrive 动作空间一致：
+
+```python
+# actions: (2,) 或 (batch, 2)
+# actions[0]: steering  ∈ [-1, 1]
+# actions[1]: acceleration ∈ [-1, 1]，正=油门，负=刹车
+```
+
+环境会在 `_preprocess_actions` 中做限幅与平滑，无需在策略内再次裁剪。
+
+---
+
+## 7. 部署为 MetaDrive 策略（背景车）
+
+### 7.1 自定义 Policy
+
+实现一个继承 `BasePolicy` 的策略，在 `act` 中调用 HBBC：
+
+```python
+from metadrive.policy.base_policy import BasePolicy
+import torch
+import numpy as np
+
+class HBBCPolicy(BasePolicy):
+    def __init__(self, control_object, random_seed=None, hbbc_path="weights/hbbc.pt", device="cpu"):
+        super().__init__(control_object, random_seed)
+        self.device = torch.device(device)
+        self.hbbc = self._load_hbbc(hbbc_path)
+        self.latent_eps = None
+        self.latent_c = None
+        self._resample_latent()
+
+    def _load_hbbc(self, path):
+        from algorithms.modules import ActorCritic  # 根据实际路径调整
+        model = ActorCritic(
+            num_actor_obs=18, num_critic_obs=18, num_actions=2,
+            latent_c_dim=4, latent_eps_dim=6, use_style_latent=True
+        ).to(self.device)
+        ckpt = torch.load(path, map_location=self.device, weights_only=False)
+        model.load_state_dict(ckpt['actor_critic'])
+        model.eval()
+        return model
+
+    def _resample_latent(self):
+        self.latent_eps = np.random.randn(6).astype(np.float32)
+        self.latent_eps = self.latent_eps / (np.linalg.norm(self.latent_eps) + 1e-8)
+        self.latent_eps = np.clip(self.latent_eps, -1.0, 1.0)
+        mode = np.random.randint(0, 4)
+        self.latent_c = np.zeros(4, dtype=np.float32)
+        self.latent_c[mode] = 1.0
+
+    def act(self, agent_id=None):
+        vehicle = self.control_object
+        base_state = build_hbbc_base_state(vehicle)
+        obs = np.concatenate([base_state, self.latent_eps, self.latent_c], axis=-1)
+        obs_t = torch.tensor(obs, dtype=torch.float32, device=self.device).unsqueeze(0)
+        with torch.no_grad():
+            actions = self.hbbc.act_inference(obs_t).cpu().numpy().squeeze()
+        self.action_info["action"] = actions.tolist()
+        return [float(actions[0]), float(actions[1])]
+
+    def reset(self):
+        super().reset()
+        self._resample_latent()
+```
+
+### 7.2 配置背景车使用 HBBC
+
+在环境配置中为背景车辆指定 `HBBCPolicy`：
+
+```python
+config = {
+    # ...
+    "agent_configs": {
+        "agent0": {
+            "policy": HBBCPolicy,
+            "policy_kwargs": {"hbbc_path": "path/to/hbbc.pt", "device": "cuda:0"},
+        }
+    },
+    # 若使用 traffic 的 policy 配置方式，则需在 traffic 管理逻辑中
+    # 将部分或全部背景车的 policy 替换为 HBBCPolicy
+}
+```
+
+若背景车由 TrafficManager 等模块统一管理，需在该模块的 policy 选择逻辑中加入对 `HBBCPolicy` 的分配。
+
+### 7.3 与 TrafficManager 集成
+
+若背景车由 `PGTrafficManager` 等生成，需在添加策略时改为使用 `HBBCPolicy`：
+
+```python
+# 原代码通常为:
+# self.add_policy(random_v.id, IDMPolicy, random_v, self.generate_seed())
+
+# 改为:
+from your_policy_module import HBBCPolicy
+self.add_policy(random_v.id, HBBCPolicy, random_v, self.generate_seed(),
+                hbbc_path="path/to/hbbc.pt", device="cuda:0")
+```
+
+`add_policy` 的额外参数会传给 Policy 的 `__init__`。若接口不支持传参，可修改 `HBBCPolicy` 从全局配置读取路径，或使用自定义 TrafficManager 子类。
+
+**注意**：HBBC 在 styledrive 中基于 scenario 轨迹训练，不包含路由逻辑。背景车若需要沿车道/路线行驶，可能需：
+- 在项目中为 HBBC 车辆配置 `navigation`，或
+- 仅对部分背景车使用 HBBC（如混合 IDM + HBBC），或
+- 在目标项目中验证 HBBC 在开放道路上的表现后决定是否全量使用。
+
+### 7.4 注意事项
+
+1. **latent 生命周期**：可为每辆车在 spawn 时采样一次，或在每个 episode reset 时重采样。
+2. **首帧 action_buffer**：首步 `action_buffer[-1]` 通常为 `(0, 0)`，由 `BaseVehicle` 初始化保证。
+3. **同步更新 buffer**：车辆必须在每步调用 `before_step` 之类接口，更新 `pos_buffer`、`action_buffer` 等，否则观测会错位。
+4. **veh_obs_len**：车辆配置中设置 `veh_obs_len >= 3`（建议 10），确保 buffer 长度足够。
+
+---
+
+## 8. ActorCritic 网络定义（可移植）
+
+若目标项目无法导入 styledrive 的 `algorithms`，可把以下简化版 `ActorCritic` 放到本项目中单独使用：
+
+```python
+import torch
+import torch.nn as nn
+
+def get_activation(name):
+    return getattr(nn, name)()
+
+class ActorCritic(nn.Module):
+    def __init__(self, num_actor_obs=18, num_critic_obs=18, num_actions=2,
+                 latent_c_dim=4, latent_eps_dim=6, use_style_latent=True,
+                 actor_hidden_dims=[512, 256, 128], activation='elu'):
+        super().__init__()
+        act_fn = getattr(nn, activation.upper())()
+        self.latent_c_dim = latent_c_dim
+        self.latent_eps_dim = latent_eps_dim
+        self.use_style_latent = use_style_latent
+
+        layers = []
+        layers.append(nn.Linear(num_actor_obs, actor_hidden_dims[0]))
+        layers.append(act_fn)
+        for i in range(len(actor_hidden_dims) - 1):
+            layers.append(nn.Linear(actor_hidden_dims[i], actor_hidden_dims[i + 1]))
+            layers.append(act_fn)
+        self.actor_trunk = nn.Sequential(*layers)
+        self.actor_head = nn.Linear(actor_hidden_dims[-1], num_actions)
+
+        if use_style_latent:
+            style_layers = [nn.Linear(latent_eps_dim, 512), act_fn,
+                           nn.Linear(512, 256), act_fn, nn.Linear(256, 128), act_fn]
+            self.style_trunk = nn.Sequential(*style_layers)
+            self.style_head = nn.Linear(128, latent_eps_dim)
+            self.style_activation = torch.tanh
+
+    def act_inference(self, observations):
+        if self.use_style_latent:
+            obs = observations[..., :-(self.latent_c_dim + self.latent_eps_dim)]
+            eps = observations[..., -self.latent_c_dim - self.latent_eps_dim:-self.latent_c_dim]
+            c = observations[..., -self.latent_c_dim:]
+            eps = self.style_activation(self.style_head(self.style_trunk(eps)))
+            observations = torch.cat([obs, eps, c], dim=-1)
+        embedding = self.actor_trunk(observations)
+        return self.actor_head(embedding)
+```
+
+加载与调用方式与前面一致。
+
+---
+
+## 9. 简要检查清单
+
+- [ ] 正确加载 `hbbc.pt` 的 `actor_critic` 权重
+- [ ] `build_hbbc_base_state` 输出 8 维，顺序与文档一致
+- [ ] `latent_eps` 6 维、L2 归一化
+- [ ] `latent_c` 4 维 one-hot
+- [ ] 车辆具备 `pos_buffer`、`rot_buffer`、`heading_buffer`、`action_buffer` 等属性
+- [ ] 策略返回 `[steering, acceleration]`，范围 [-1, 1]
+- [ ] 每步更新上述 buffer，保证观测连续
+
+---
+
+## 10. 本仓库配置项与 JSON 示例
+
+可通过环境配置控制 HBBC 背景车行为：
+
+- `enable_hbbc_background`：是否启用动态背景车 HBBC（`True/False`）
+- `hbbc_model_path`：模型路径（默认 `models/hbbc/hbbc.pt`）
+- `hbbc_inference_device`：推理设备（如 `cpu` / `cuda:0`）
+- `hbbc_latent_mode`：`per_vehicle_fixed` 或 `per_episode_reset`
+- `hbbc_latent_json_path`：可选，手动 latent JSON 路径
+
+`hbbc_latent_json_path` 内容格式（优先按 `object_id` 匹配，失败回退 `agent_id`）：
+
+```json
+{
+  "global": {
+    "latent_eps": [0.35, -0.12, 0.28, 0.46, -0.22, 0.18],
+    "latent_c": [0, 0, 1, 0]
+  },
+  "object_id": {
+    "12345": {
+      "latent_eps": [0.2, -0.1, 0.3, 0.4, -0.2, 0.1],
+      "latent_c": [0, 1, 0, 0]
+    }
+  },
+  "agent_id": {
+    "controlled_abcde": {
+      "latent_eps": [0.5, 0.1, -0.1, 0.2, -0.3, 0.4],
+      "latent_c": [1, 0, 0, 0]
+    }
+  }
+}
+```
+
+匹配优先级为：`object_id` > `agent_id` > `global` > 随机采样。  
+`latent_eps` 会做 L2 归一化，`latent_c` 会强制 one-hot；非法输入会告警并回退随机采样。
+
+---
+
+## 11. 参考来源
+
+- 策略与观测：`envs/ad_hbbc_gym.py` 中的 `ADObservation.vehicle_state`
+- 模型：`algorithms/modules/actor_critic.py` 中 `ActorCritic`
+- 工具：`envs/utils.py` 中的 `relative_pos_local`、`rot_matrix`、`rot_matrix_inv`

docs/examples/hbbc_latent_example.json

@@ -0,0 +1,18 @@
+{
+  "global": {
+    "latent_eps": [0.35, -0.12, 0.28, 0.46, -0.22, 0.18],
+    "latent_c": [0, 1,1, 0]
+  },
+  "object_id": {
+    "12345": {
+      "latent_eps": [0.2, -0.1, 0.3, 0.4, -0.2, 0.1],
+      "latent_c": [0, 1, 0, 0]
+    }
+  },
+  "agent_id": {
+    "controlled_abcde": {
+      "latent_eps": [0.5, 0.1, -0.1, 0.2, -0.3, 0.4],
+      "latent_c": [1, 0, 0, 0]
+    }
+  }
+}

scripts/README.md

@@ -0,0 +1,105 @@
+# scripts 工具脚本说明
+
+本目录包含数据生成、回放、可视化与分析等工具脚本。训练脚本（`train_bc.py`、`train_magail.py`）位于项目根目录。
+
+## 路径约定（相对项目根）
+
+- **数据**：`data/exp_filtered`（Waymo 场景）、`data/training_data`（专家 pkl 输出）
+- **模型**：`models/bc/`（BC）、`models/magail/`（MAGAIL）
+- **日志**：`logs/bc/`、`logs/magail/`（TensorBoard）
+
+---
+
+## 脚本列表与用法
+
+### 数据生成
+
+| 脚本 | 用途 | 用法示例 |
+|------|------|----------|
+| [generate_expert_data.py](generate_expert_data.py) | 从 Waymo 数据生成专家 (obs, act) 的 pkl | 见下方 |
+
+**多智能体**（输出 `expert_data_{start_index}_{num_scenarios}.pkl`）：
+```bash
+python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100 --start_index 0
+```
+
+**单智能体**（仅采集 ego 车轨迹，输出 `expert_data_ego_{start_index}_{num_scenarios}.pkl`，用于单智能体 BC）：
+```bash
+python scripts/generate_expert_data.py --data_dir data/exp_filtered --output_dir data/training_data --num_scenarios 100 --start_index 0 --ego_only
+```
+
+**常用参数**：`--data_dir`（默认 `data/exp_filtered`）、`--output_dir`（默认 `data/training_data`）、`--start_index`、`--num_scenarios`、`--ego_only`（仅保存 default_agent 轨迹，输出使用 `expert_data_ego_*.pkl` 前缀）。
+
+---
+
+### 可视化（统一入口）
+
+| 脚本 | 用途 | 用法示例 |
+|------|------|----------|
+| [visualize.py](visualize.py) | **replay**：场景回放（ExpertReplayEnv）；**policy**：BC/MAGAIL 策略；**trajectory**：专家轨迹 2D 动画 | 见下方 |
+
+**子命令**：
+
+- **replay**（原始专家轨迹回放）：
+```bash
+python scripts/visualize.py replay --data_dir data/exp_filtered --num_scenarios 1 --horizon 500
+```
+
+- **policy**（BC 或 MAGAIL 训练策略）：与专家数据生成/回放一致——同一套车道+静态筛选、且会生成背景车（bg_*），使观测分布与训练集一致，便于在训练集上公平演示。
+```bash
+python scripts/visualize.py policy --policy_type bc --model_path models/bc/policy_best.pt --data_dir data/exp_filtered --num_scenarios 1
+python scripts/visualize.py policy --policy_type magail --model_path models/magail/model_50_actor.pth --num_scenarios 1 --deterministic
+```
+- **policy + 仅自车策略、其他车回放**（BC 单智能体模型）：加 `--ego_only`，自车由策略控制，其余车辆按专家轨迹回放。
+```bash
+python scripts/visualize.py policy --policy_type bc --model_path models/bc/policy_best.pt --data_dir data/exp_filtered --num_scenarios 1 --ego_only
+```
+
+- **policy + HBBC 动态背景车**（仅动态背景车启用，静态背景车保持原样）：
+```bash
+python scripts/visualize.py policy \
+  --policy_type bc \
+  --model_path models/bc/policy_best.pt \
+  --data_dir data/exp_filtered \
+  --num_scenarios 1 \
+  --ego_only \
+  --enable_hbbc_background \
+  --hbbc_model_path models/hbbc/hbbc.pt \
+  --hbbc_inference_device cpu \
+  --hbbc_latent_mode per_vehicle_fixed \
+  --hbbc_latent_json_path docs/examples/hbbc_latent_example.json
+```
+
+- **trajectory**（专家轨迹 matplotlib 俯视图动画）：
+```bash
+python scripts/visualize.py trajectory --data_dir data/exp_filtered --scenario_idx 0
+```
+
+**公共参数**：`--data_dir`（默认 `data/exp_filtered`）、`--start_index`、`--num_scenarios`、`--horizon`。policy 模式另有 `--policy_type`（auto/bc/magail）、`--model_path`、`--deterministic`（仅 MAGAIL）、`--ego_only`（仅 BC：自车用策略，其他车专家回放）、`--enable_hbbc_background`、`--hbbc_model_path`、`--hbbc_inference_device`、`--hbbc_latent_mode`、`--hbbc_latent_json_path`。
+
+---
+
+### 数据分析与检查
+
+| 脚本 | 用途 | 用法示例 |
+|------|------|----------|
+| [analyze_expert_data.py](analyze_expert_data.py) | 分析专家数据分布与统计 | 见脚本内 `__main__`（依赖 env 与数据目录配置） |
+| [check_track_fields.py](check_track_fields.py) | 检查 Waymo 轨迹字段 | 见脚本内 `__main__` |
+| [check_database_info.py](check_database_info.py) | 检查数据库/场景信息 | 见脚本内 `__main__`（含硬编码路径，可按需改为 `data/exp_filtered`） |
+
+---
+
+### 其他
+
+| 脚本 | 用途 | 用法示例 |
+|------|------|----------|
+| [launch_tensorboard.py](launch_tensorboard.py) | 启动 TensorBoard | `python scripts/launch_tensorboard.py --logdir logs`（或 `logs/bc` / `logs/magail`） |
+
+---
+
+## 与训练流程的对应关系
+
+1. **数据准备**：`generate_expert_data.py` → 输出到 `data/training_data/*.pkl`（多智能体 `expert_data_*.pkl`，单智能体 `expert_data_ego_*.pkl`）
+2. **BC 训练**：根目录 `train_bc.py` → 模型保存到 `models/bc/`，日志到 `logs/bc/`。单智能体模式加 `--single_agent` 并指定 ego-only 的 pkl。
+3. **MAGAIL 训练**：根目录 `train_magail.py` → 模型保存到 `models/magail/`，日志到 `logs/magail/`
+4. **可视化**：`scripts/visualize.py`（子命令 replay / policy / trajectory）→ 数据目录默认 `data/exp_filtered`

scripts/generate_expert_data.py

@@ -102,7 +102,9 @@ def generate_data(args):
             
             # Post-process episode data
             for agent_id, data in episode_data.items():
-                if len(data['obs']) > 10: # Minimum length filter
+                if args.ego_only and agent_id != "default_agent":
+                    continue
+                if len(data['obs']) > 10:  # Minimum length filter
                     expert_trajectories.append({
                         'obs': np.array(data['obs']),
                         'acts': np.array(data['acts']),
@@ -120,9 +122,14 @@ def generate_data(args):
         pass
         
     # Save data
-    output_file = os.path.join(args.output_dir, f"expert_data_{args.start_index}_{args.num_scenarios}.pkl")
+    if args.ego_only:
+        output_file = os.path.join(args.output_dir, f"expert_data_ego_{args.start_index}_{args.num_scenarios}.pkl")
+    else:
+        output_file = os.path.join(args.output_dir, f"expert_data_{args.start_index}_{args.num_scenarios}.pkl")
     os.makedirs(args.output_dir, exist_ok=True)
 
+    if args.ego_only:
+        print("Ego-only mode: saved trajectories are SDC (default_agent) only.")
     print(f"Saving {len(expert_trajectories)} trajectories to {output_file}")
     with open(output_file, 'wb') as f:
         pickle.dump(expert_trajectories, f)
@@ -153,10 +160,10 @@ def generate_data(args):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--data_dir", type=str, default="/home/huangfukk/MAGAIL4AutoDrive/data/exp_filtered", help="Path to Waymo pickles (or filtered index)")
-    parser.add_argument("--output_dir", type=str, default="/home/huangfukk/MAGAIL4AutoDrive/data/training", help="Output directory")
+    parser.add_argument("--data_dir", type=str, default="data/exp_filtered", help="Path to Waymo pickles (or filtered index)")
+    parser.add_argument("--output_dir", type=str, default="data/training_data", help="Output directory")
     parser.add_argument("--start_index", type=int, default=0)
     parser.add_argument("--num_scenarios", type=int, default=10)
-    
+    parser.add_argument("--ego_only", action="store_true", help="Only collect and save ego (default_agent) trajectories; output uses expert_data_ego_*.pkl prefix")
     args = parser.parse_args()
     generate_data(args)

scripts/launch_tensorboard.py

@@ -0,0 +1,18 @@
+import sys
+import types
+import os
+
+# Mock imghdr module for Python 3.13 compatibility
+# TensorBoard depends on imghdr which was removed in Python 3.13
+if sys.version_info >= (3, 13):
+    if 'imghdr' not in sys.modules:
+        imghdr_mock = types.ModuleType('imghdr')
+        imghdr_mock.what = lambda filename, h=None: None
+        # Mock tests list which tensorboard appends to
+        imghdr_mock.tests = []
+        sys.modules['imghdr'] = imghdr_mock
+
+from tensorboard import main as tb_main
+
+if __name__ == '__main__':
+    sys.exit(tb_main.run_main())

scripts/visualize.py

@@ -0,0 +1,432 @@
+"""
+Unified visualization: replay (scenario replay), policy (BC/MAGAIL), trajectory (2D expert trajectory animation).
+Usage: python scripts/visualize.py <replay|policy|trajectory> [args...]
+"""
+import argparse
+import os
+import sys
+import time
+import numpy as np
+import torch
+
+project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if project_root not in sys.path:
+    sys.path.insert(0, project_root)
+
+# --- Replay ---
+def _run_replay(args):
+    from Env.expert_replay_env import ExpertReplayEnv
+
+    data_path = os.path.abspath(args.data_dir)
+    if not os.path.exists(data_path):
+        raise ValueError(f"Data directory {data_path} not found")
+
+    from metadrive.scenario.utils import read_dataset_summary
+    _, summary_lookup, _ = read_dataset_summary(data_path)
+    if args.start_index >= len(summary_lookup):
+        raise ValueError(
+            f"start_index={args.start_index} out of range. Dataset has {len(summary_lookup)} scenarios."
+        )
+    max_available = len(summary_lookup) - args.start_index
+    num_to_run = min(args.num_scenarios, max_available)
+
+    print(f"Initializing ExpertReplayEnv with data from {data_path}...")
+
+    try:
+        for i in range(args.start_index, args.start_index + num_to_run):
+            print(f"\n--- Playing Scenario {i} ---")
+            # Each scenario uses a fresh env (start_scenario_index=i, num_scenarios=1) so the second
+            # scenario and beyond are fully cleaned and loaded like a single-scenario run.
+            env_config = {
+                "data_directory": data_path,
+                "is_multi_agent": True,
+                "num_controlled_agents": 100,
+                "horizon": args.horizon,
+                "use_render": True,
+                "sequential_seed": True,
+                "reactive_traffic": False,
+                "start_scenario_index": i,
+                "num_scenarios": 1,
+                "log_level": 40,
+            }
+            env = ExpertReplayEnv(config=env_config)
+            try:
+                obs = env.reset(seed=i)
+            except Exception as e:
+                print(f"Error resetting scenario {i}: {e}")
+                env.close()
+                continue
+
+            print(f"Scenario loaded. Controlled agents (current): {len(env.controlled_agents)}, total in scenario: {env.num_controlled_in_scenario}")
+
+            for step in range(args.horizon):
+                obs, rewards, dones, infos = env.step(None)
+                env.render(
+                    mode="top_down",
+                    text={"Step": step, "Agents": len(env.controlled_agents), "Scenario": i},
+                )
+                time.sleep(0.05)
+                if dones["__all__"]:
+                    print(f"Scenario {i} finished at step {step}")
+                    break
+            env.close()
+    except KeyboardInterrupt:
+        print("Interrupted by user")
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        print(f"Global error: {e}")
+    finally:
+        print("Environment closed.")
+
+
+# --- Policy (BC / MAGAIL) ---
+def _resolve_data_dir(data_dir_arg):
+    if data_dir_arg:
+        data_dir = data_dir_arg
+    else:
+        data_dir = os.path.join(project_root, "data", "exp_filtered")
+        if not os.path.exists(data_dir):
+            data_dir = os.path.join(project_root, "data", "exp_converted")
+    if not os.path.exists(data_dir):
+        raise FileNotFoundError(f"Data directory not found at {data_dir}. Please specify --data_dir.")
+    return data_dir
+
+
+def _resolve_model_path(model_path, policy_type):
+    if os.path.exists(model_path):
+        return model_path
+    if policy_type == "bc":
+        candidate = os.path.join(project_root, "models", "bc", os.path.basename(model_path))
+    else:
+        candidate = os.path.join(project_root, "models", "magail", os.path.basename(model_path))
+    if os.path.exists(candidate):
+        return candidate
+    if policy_type == "magail" and not model_path.endswith("_actor.pth"):
+        candidate = os.path.join(project_root, "models", "magail", os.path.basename(model_path) + "_actor.pth")
+        if os.path.exists(candidate):
+            return candidate
+    raise FileNotFoundError(f"Model path {model_path} not found.")
+
+
+def _run_policy(args):
+    from Env.bc_env import BCScenarioEnv
+    from Env.bc_ego_replay_env import BCEgoReplayEnv
+    from metadrive.engine.engine_utils import close_engine
+
+    policy_type = (args.policy_type or "auto").lower()
+    if policy_type == "auto":
+        policy_type = "bc" if args.model_path.endswith(".pt") else "magail"
+    ego_only = getattr(args, "ego_only", False)
+    if ego_only and policy_type != "bc":
+        print("[WARN] --ego_only is supported for BC policy only; MAGAIL will run in multi-agent mode.")
+
+    data_dir = _resolve_data_dir(args.data_dir)
+    data_path = os.path.abspath(data_dir)
+    env_config = {
+        "data_directory": data_path,
+        "is_multi_agent": True,
+        "num_controlled_agents": 100 if ego_only else 3,
+        "horizon": args.horizon,
+        "use_render": True,
+        "sequential_seed": True,
+        "start_scenario_index": args.start_index,
+        "num_scenarios": args.num_scenarios,
+        "log_level": 40,
+        "enable_hbbc_background": bool(getattr(args, "enable_hbbc_background", False)),
+        "hbbc_model_path": getattr(args, "hbbc_model_path", "models/hbbc/hbbc.pt"),
+        "hbbc_inference_device": getattr(args, "hbbc_inference_device", "cpu"),
+        "hbbc_latent_mode": getattr(args, "hbbc_latent_mode", "per_vehicle_fixed"),
+        "hbbc_latent_json_path": getattr(args, "hbbc_latent_json_path", None),
+    }
+
+    if ego_only and policy_type == "bc":
+        print("Initializing BCEgoReplayEnv (ego-only: policy on self, others replayed)...")
+    else:
+        print(f"Initializing BCScenarioEnv (policy_type={policy_type})...")
+
+    try:
+        env = BCEgoReplayEnv(config=env_config) if (ego_only and policy_type == "bc") else BCScenarioEnv(env_config, agent2policy={})
+    except Exception as e:
+        print(f"Error init env: {e}. Trying to close lingering engine...")
+        try:
+            close_engine()
+        except Exception:
+            pass
+        env = BCEgoReplayEnv(config=env_config) if (ego_only and policy_type == "bc") else BCScenarioEnv(env_config, agent2policy={})
+
+    state_dim = 45
+    action_dim = 2
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model_path = _resolve_model_path(args.model_path, policy_type)
+    print(f"Loading model from {model_path}...")
+
+    if policy_type == "bc":
+        from Algorithm.policy import StateIndependentPolicy
+        policy = StateIndependentPolicy(
+            state_shape=(state_dim,),
+            action_shape=(action_dim,),
+            hidden_units=(256, 256),
+            hidden_activation=torch.nn.Tanh(),
+        ).to(device)
+        try:
+            state = torch.load(model_path, map_location=device, weights_only=True)
+        except TypeError:
+            state = torch.load(model_path, map_location=device)
+        policy.load_state_dict(state)
+        policy.eval()
+    else:
+        from train_magail import Actor
+        actor = Actor(state_dim, action_dim).to(device)
+        actor.load_state_dict(torch.load(model_path, map_location=device))
+        actor.eval()
+
+    try:
+        for i in range(args.start_index, args.start_index + args.num_scenarios):
+            print(f"\n--- Playing Scenario {i} ---")
+            try:
+                obs_dict = env.reset(seed=i)
+            except Exception as e:
+                print(f"Error resetting {i}: {e}. Skipping.")
+                try:
+                    close_engine()
+                    env = BCScenarioEnv(env_config, agent2policy={})
+                except Exception:
+                    pass
+                continue
+
+            n_total = getattr(env, "num_controlled_in_scenario", len(obs_dict))
+            mode_note = " (ego only, others replayed)" if (ego_only and policy_type == "bc") else ""
+            if ego_only and policy_type == "bc" and bool(env_config.get("enable_hbbc_background", False)):
+                mode_note = " (ego only, dynamic background via HBBC)"
+            print(f"Scenario loaded. Controlled agents (current): {len(obs_dict)}, total in scenario: {n_total}{mode_note}")
+            if ego_only and policy_type == "bc" and len(obs_dict) == 1:
+                if bool(env_config.get("enable_hbbc_background", False)):
+                    print("  [Ego control: policy injected — dynamic background vehicles use HBBC; static background stays static.]")
+                else:
+                    print("  [Ego control: policy injected — ego uses model output each step; other vehicles expert replay.]")
+            if len(obs_dict) == 0:
+                print(f"Scenario {i} has no controlled agents (all filtered out). Skipping.")
+                continue
+            step_count = 0
+            episode_reward = 0.0
+
+            while True:
+                agent_ids = list(obs_dict.keys())
+                obs_list = [obs_dict[aid] for aid in agent_ids]
+                obs_tensor = torch.FloatTensor(np.array(obs_list)).to(device)
+
+                with torch.no_grad():
+                    if policy_type == "bc":
+                        actions_np = policy(obs_tensor).cpu().numpy()
+                    else:
+                        dist = actor(obs_tensor)
+                        if args.deterministic:
+                            actions_np = torch.tanh(dist.mean).cpu().numpy()
+                        else:
+                            actions_np = torch.tanh(dist.sample()).cpu().numpy()
+
+                actions = {aid: actions_np[idx].flatten() for idx, aid in enumerate(agent_ids)}
+                obs_dict, rewards, dones, infos = env.step(actions)
+                episode_reward += sum(rewards.values())
+
+                env.render(
+                    mode="top_down",
+                    text={
+                        "Scenario": i,
+                        "Step": step_count,
+                        "Agents": len(obs_dict),
+                        "Total Reward": f"{episode_reward:.2f}",
+                    },
+                )
+                step_count += 1
+
+                if dones["__all__"] or step_count >= args.horizon:
+                    print(f"Scenario finished at step {step_count}, reward {episode_reward:.2f}")
+                    break
+    except KeyboardInterrupt:
+        print("Interrupted.")
+    finally:
+        env.close()
+
+
+# --- Trajectory (matplotlib 2D animation) ---
+def _build_expert_trajectories_from_env(env):
+    """Build expert_trajectories dict from env (ExpertReplayEnv has traffic_manager.current_traffic_data)."""
+    if hasattr(env, "expert_trajectories") and env.expert_trajectories:
+        return env.expert_trajectories
+    if not hasattr(env, "engine") or not hasattr(env.engine, "traffic_manager"):
+        return {}
+    from metadrive.type import MetaDriveType
+    data = getattr(env.engine.traffic_manager, "current_traffic_data", None)
+    if not data:
+        return {}
+    expert_trajs = {}
+    for scenario_id, track in data.items():
+        if track.get("type") != MetaDriveType.VEHICLE or "state" not in track:
+            continue
+        state = track["state"]
+        positions = state.get("position")
+        if positions is None:
+            continue
+        valid = state.get("valid", np.ones(len(positions), dtype=bool))
+        valid = np.asarray(valid).flatten()
+        if valid.size != len(positions):
+            valid = np.ones(len(positions), dtype=bool)
+        first_show = int(np.argmax(valid)) if valid.any() else 0
+        last_show = len(valid) - 1 - int(np.argmax(valid[::-1])) if valid.any() else len(positions) - 1
+        obj_id = track.get("metadata", {}).get("object_id", str(scenario_id))
+        expert_trajs[obj_id] = {
+            "positions": np.asarray(positions),
+            "start_timestep": first_show,
+            "end_timestep": last_show,
+        }
+    return expert_trajs
+
+
+def _run_trajectory_animation(expert_trajs, scenario_idx):
+    import matplotlib.pyplot as plt
+    from matplotlib.animation import FuncAnimation
+
+    if len(expert_trajs) == 0:
+        print("No expert trajectories to visualize.")
+        return
+
+    fig, ax = plt.subplots(figsize=(12, 12))
+    max_timestep = max(t["end_timestep"] for t in expert_trajs.values())
+    min_timestep = min(t["start_timestep"] for t in expert_trajs.values())
+
+    colors = plt.cm.tab10(np.linspace(0, 1, len(expert_trajs)))
+    for idx, (obj_id, traj) in enumerate(expert_trajs.items()):
+        positions = np.asarray(traj["positions"])
+        if positions.ndim >= 2:
+            positions = positions[:, :2]
+        else:
+            continue
+        ax.plot(
+            positions[:, 0], positions[:, 1],
+            color=colors[idx], alpha=0.3, linewidth=1,
+            label=f"Vehicle {str(obj_id)[:6]}",
+        )
+
+    scatter = ax.scatter([], [], s=200, c="red", marker="o", edgecolors="black", linewidths=2)
+    time_text = ax.text(0.02, 0.95, "", transform=ax.transAxes, fontsize=14)
+    ax.set_xlabel("X (m)")
+    ax.set_ylabel("Y (m)")
+    ax.set_title(f"Expert Trajectory Visualization - Scenario {scenario_idx}")
+    ax.legend(loc="upper right", fontsize=8)
+    ax.grid(True, alpha=0.3)
+    ax.axis("equal")
+
+    def update(frame):
+        current_time = min_timestep + frame
+        current_positions = []
+        for traj in expert_trajs.values():
+            st, et = traj["start_timestep"], traj["end_timestep"]
+            if st <= current_time <= et:
+                pos = np.asarray(traj["positions"])
+                if pos.ndim >= 2:
+                    pos = pos[current_time - st, :2]
+                else:
+                    continue
+                current_positions.append(pos)
+        if current_positions:
+            scatter.set_offsets(np.array(current_positions))
+        time_text.set_text(f"Time: {frame * 0.1:.1f}s (Frame {frame})")
+        return scatter, time_text
+
+    anim = FuncAnimation(
+        fig, update, frames=max_timestep - min_timestep + 1,
+        interval=100, blit=True, repeat=True,
+    )
+    plt.tight_layout()
+    plt.show()
+    return anim
+
+
+def _run_trajectory(args):
+    from Env.expert_replay_env import ExpertReplayEnv
+
+    data_dir = _resolve_data_dir(args.data_dir)
+    data_path = os.path.abspath(data_dir)
+    env_config = {
+        "data_directory": data_path,
+        "is_multi_agent": True,
+        "num_controlled_agents": 100,
+        "horizon": 500,
+        "use_render": False,
+        "sequential_seed": True,
+        "reactive_traffic": False,
+        "start_scenario_index": args.scenario_idx,
+        "num_scenarios": 1,
+        "log_level": 40,
+    }
+
+    env = ExpertReplayEnv(config=env_config)
+    try:
+        env.reset(seed=args.scenario_idx)
+        expert_trajs = _build_expert_trajectories_from_env(env)
+        _run_trajectory_animation(expert_trajs, args.scenario_idx)
+    finally:
+        env.close()
+
+
+# --- Main ---
+def main():
+    parser = argparse.ArgumentParser(
+        description="Unified visualization: replay, policy (BC/MAGAIL), trajectory.",
+    )
+    subparsers = parser.add_subparsers(dest="mode", required=True, help="replay | policy | trajectory")
+
+    # Common args for data_dir (used by all)
+    def add_common_data_args(p):
+        p.add_argument("--data_dir", type=str, default="data/exp_filtered", help="Waymo scenario directory")
+        p.add_argument("--start_index", type=int, default=0)
+        p.add_argument("--num_scenarios", type=int, default=1)
+        p.add_argument("--horizon", type=int, default=200)
+
+    # replay
+    pr = subparsers.add_parser("replay", help="Replay scenario with ExpertReplayEnv (no policy)")
+    add_common_data_args(pr)
+    pr.set_defaults(horizon=500)
+
+    # policy
+    pp = subparsers.add_parser("policy", help="Visualize BC or MAGAIL trained policy")
+    add_common_data_args(pp)
+    pp.add_argument("--policy_type", type=str, default="auto", choices=["auto", "bc", "magail"])
+    pp.add_argument("--model_path", type=str, default="models/bc/policy_best.pt")
+    pp.add_argument("--deterministic", action="store_true", help="MAGAIL: use mean action")
+    pp.add_argument("--ego_only", action="store_true", help="BC only: inject policy into ego only; other vehicles use expert replay")
+    pp.add_argument("--enable_hbbc_background", action="store_true", help="Enable HBBC policy for dynamic background vehicles")
+    pp.add_argument("--hbbc_model_path", type=str, default="models/hbbc/hbbc.pt")
+    pp.add_argument("--hbbc_inference_device", type=str, default="cpu")
+    pp.add_argument("--hbbc_latent_mode", type=str, default="per_vehicle_fixed", choices=["per_vehicle_fixed", "per_episode_reset"])
+    pp.add_argument("--hbbc_latent_json_path", type=str, default=None, help="Optional JSON for per-vehicle latent override")
+
+    # trajectory
+    pt = subparsers.add_parser("trajectory", help="2D matplotlib animation of expert trajectories")
+    pt.add_argument("--data_dir", type=str, default="data/exp_filtered")
+    pt.add_argument("--scenario_idx", type=int, default=0)
+
+    args = parser.parse_args()
+
+    # Resolve data_dir relative to project root when default
+    if args.mode != "trajectory":
+        if args.data_dir in ("data/exp_filtered", "data/exp_converted"):
+            args.data_dir = os.path.join(project_root, args.data_dir)
+    else:
+        if args.data_dir in ("data/exp_filtered", "data/exp_converted"):
+            args.data_dir = os.path.join(project_root, args.data_dir)
+
+    if args.mode == "replay":
+        _run_replay(args)
+    elif args.mode == "policy":
+        _run_policy(args)
+    elif args.mode == "trajectory":
+        _run_trajectory(args)
+    else:
+        parser.error(f"Unknown mode: {args.mode}")
+
+
+if __name__ == "__main__":
+    main()

scripts/visualize_expert_trajectory.py

@@ -1,105 +0,0 @@
-import sys
-import os
-
-# 添加路径
-current_dir = os.path.dirname(os.path.abspath(__file__))
-project_root = os.path.dirname(current_dir)
-env_dir = os.path.join(project_root, "Env")
-sys.path.insert(0, project_root)
-sys.path.insert(0, env_dir)
-
-# 现在可以导入了
-from scenario_env import MultiAgentScenarioEnv
-from metadrive.engine.asset_loader import AssetLoader
-import numpy as np
-import matplotlib.pyplot as plt
-from matplotlib.animation import FuncAnimation
-
-class DummyPolicy:
-    """
-    占位策略,用于数据检查时初始化环境
-    不需要实际执行动作,只是为了满足环境初始化要求
-    """
-    def act(self, *args, **kwargs):
-        # 返回零动作 [throttle, steering]
-        return np.array([0.0, 0.0])
-
-def visualize_expert_trajectory(env, scenario_idx=0):
-    """
-    可视化专家轨迹的俯视图动画
-    """
-    env.reset()
-    expert_trajs = env.expert_trajectories
-    
-    if len(expert_trajs) == 0:
-        print("当前场景无专家轨迹")
-        return
-    
-    # 设置绘图
-    fig, ax = plt.subplots(figsize=(12, 12))
-    
-    # 获取所有轨迹的最大时间长度
-    max_timestep = max(traj["end_timestep"] for traj in expert_trajs.values())
-    min_timestep = min(traj["start_timestep"] for traj in expert_trajs.values())
-    
-    # 绘制完整轨迹(淡色)
-    colors = plt.cm.tab10(np.linspace(0, 1, len(expert_trajs)))
-    for idx, (obj_id, traj) in enumerate(expert_trajs.items()):
-        positions = traj["positions"][:, :2]
-        ax.plot(positions[:, 0], positions[:, 1], 
-                color=colors[idx], alpha=0.3, linewidth=1,
-                label=f'Vehicle {obj_id[:6]}')
-    
-    # 初始化当前位置标记
-    scatter = ax.scatter([], [], s=200, c='red', marker='o', edgecolors='black', linewidths=2)
-    time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes, fontsize=14)
-    
-    ax.set_xlabel('X (m)')
-    ax.set_ylabel('Y (m)')
-    ax.set_title(f'Expert Trajectory Visualization - Scenario {scenario_idx}')
-    ax.legend(loc='upper right', fontsize=8)
-    ax.grid(True, alpha=0.3)
-    ax.axis('equal')
-    
-    def update(frame):
-        current_time = min_timestep + frame
-        
-        # 收集当前时间所有车辆的位置
-        current_positions = []
-        for traj in expert_trajs.values():
-            if traj["start_timestep"] <= current_time <= traj["end_timestep"]:
-                idx = current_time - traj["start_timestep"]
-                pos = traj["positions"][idx, :2]
-                current_positions.append(pos)
-        
-        if len(current_positions) > 0:
-            current_positions = np.array(current_positions)
-            scatter.set_offsets(current_positions)
-        
-        time_text.set_text(f'Time: {frame * 0.1:.1f}s (Frame {frame})')
-        return scatter, time_text
-    
-    anim = FuncAnimation(fig, update, frames=max_timestep-min_timestep+1,
-                        interval=100, blit=True, repeat=True)
-    
-    plt.tight_layout()
-    plt.show()
-    
-    return anim
-
-if __name__ == "__main__":
-    WAYMO_DATA_DIR = r"/home/huangfukk/mdsn"
-    data_dir = AssetLoader.file_path(WAYMO_DATA_DIR, "exp_filtered", unix_style=False)
-    
-    env = MultiAgentScenarioEnv(
-        config={
-            "data_directory": data_dir,
-            "is_multi_agent": True,
-            "num_controlled_agents": 3,
-            "use_render": False,
-        },
-        agent2policy=DummyPolicy()
-    )
-    
-    # 可视化第一个场景
-    anim = visualize_expert_trajectory(env, scenario_idx=0)

scripts/visualize_replay.py

@@ -1,93 +0,0 @@
-import argparse
-import os
-import sys
-import time
-
-# Add project root to Python path so we can import Env module
-project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-if project_root not in sys.path:
-    sys.path.insert(0, project_root)
-
-from Env.expert_replay_env import ExpertReplayEnv
-
-def visualize_replay(args):
-    data_path = os.path.abspath(args.data_dir)
-    if not os.path.exists(data_path):
-        raise ValueError(f"Data directory {data_path} not found")
-
-    # Same as data generation: avoid MetaDrive assertion when requested num_scenarios > available.
-    from metadrive.scenario.utils import read_dataset_summary
-    _, summary_lookup, _ = read_dataset_summary(data_path)
-    if args.start_index >= len(summary_lookup):
-        raise ValueError(
-            f"start_index={args.start_index} out of range. Dataset has {len(summary_lookup)} scenarios."
-        )
-    max_available = len(summary_lookup) - args.start_index
-    num_to_run = min(args.num_scenarios, max_available)
-
-    env_config = {
-        "data_directory": data_path,
-        "is_multi_agent": True,
-        "num_controlled_agents": 100,
-        "horizon": args.horizon,
-        "use_render": True,  # Enable rendering
-        "sequential_seed": True,
-        "reactive_traffic": False,
-        "start_scenario_index": args.start_index,
-        "num_scenarios": -1,
-        "log_level": 40, # ERROR
-        # "pstats": True, # For performance debugging
-    }
-
-    print(f"Initializing ExpertReplayEnv with data from {data_path}...")
-    env = ExpertReplayEnv(config=env_config)
-
-    try:
-        for i in range(args.start_index, args.start_index + num_to_run):
-            print(f"\n--- Playing Scenario {i} ---")
-            try:
-                obs = env.reset(seed=i)
-            except Exception as e:
-                print(f"Error resetting scenario {i}: {e}")
-                continue
-            
-            print(f"Scenario loaded. Controlled agents: {len(env.controlled_agents)}")
-            
-            for step in range(args.horizon):
-                # Step
-                obs, rewards, dones, infos = env.step(None)
-                
-                # Render
-                env.render(mode="top_down", 
-                           text={
-                               "Step": step, 
-                               "Agents": len(env.controlled_agents),
-                               "Scenario": i
-                           })
-                
-                # Sleep to control playback speed
-                time.sleep(0.05)
-                
-                if dones["__all__"]:
-                    print(f"Scenario {i} finished at step {step}")
-                    break
-                    
-    except KeyboardInterrupt:
-        print("Interrupted by user")
-    except Exception as e:
-        import traceback
-        traceback.print_exc()
-        print(f"Global error: {e}")
-    finally:
-        env.close()
-        print("Environment closed.")
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--data_dir", type=str, default="/home/huangfukk/MAGAIL4AutoDrive/data/exp_filtered", help="Path to Waymo data")
-    parser.add_argument("--start_index", type=int, default=0)
-    parser.add_argument("--num_scenarios", type=int, default=1)
-    parser.add_argument("--horizon", type=int, default=500)
-    
-    args = parser.parse_args()
-    visualize_replay(args)

train_bc.py

@@ -0,0 +1,212 @@
+"""
+BC 训练脚本：负责数据加载、环境评估、日志与保存；BC 算法由 Algorithm.bc 提供。
+使用方式不变：python train_bc.py [--expert_data_path data/training_data] [--save_dir models/bc] ...
+"""
+import os
+import numpy as np
+import torch
+import argparse
+from torch.utils.data import DataLoader, TensorDataset
+from torch.optim import Adam
+from torch.optim.lr_scheduler import ExponentialLR
+from datetime import datetime
+from torch.utils.tensorboard import SummaryWriter
+
+from Algorithm.policy import StateIndependentPolicy
+from Algorithm.bc import train_bc_epoch, eval_bc_epoch
+from Env.bc_env import BCScenarioEnv
+from Env.bc_ego_replay_env import BCEgoReplayEnv
+from dataset.loader import load_expert_pkl, get_expert_scenario_ids
+
+
+def evaluate_policy(policy, args, device):
+    """在 BCScenarioEnv（多智能体）或 BCEgoReplayEnv（单智能体）中评估策略。
+    仅使用专家数据中出现过的 scenario_id。单智能体模式下仅 ego 受策略控制，其他车专家回放。"""
+    waymo_data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
+    data_dir = os.path.join(waymo_data_dir, "exp_filtered")
+    if not os.path.exists(data_dir):
+        data_dir = os.path.join(waymo_data_dir, "exp_converted")
+    if not os.path.exists(data_dir):
+        print(f"[ERROR] Could not find scenario data in {waymo_data_dir}. Evaluation skipped.")
+        return 0.0, 0.0, 0.0
+
+    scenario_ids = get_expert_scenario_ids(args.expert_data_path, max_ids=5)
+    if not scenario_ids:
+        print("[WARN] No scenario_id in expert pkl, falling back to scenarios [0,1,2]. Eval may have 0 controlled agents.")
+        scenario_ids = [0, 1, 2]
+
+    total_rewards = []
+    total_steps = []
+    collision_episodes = 0
+    horizon = 200
+    single_agent = getattr(args, "single_agent", False)
+
+    for idx, scenario_id in enumerate(scenario_ids):
+        env_config = {
+            "data_directory": data_dir,
+            "is_multi_agent": True,
+            "num_controlled_agents": 100,
+            "use_render": False,
+            "sequential_seed": True,
+            "horizon": horizon,
+            "start_scenario_index": scenario_id,
+            "num_scenarios": 1,
+            "log_level": 50,
+        }
+        if single_agent:
+            env = BCEgoReplayEnv(config=env_config)
+        else:
+            env = BCScenarioEnv(env_config, agent2policy=None)
+        try:
+            obs_dict = env.reset(seed=scenario_id)
+        except Exception as e:
+            print(f"  Eval Episode {idx} (scenario {scenario_id}): reset failed: {e}")
+            env.close()
+            continue
+
+        n_controlled = len(env.controlled_agents)
+        n_total_in_scenario = getattr(env, "num_controlled_in_scenario", n_controlled) if not single_agent else 1
+        if n_controlled == 0:
+            print(
+                f"  Eval Episode {idx} (scenario {scenario_id}): 0 controlled agents (total in scenario: {n_total_in_scenario}), skip."
+            )
+            env.close()
+            continue
+
+        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
+            if not obs_dict:
+                obs_dict, _, dones, _ = env.step({})
+                continue
+            agent_ids = list(obs_dict.keys())
+            obs_list = [obs_dict[aid] for aid in agent_ids]
+            obs_tensor = torch.FloatTensor(np.array(obs_list)).to(device)
+            with torch.no_grad():
+                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, 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
+        mode_str = "single-agent (ego)" if single_agent else f"agents (current): {n_controlled}, total in scenario: {n_total_in_scenario}"
+        print(
+            f"  Eval Episode {idx} (scenario {scenario_id}): Total Reward {episode_reward:.2f}, steps {step_count}, {mode_str}"
+        )
+        env.close()
+
+    if not total_rewards:
+        print("  No valid eval episodes (all skipped or failed).")
+        return 0.0, 0.0, 0.0
+    avg_reward = float(np.mean(total_rewards))
+    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):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    os.makedirs("logs/bc", exist_ok=True)
+    log_dir = os.path.join("logs", "bc", datetime.now().strftime("%Y%m%d-%H%M%S"))
+    writer = SummaryWriter(log_dir)
+    print(f"TensorBoard logging to: {log_dir}")
+    os.makedirs(args.save_dir, exist_ok=True)
+
+    agent_id_filter = "default_agent" if getattr(args, "single_agent", False) else None
+    obs_data, act_data = load_expert_pkl(
+        args.expert_data_path,
+        filter_terminal_last_step=args.filter_terminal_last_step,
+        agent_id_filter=agent_id_filter,
+    )
+    obs_tensor = torch.FloatTensor(obs_data)
+    act_tensor = torch.FloatTensor(act_data)
+    dataset = TensorDataset(obs_tensor, act_tensor)
+    train_size = int(0.8 * len(dataset))
+    val_size = len(dataset) - train_size
+    train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])
+    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False)
+    print(f"Dataset loaded. Train size: {len(train_dataset)}, Val size: {len(val_dataset)}")
+
+    state_dim = obs_data.shape[1]
+    action_dim = act_data.shape[1]
+    print(f"State Dim: {state_dim}, Action Dim: {action_dim}")
+
+    policy = StateIndependentPolicy(
+        state_shape=(state_dim,),
+        action_shape=(action_dim,),
+        hidden_units=(256, 256),
+        hidden_activation=torch.nn.Tanh(),
+    ).to(device)
+    optimizer = Adam(policy.parameters(), lr=args.lr)
+    scheduler = ExponentialLR(optimizer, gamma=0.99)
+
+    best_val_loss = float("inf")
+    for epoch in range(args.epochs):
+        avg_train_loss = train_bc_epoch(policy, train_loader, optimizer, device)
+        scheduler.step()
+        avg_val_loss = eval_bc_epoch(policy, val_loader, device)
+
+        print(f"Epoch {epoch+1}/{args.epochs} | Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f}")
+        writer.add_scalar("Loss/train", avg_train_loss, epoch)
+        writer.add_scalar("Loss/val", avg_val_loss, epoch)
+        writer.add_scalar("Learning_rate", scheduler.get_last_lr()[0], epoch)
+
+        if avg_val_loss < best_val_loss:
+            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, 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()
+    print("Training finished.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--expert_data_path", type=str, default="data/training_data", help="Path to expert data pickle or directory")
+    parser.add_argument("--save_dir", type=str, default="models/bc", help="Directory to save models")
+    parser.add_argument("--epochs", type=int, default=100)
+    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).",
+    )
+    parser.add_argument(
+        "--single_agent",
+        action="store_true",
+        help="Use single-agent (ego) expert data and evaluation; load only default_agent trajectories and evaluate with BCEgoReplayEnv.",
+    )
+    args = parser.parse_args()
+    main(args)

train_magail.py

@@ -9,7 +9,8 @@ import argparse
 import signal
 import sys
 from torch.utils.data import DataLoader
-from dataset.magail_dataset import MAGAILExpertDataset
+from dataset.loader import MAGAILExpertDataset
+from Env.bc_env import BCScenarioEnv
 
 # --- Networks ---
 
@@ -79,18 +80,30 @@ class PPO:
         self.K_epochs = K_epochs
         self.mse_loss = nn.MSELoss()
 
+    def _log_prob_from_dist(self, dist, pre_tanh_action):
+        # Tanh-squashed Gaussian log-prob with correction term.
+        log_prob = dist.log_prob(pre_tanh_action)
+        correction = torch.log(1 - torch.tanh(pre_tanh_action) ** 2 + 1e-6)
+        return (log_prob - correction).sum(dim=-1)
+
     def select_action(self, state):
         with torch.no_grad():
             state = torch.FloatTensor(state).cuda()
             dist = self.actor(state)
-            action = dist.sample()
-            action_logprob = dist.log_prob(action).sum(dim=-1)
-        return action.cpu().numpy(), action_logprob.cpu().numpy()
+            pre_tanh_action = dist.sample()
+            action = torch.tanh(pre_tanh_action)
+            action_logprob = self._log_prob_from_dist(dist, pre_tanh_action)
+        return (
+            action.cpu().numpy(),
+            action_logprob.cpu().numpy(),
+            pre_tanh_action.cpu().numpy()
+        )
 
     def update(self, memory):
         # Convert memory to tensors
         states = torch.FloatTensor(np.array(memory['states'])).cuda()
         actions = torch.FloatTensor(np.array(memory['actions'])).cuda()
+        pre_tanh_actions = torch.FloatTensor(np.array(memory['pre_tanh_actions'])).cuda()
         logprobs = torch.FloatTensor(np.array(memory['logprobs'])).cuda()
         rewards = torch.FloatTensor(np.array(memory['rewards'])).cuda()
         next_states = torch.FloatTensor(np.array(memory['next_states'])).cuda()
@@ -124,7 +137,7 @@ class PPO:
         for _ in range(self.K_epochs):
             # Evaluating old actions and values :
             dist = self.actor(states)
-            action_logprobs = dist.log_prob(actions).sum(dim=-1)
+            action_logprobs = self._log_prob_from_dist(dist, pre_tanh_actions)
             dist_entropy = dist.entropy().sum(dim=-1)
             state_values = self.critic(states).squeeze()
             
@@ -152,12 +165,7 @@ class PPO:
 # --- Training Loop ---
 
 def train(args):
-    # 1. Setup Environment (Dummy for now, usually you run simulation here)
-    # But for MAGAIL we need to collect generated trajectories.
-    # We need the Env class to be importable.
-    from Env.scenario_env import MultiAgentScenarioEnv
-    from Env.simple_idm_policy import ConstantVelocityPolicy # Just for init
-    
+    # 1. Setup Environment (45-dim obs via BCScenarioEnv)
     # Config for Env
     env_config = {
         "data_directory": args.data_dir,
@@ -200,12 +208,7 @@ def train(args):
                 yield batch
     expert_iter = cycle(expert_loader)
     
-    # 4. Initialize Env
-    from Env.expert_replay_env import ExpertReplayEnv # Using ReplayEnv for config, but we need ScenarioEnv for simulation?
-    # Actually we need MultiAgentScenarioEnv for interactive training, not Replay.
-    from Env.scenario_env import MultiAgentScenarioEnv
-    from Env.simple_idm_policy import ConstantVelocityPolicy # Placeholder policy for init
-    
+    # 4. Initialize Env (BCScenarioEnv provides 45-dim obs)
     # 2. Setup Models
     # Determine state dim from environment if possible, or use fixed
     # Expert data has 45 dim? 
@@ -220,48 +223,48 @@ def train(args):
     # We need to inject that same logic into the training env, OR
     # subclass MultiAgentScenarioEnv in the training script to override observation.
     
-    class MAGAILScenarioEnv(MultiAgentScenarioEnv):
-        def _get_all_obs(self):
-            # Same logic as ExpertReplayEnv to ensure compatibility
-            obs_dict = {}
-            for agent_id, vehicle in self.controlled_agents.items():
-                # 1. Ego State
-                ego_state = [
-                    vehicle.position[0], vehicle.position[1],
-                    vehicle.velocity[0], vehicle.velocity[1],
-                    vehicle.heading_theta
-                ]
+    # class MAGAILScenarioEnv(MultiAgentScenarioEnv):
+    #     def _get_all_obs(self):
+    #         # Same logic as ExpertReplayEnv to ensure compatibility
+    #         obs_dict = {}
+    #         for agent_id, vehicle in self.controlled_agents.items():
+    #             # 1. Ego State
+    #             ego_state = [
+    #                 vehicle.position[0], vehicle.position[1],
+    #                 vehicle.velocity[0], vehicle.velocity[1],
+    #                 vehicle.heading_theta
+    #             ]
+    #             
+    #             # 2. Neighbors
+    #             candidates = []
+    #             for other_id, other_vehicle in self.engine.agent_manager.active_agents.items():
+    #                 if other_id == agent_id:
+    #                     continue
+    #                 dist = np.linalg.norm(vehicle.position - other_vehicle.position)
+    #                 if dist < 30.0:
+    #                     candidates.append((dist, other_vehicle))
+    #             
+    #             candidates.sort(key=lambda x: x[0])
+    #             top_10 = candidates[:10]
+    #             
+    #             neighbor_feats = []
+    #             for _, neighbor in top_10:
+    #                 neighbor_feats.extend([
+    #                     neighbor.position[0] - vehicle.position[0], 
+    #                     neighbor.position[1] - vehicle.position[1],
+    #                     neighbor.velocity[0], 
+    #                     neighbor.velocity[1]
+    #                 ])
+    #                 
+    #             missing = 10 - len(top_10)
+    #             if missing > 0:
+    #                 neighbor_feats.extend([0.0] * (4 * missing))
+    #                 
+    #             obs = np.array(ego_state + neighbor_feats, dtype=np.float32)
+    #             obs_dict[agent_id] = obs
+    #         return obs_dict
     
-                # 2. Neighbors
-                candidates = []
-                for other_id, other_vehicle in self.engine.agent_manager.active_agents.items():
-                    if other_id == agent_id:
-                        continue
-                    dist = np.linalg.norm(vehicle.position - other_vehicle.position)
-                    if dist < 30.0:
-                        candidates.append((dist, other_vehicle))
-                
-                candidates.sort(key=lambda x: x[0])
-                top_10 = candidates[:10]
-                
-                neighbor_feats = []
-                for _, neighbor in top_10:
-                    neighbor_feats.extend([
-                        neighbor.position[0] - vehicle.position[0], 
-                        neighbor.position[1] - vehicle.position[1],
-                        neighbor.velocity[0], 
-                        neighbor.velocity[1]
-                    ])
-                    
-                missing = 10 - len(top_10)
-                if missing > 0:
-                    neighbor_feats.extend([0.0] * (4 * missing))
-                    
-                obs = np.array(ego_state + neighbor_feats, dtype=np.float32)
-                obs_dict[agent_id] = obs
-            return obs_dict
-
-    env = MAGAILScenarioEnv(config=env_config, agent2policy={}) # Pass empty dict if we control all externally
+    env = BCScenarioEnv(env_config, agent2policy={})  # 45-dim obs
     
     print("Starting training...")
     
@@ -277,7 +280,15 @@ def train(args):
     
     for i_episode in range(args.max_episodes):
         # --- 1. Collect Rollouts (Interaction) ---
-        memory = {'states': [], 'actions': [], 'logprobs': [], 'rewards': [], 'next_states': [], 'dones': []}
+        memory = {
+            'states': [],
+            'actions': [],
+            'pre_tanh_actions': [],
+            'logprobs': [],
+            'rewards': [],
+            'next_states': [],
+            'dones': []
+        }
         
         # Prepare seed
         available_scenarios = env.config["num_scenarios"]
@@ -338,7 +349,7 @@ def train(args):
             import gc
             gc.collect()
             
-            env = MAGAILScenarioEnv(config=env_config, agent2policy={})
+            env = BCScenarioEnv(env_config, agent2policy={})
             obs_dict = env.reset(seed=seed)
         
         episode_reward = 0
@@ -349,6 +360,7 @@ def train(args):
             # Select actions for all agents
             actions = {}
             action_logprobs = {}
+            pre_tanh_actions = {}
             
             # obs_dict: {agent_id: obs}
             # MultiAgentScenarioEnv usually returns a dict {agent_id: obs}
@@ -386,9 +398,10 @@ def train(args):
                 obs_dict = new_obs_dict
 
             for agent_id, obs in obs_dict.items():
-                act, logprob = ppo_agent.select_action(obs) # Select action returns numpy
+                act, logprob, pre_tanh = ppo_agent.select_action(obs) # Select action returns numpy
                 actions[agent_id] = act.flatten() # (2,)
                 action_logprobs[agent_id] = logprob # scalar
+                pre_tanh_actions[agent_id] = pre_tanh.flatten()
                 
             # Step Env
             next_obs_dict, rewards, dones, infos = env.step(actions)
@@ -398,6 +411,7 @@ def train(args):
                 if agent_id in actions:
                     memory['states'].append(obs)
                     memory['actions'].append(actions[agent_id])
+                    memory['pre_tanh_actions'].append(pre_tanh_actions[agent_id])
                     memory['logprobs'].append(action_logprobs[agent_id])
                     
                     # Store standard environmental reward for logging (not used for update in GAIL)
@@ -407,7 +421,7 @@ def train(args):
                     # Next state
                     if agent_id in next_obs_dict:
                         memory['next_states'].append(next_obs_dict[agent_id])
-                        memory['dones'].append(False)
+                        memory['dones'].append(dones.get("__all__", False))
                     else:
                         # Agent finished/vanished
                         # We need a dummy next state or handle done correctly
@@ -461,6 +475,10 @@ def train(args):
             disc_loss.backward()
             disc_optimizer.step()
 
+            with torch.no_grad():
+                disc_acc_exp = (exp_preds > 0.5).float().mean().item()
+                disc_acc_pol = (pol_preds < 0.5).float().mean().item()
+            
             # --- 3. Update Policy with GAIL Rewards ---
             # Reward = -log(1 - D(s, a))
             # Or more stable: log(D(s, a)) ? Original GAIL uses -log(1-D) which is log(D) roughly.
@@ -494,11 +512,18 @@ def train(args):
             writer.add_scalar('Loss/Discriminator', disc_loss.item(), i_episode)
             writer.add_scalar('Loss/Policy', ppo_loss, i_episode)
             writer.add_scalar('Reward/Mean_GAIL', np.mean(all_gail_rewards), i_episode)
+            if batch_size > 0:
+                writer.add_scalar('Acc/Disc_Expert', disc_acc_exp, i_episode)
+                writer.add_scalar('Acc/Disc_Policy', disc_acc_pol, i_episode)
+            if len(memory['actions']) > 0:
+                action_arr = np.array(memory['actions'])
+                action_clip_ratio = (np.abs(action_arr) > 0.98).mean()
+                writer.add_scalar('Policy/ActionClipRatio', action_clip_ratio, i_episode)
         
         print(f"Episode {i_episode}: Disc Loss {disc_loss.item():.4f} | PPO Loss {ppo_loss:.4f} | Mean Reward {np.mean(all_gail_rewards):.4f}")
         
         if i_episode % 50 == 0:
-            ppo_agent.save(os.path.join(args.log_dir, f"model_{i_episode}"))
+            ppo_agent.save(os.path.join(args.save_dir, f"model_{i_episode}"))
 
     env.close()
     if writer:
@@ -511,11 +536,13 @@ if __name__ == '__main__':
     parser.add_argument("--batch_size", type=int, default=1024)
     parser.add_argument("--max_episodes", type=int, default=1000)
     parser.add_argument("--num_scenarios", type=int, default=100)
-    parser.add_argument("--log_dir", type=str, default="runs/magail_exp")
+    parser.add_argument("--log_dir", type=str, default="logs/magail", help="TensorBoard log directory")
+    parser.add_argument("--save_dir", type=str, default="models/magail", help="Directory to save model checkpoints")
     
     args = parser.parse_args()
     
-    # Create log dir
+    # Create log dir and save dir
     os.makedirs(args.log_dir, exist_ok=True)
+    os.makedirs(args.save_dir, exist_ok=True)
     
     train(args)