完成回放模式与仿真模式,过滤非车道生成车辆,增加对于行人自行车的过滤功能

Env/scenario_env.py

@@ -15,6 +15,7 @@ class PolicyVehicle(DefaultVehicle):
         super().__init__(*args, **kwargs)
         self.policy = None
         self.destination = None
+        self.expert_vehicle_id = None  # 关联专家车辆ID
 
     def set_policy(self, policy):
         self.policy = policy
@@ -22,6 +23,9 @@ class PolicyVehicle(DefaultVehicle):
     def set_destination(self, des):
         self.destination = des
 
+    def set_expert_vehicle_id(self, vid):
+        self.expert_vehicle_id = vid
+
     def act(self, observation, policy=None):
         if self.policy is not None:
             return self.policy.act(observation)
@@ -53,6 +57,15 @@ class MultiAgentScenarioEnv(ScenarioEnv):
             data_directory=None,
             num_controlled_agents=3,
             horizon=1000,
+            filter_offroad_vehicles=True,  # 车道过滤开关
+            lane_tolerance=3.0,  # 车道检测容差（米）
+            replay_mode=False,  # 回放模式开关
+            specific_scenario_id=None,  # 新增：指定场景ID（仅回放模式）
+            use_scenario_duration=False,  # 新增：使用场景原始时长作为horizon
+            # 对象类型过滤选项
+            spawn_vehicles=True,  # 是否生成车辆
+            spawn_pedestrians=True,  # 是否生成行人
+            spawn_cyclists=True,  # 是否生成自行车
         ))
         return config
 
@@ -62,50 +75,180 @@ class MultiAgentScenarioEnv(ScenarioEnv):
         self.controlled_agent_ids = []
         self.obs_list = []
         self.round = 0
+        self.expert_trajectories = {}  # 存储完整专家轨迹
+        self.replay_mode = config.get("replay_mode", False)
+        self.scenario_max_duration = 0  # 场景实际最大时长
         super().__init__(config)
 
     def reset(self, seed: Union[None, int] = None):
         self.round = 0
+        
         if self.logger is None:
             self.logger = get_logger()
-            log_level = self.config.get("log_level", logging.DEBUG if self.config.get("debug", False) else logging.INFO)
-            set_log_level(log_level)
+        log_level = self.config.get("log_level", logging.DEBUG if self.config.get("debug", False) else logging.INFO)
+        set_log_level(log_level)
 
+        # ✅ 关键修复：在每次 reset 前清理所有自定义生成的对象
+        if hasattr(self, 'engine') and self.engine is not None:
+            if hasattr(self, 'controlled_agents') and self.controlled_agents:
+                # 先从 agent_manager 中移除
+                if hasattr(self.engine, 'agent_manager'):
+                    for agent_id in list(self.controlled_agents.keys()):
+                        if agent_id in self.engine.agent_manager.active_agents:
+                            self.engine.agent_manager.active_agents.pop(agent_id)
+                
+                # 然后清理对象
+                for agent_id, vehicle in list(self.controlled_agents.items()):
+                    try:
+                        self.engine.clear_objects([vehicle.id])
+                    except:
+                        pass
+                
+                self.controlled_agents.clear()
+                self.controlled_agent_ids.clear()
+        
         self.lazy_init()
         self._reset_global_seed(seed)
+
         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:
-            self.engine.traffic_manager.current_traffic_data.pop(scenario_id)
+        # 如果指定了场景ID，修改start_scenario_index
+        if self.config.get("specific_scenario_id") is not None:
+            scenario_id = self.config.get("specific_scenario_id")
+            self.config["start_scenario_index"] = scenario_id
+            if self.config.get("debug", False):
+                self.logger.info(f"Using specific scenario ID: {scenario_id}")
 
+        # ✅ 先初始化引擎和 lanes
         self.engine.reset()
         self.reset_sensors()
         self.engine.taskMgr.step()
-
         self.lanes = self.engine.map_manager.current_map.road_network.graph
 
