import numpy as np from metadrive.component.navigation_module.node_network_navigation import NodeNetworkNavigation from metadrive.envs.scenario_env import ScenarioEnv from metadrive.component.vehicle.vehicle_type import DefaultVehicle, vehicle_class_to_type import math import logging from collections import defaultdict from typing import Union, Dict, AnyStr from metadrive.engine.logger import get_logger, set_log_level from metadrive.type import MetaDriveType class PolicyVehicle(DefaultVehicle): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.policy = None self.destination = None self.expert_vehicle_id = None # 关联专家车辆ID def set_policy(self, policy): self.policy = policy 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) else: return self.action_space.sample() def before_step(self, action): self.last_position = self.position # 2D vector self.last_velocity = self.velocity # 2D vector self.last_speed = self.speed # Scalar self.last_heading_dir = self.heading if action is not None: self.last_current_action.append(action) self._set_action(action) def is_done(self): # arrive or crash pass vehicle_class_to_type[PolicyVehicle] = "default" class MultiAgentScenarioEnv(ScenarioEnv): @classmethod def default_config(cls): config = super().default_config() config.update(dict( 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 def __init__(self, config, agent2policy): self.policy = agent2policy self.controlled_agents = {} 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) # ✅ 关键修复：在每次 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.") # 如果指定了场景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 self.dones = {} 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 _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): 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, 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 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): 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): traffic_light = self.engine.light_manager._lane_index_to_obj[lane.lane.index].status if traffic_light == 'TRAFFIC_LIGHT_GREEN': traffic_light = 1 elif traffic_light == 'TRAFFIC_LIGHT_YELLOW': traffic_light = 2 elif traffic_light == 'TRAFFIC_LIGHT_RED': traffic_light = 3 else: traffic_light = 0 break lidar = self.engine.get_sensor("lidar").perceive(num_lasers=80, distance=30, base_vehicle=vehicle, 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) 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) 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} # ✅ 修复：添加回放模式的完成检查 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()