HBBC部署到代码中

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()