+        # 记录专家数据（现在 self.lanes 已经初始化）
+        _obj_to_clean_this_frame = []
+        self.car_birth_info_list = []
+        self.expert_trajectories.clear()
+        total_vehicles = 0
+        total_pedestrians = 0
+        total_cyclists = 0
+        filtered_vehicles = 0
+        filtered_by_type = 0
+        self.scenario_max_duration = 0  # 重置场景时长
+
+        for scenario_id, track in self.engine.traffic_manager.current_traffic_data.items():
+            if scenario_id == self.engine.traffic_manager.sdc_scenario_id:
+                continue
+
+            # 对象类型过滤
+            obj_type = track["type"]
+            
+            # 统计对象类型
+            if obj_type == MetaDriveType.VEHICLE:
+                total_vehicles += 1
+            elif obj_type == MetaDriveType.PEDESTRIAN:
+                total_pedestrians += 1
+            elif obj_type == MetaDriveType.CYCLIST:
+                total_cyclists += 1
+            
+            # 根据配置过滤对象类型
+            if obj_type == MetaDriveType.VEHICLE and not self.config.get("spawn_vehicles", True):
+                _obj_to_clean_this_frame.append(scenario_id)
+                filtered_by_type += 1
+                if self.config.get("debug", False):
+                    self.logger.debug(f"Filtering VEHICLE {track['metadata']['object_id']} - spawn_vehicles=False")
+                continue
+            
+            if obj_type == MetaDriveType.PEDESTRIAN and not self.config.get("spawn_pedestrians", True):
+                _obj_to_clean_this_frame.append(scenario_id)
+                filtered_by_type += 1
+                if self.config.get("debug", False):
+                    self.logger.debug(f"Filtering PEDESTRIAN {track['metadata']['object_id']} - spawn_pedestrians=False")
+                continue
+            
+            if obj_type == MetaDriveType.CYCLIST and not self.config.get("spawn_cyclists", True):
+                _obj_to_clean_this_frame.append(scenario_id)
+                filtered_by_type += 1
+                if self.config.get("debug", False):
+                    self.logger.debug(f"Filtering CYCLIST {track['metadata']['object_id']} - spawn_cyclists=False")
+                continue
+
+            # 只处理车辆类型（行人和自行车暂时只做过滤）
+            if track["type"] == MetaDriveType.VEHICLE:
+                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
+
+                if first_show == -1 or last_show == -1:
+                    continue
+
+                # 更新场景最大时长
+                self.scenario_max_duration = max(self.scenario_max_duration, last_show + 1)
+
+                # 获取车辆初始位置
+                initial_position = (
+                    track['state']['position'][first_show, 0],
+                    track['state']['position'][first_show, 1]
+                )
+
+                # 车道过滤
+                if self.config.get("filter_offroad_vehicles", True):
+                    if not self._is_position_on_lane(initial_position):
+                        filtered_vehicles += 1
+                        _obj_to_clean_this_frame.append(scenario_id)
+                        if self.config.get("debug", False):
+                            self.logger.debug(
+                                f"Filtering vehicle {track['metadata']['object_id']} - "
+                                f"not on lane at position {initial_position}"
+                            )
+                        continue
+
+                # 存储完整专家轨迹（只使用2D位置，避免高度问题）
+                object_id = track['metadata']['object_id']
+                positions_2d = track['state']['position'].copy()
+                positions_2d[:, 2] = 0  # 将z坐标设为0，让MetaDrive自动处理高度
+                
+                self.expert_trajectories[object_id] = {
+                    'positions': positions_2d,
+                    'headings': track['state']['heading'].copy(),
+                    'velocities': track['state']['velocity'].copy(),
+                    'valid': track['state']['valid'].copy(),
+                }
+
+                # 保存车辆生成信息
+                self.car_birth_info_list.append({
+                    'id': object_id,
+                    'show_time': first_show,
+                    'begin': initial_position,
+                    'heading': track['state']['heading'][first_show],
+                    'velocity': track['state']['velocity'][first_show] if self.config.get("inherit_expert_velocity", False) else None,
+                    'end': (track['state']['position'][last_show, 0], track['state']['position'][last_show, 1])
+                })
+
+                # 在回放和仿真模式下都清除原始专家车辆
+                _obj_to_clean_this_frame.append(scenario_id)
+
+        # 清除专家车辆和过滤的对象
+        for scenario_id in _obj_to_clean_this_frame:
+            self.engine.traffic_manager.current_traffic_data.pop(scenario_id)
+
+        # 输出统计信息
+        if self.config.get("debug", False):
+            self.logger.info(f"=== 对象统计 ===")
+            self.logger.info(f"车辆 (VEHICLE): 总数={total_vehicles}, 车道过滤={filtered_vehicles}, 保留={total_vehicles - filtered_vehicles}")
+            self.logger.info(f"行人 (PEDESTRIAN): 总数={total_pedestrians}")
+            self.logger.info(f"自行车 (CYCLIST): 总数={total_cyclists}")
+            self.logger.info(f"类型过滤: {filtered_by_type} 个对象")
+            self.logger.info(f"场景时长: {self.scenario_max_duration} 步")
+
+        # 如果启用场景时长控制，更新horizon
+        if self.config.get("use_scenario_duration", False) and self.scenario_max_duration > 0:
+            original_horizon = self.config["horizon"]
+            self.config["horizon"] = self.scenario_max_duration
+            if self.config.get("debug", False):
+                self.logger.info(f"Horizon updated from {original_horizon} to {self.scenario_max_duration} (scenario duration)")
+
         if self.top_down_renderer is not None:
             self.top_down_renderer.clear()
             self.engine.top_down_renderer = None
