HBBC部署到代码中

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

@@ -1,4 +1,5 @@
 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
@@ -15,19 +16,71 @@ class BCScenarioEnv(MultiAgentScenarioEnv):
     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_")
+                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_"):
+                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()
 
@@ -76,16 +129,21 @@ class BCScenarioEnv(MultiAgentScenarioEnv):
         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()
 

Env/expert_replay_env.py

@@ -293,61 +293,53 @@ 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

@@ -61,5 +61,33 @@ class InverseDynamics:
         # Normalize actions to [-1, 1]
         norm_acc = np.clip(acc / self.max_acc, -1.0, 1.0)
         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