@@ -113,7 +256,6 @@ class MultiAgentScenarioEnv(ScenarioEnv):
         self.dones = {}
         self.episode_rewards = defaultdict(float)
         self.episode_lengths = defaultdict(int)
-
         self.controlled_agents.clear()
         self.controlled_agent_ids.clear()
 
@@ -122,37 +264,92 @@ class MultiAgentScenarioEnv(ScenarioEnv):
 
         return self._get_all_obs()
 
+    def _is_position_on_lane(self, position, tolerance=None):
+        if tolerance is None:
+            tolerance = self.config.get("lane_tolerance", 3.0)
+
+        # 确保 self.lanes 已初始化
+        if not hasattr(self, 'lanes') or self.lanes is None:
+            if self.config.get("debug", False):
+                self.logger.warning("Lanes not initialized, skipping lane check")
+            return True
+
+        position_2d = np.array(position[:2]) if len(position) > 2 else np.array(position)
+
+        try:
+            for lane in self.lanes.values():
+                if lane.lane.point_on_lane(position_2d):
+                    return True
+
+                lane_start = np.array(lane.lane.start)[:2]
+                lane_end = np.array(lane.lane.end)[:2]
+                lane_vec = lane_end - lane_start
+                lane_length = np.linalg.norm(lane_vec)
+
+                if lane_length < 1e-6:
+                    continue
+
+                lane_vec_normalized = lane_vec / lane_length
+                point_vec = position_2d - lane_start
+                projection = np.dot(point_vec, lane_vec_normalized)
+
+                if 0 <= projection <= lane_length:
+                    closest_point = lane_start + projection * lane_vec_normalized
+                    distance = np.linalg.norm(position_2d - closest_point)
+                    if distance <= tolerance:
+                        return True
+        except Exception as e:
+            if self.config.get("debug", False):
+                self.logger.warning(f"Lane check error: {e}")
+            return False
+
+        return False
+
     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])
         for car in self.car_birth_info_list:
             if car['show_time'] == self.round:
                 agent_id = f"controlled_{car['id']}"
-
+                vehicle_config = {}
                 vehicle = self.engine.spawn_object(
                     PolicyVehicle,
-                    vehicle_config={},
+                    vehicle_config=vehicle_config,
                     position=car['begin'],
                     heading=car['heading']
                 )
-                vehicle.reset(position=car['begin'], heading=car['heading'])
 
+                # 重置车辆状态
+                reset_kwargs = {
+                    'position': car['begin'],
+                    'heading': car['heading']
+                }
+
+                # 如果启用速度继承，设置初始速度
+                if car.get('velocity') is not None:
+                    reset_kwargs['velocity'] = car['velocity']
+
+                vehicle.reset(**reset_kwargs)
+
+                # 设置策略和目的地
                 vehicle.set_policy(self.policy)
                 vehicle.set_destination(car['end'])
+                vehicle.set_expert_vehicle_id(car['id'])
 
                 self.controlled_agents[agent_id] = vehicle
                 self.controlled_agent_ids.append(agent_id)
 
-                # ✅ 关键：注册到引擎的 active_agents，才能参与物理更新
+                # 注册到引擎的 active_agents
                 self.engine.agent_manager.active_agents[agent_id] = vehicle
 
+                if self.config.get("debug", False):
+                    self.logger.debug(f"Spawned vehicle {agent_id} at round {self.round}, position {car['begin']}")
+
     def _get_all_obs(self):
-        # position, velocity, heading, lidar, navigation, TODO: trafficlight -> list
         self.obs_list = []
+
         for agent_id, vehicle in self.controlled_agents.items():
             state = vehicle.get_state()
-
             traffic_light = 0
+
             for lane in self.lanes.values():
                 if lane.lane.point_on_lane(state['position'][:2]):
                     if self.engine.light_manager.has_traffic_light(lane.lane.index):
@@ -168,37 +365,69 @@ class MultiAgentScenarioEnv(ScenarioEnv):
                     break
 
             lidar = self.engine.get_sensor("lidar").perceive(num_lasers=80, distance=30, base_vehicle=vehicle,
-                                                             physics_world=self.engine.physics_world.dynamic_world)
+                                                              physics_world=self.engine.physics_world.dynamic_world)
             side_lidar = self.engine.get_sensor("side_detector").perceive(num_lasers=10, distance=8,
-                                                                          base_vehicle=vehicle,
-                                                                          physics_world=self.engine.physics_world.static_world)
+                                                                           base_vehicle=vehicle,
+                                                                           physics_world=self.engine.physics_world.static_world)
             lane_line_lidar = self.engine.get_sensor("lane_line_detector").perceive(num_lasers=10, distance=3,
-                                                                                    base_vehicle=vehicle,
-                                                                                    physics_world=self.engine.physics_world.static_world)
+                                                                                      base_vehicle=vehicle,
+                                                                                      physics_world=self.engine.physics_world.static_world)
 
-            obs = (state['position'][:2] + list(state['velocity']) + [state['heading_theta']]
+            obs = (list(state['position'][:2]) + list(state['velocity']) + [state['heading_theta']]
                    + lidar[0] + side_lidar[0] + lane_line_lidar[0] + [traffic_light]
                    + list(vehicle.destination))
+
             self.obs_list.append(obs)
+
         return self.obs_list
 
     def step(self, action_dict: Dict[AnyStr, Union[list, np.ndarray]]):
         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.engine.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()
+
+        # 获取观测
         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.episode_step >= self.config["horizon"]
+        
+        # ✅ 修复：添加回放模式的完成检查
+        replay_finished = False
+        if self.replay_mode and self.config.get("use_scenario_duration", False):
+            # 检查是否所有专家轨迹都已播放完毕
+            if self.round >= self.scenario_max_duration:
+                replay_finished = True
+                if self.config.get("debug", False):
+                    self.logger.info(f"Replay finished at step {self.round}/{self.scenario_max_duration}")
+        
+        dones["__all__"] = self.episode_step >= self.config["horizon"] or replay_finished
+
         infos = {aid: {} for aid in self.controlled_agents}
+
         return obs, rewards, dones, infos
+
+    def close(self):
+        # ✅ 清理所有生成的车辆
+        if hasattr(self, 'controlled_agents') and self.controlled_agents:
+            for agent_id, vehicle in list(self.controlled_agents.items()):
+                if vehicle in self.engine.get_objects():
+                    self.engine.clear_objects([vehicle.id])
+            self.controlled_agents.clear()
+            self.controlled_agent_ids.clear()
+        
+        super().close()