346 lines
16 KiB
Python
346 lines
16 KiB
Python
import logging
|
|
import numpy as np
|
|
from collections import defaultdict
|
|
from metadrive.component.vehicle.vehicle_type import DefaultVehicle
|
|
from metadrive.type import MetaDriveType
|
|
from Env.scenario_env import MultiAgentScenarioEnv, PolicyVehicle
|
|
from Env.inverse_dynamics import InverseDynamics
|
|
|
|
class ExpertReplayEnv(MultiAgentScenarioEnv):
|
|
def __init__(self, config=None):
|
|
# Allow passing config without agent2policy since we don't use policies for replay
|
|
if config is None:
|
|
config = {}
|
|
# Ensure we don't simulate physics for the controlled agents in the traditional sense
|
|
# but we still need the engine to run
|
|
super().__init__(config, agent2policy={})
|
|
self.inverse_dynamics = InverseDynamics()
|
|
self.expert_tracks = {}
|
|
# Replay SDC/ego ("default_agent" in MetaDrive) as well; otherwise it will keep default action=0 and look stuck.
|
|
self.replay_sdc = self.config.get("replay_sdc", True)
|
|
self.sdc_track = None
|
|
self.sdc_vehicle = None
|
|
self.sdc_agent_id = "default_agent"
|
|
|
|
def reset(self, seed=None):
|
|
self.round = 0
|
|
if self.logger is None:
|
|
from metadrive.engine.logger import get_logger, set_log_level
|
|
self.logger = get_logger()
|
|
log_level = self.config.get("log_level", logging.INFO)
|
|
set_log_level(log_level)
|
|
|
|
self.lazy_init()
|
|
self._reset_global_seed(seed)
|
|
if self.engine is None:
|
|
raise ValueError("Broken MetaDrive instance.")
|
|
|
|
self.background_vehicles = {}
|
|
self.expert_tracks = {}
|
|
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)
|
|
|
|
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)
|
|
|
|
self.engine.reset()
|
|
self.reset_sensors()
|
|
self.engine.taskMgr.step()
|
|
|
|
self.lanes = self.engine.map_manager.current_map.road_network.graph
|
|
|
|
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()
|
|
|
|
# We skip calling super().reset() to avoid double reset
|
|
# But we need to ensure ScenarioEnv-specific setup is done if any.
|
|
# ScenarioEnv.reset() basically does engine.reset() and some cleanup.
|
|
# We covered most of it.
|
|
|
|
self._spawn_controlled_agents()
|
|
self._spawn_all_background_vehicles_at_init()
|
|
|
|
# Ensure SDC/ego is moved to the correct initial expert state.
|
|
if self.replay_sdc:
|
|
self.sdc_vehicle = self.engine.agent_manager.active_agents.get(self.sdc_agent_id, None)
|
|
if self.sdc_vehicle is not None and self.sdc_track is not None:
|
|
valid = self.sdc_track["state"]["valid"]
|
|
t0 = int(np.argmax(valid)) if valid.any() else 0
|
|
pos0 = self.sdc_track["state"]["position"][t0]
|
|
heading0 = self.sdc_track["state"]["heading"][t0]
|
|
vel0 = self.sdc_track["state"]["velocity"][t0]
|
|
self.sdc_vehicle.set_position(pos0)
|
|
self.sdc_vehicle.set_heading_theta(heading0)
|
|
self.sdc_vehicle.set_velocity(vel0)
|
|
|
|
return self._get_all_obs()
|
|
|
|
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():
|
|
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):
|
|
# 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:
|
|
if car['show_time'] == self.round:
|
|
agent_id = f"controlled_{car['id']}"
|
|
|
|
# Check if we already have this agent (shouldn't happen with unique IDs but safety check)
|
|
if agent_id in self.controlled_agents:
|
|
continue
|
|
|
|
# Handling ID flickering / merging
|
|
# If this ID is new, check if there's an existing agent very close to its start position
|
|
# that just disappeared? (Not implemented here, complex logic)
|
|
# But we can check if there's an overlap with existing agents?
|
|
# For now, just spawn.
|
|
|
|
# Read vehicle type/size if available
|
|
vehicle_config = {}
|
|
if 'length' in car and 'width' in car:
|
|
vehicle_config = {
|
|
"length": car['length'],
|
|
"width": car['width']
|
|
}
|
|
|
|
vehicle = self.engine.spawn_object(
|
|
PolicyVehicle,
|
|
name=agent_id,
|
|
vehicle_config=vehicle_config,
|
|
position=car['begin'],
|
|
heading=car['heading']
|
|
)
|
|
vehicle.reset(position=car['begin'], heading=car['heading'])
|
|
|
|
# We don't set policy or destination in the same way, or maybe we do for compatibility
|
|
vehicle.set_destination(car['end'])
|
|
|
|
# Store extra info for replay
|
|
vehicle.expert_track = self.expert_tracks[car['scenario_id']]
|
|
vehicle.original_id = car['id']
|
|
|
|
self.controlled_agents[agent_id] = vehicle
|
|
self.controlled_agent_ids.append(agent_id)
|
|
self.engine.agent_manager.active_agents[agent_id] = vehicle
|
|
|
|
def step(self, action_dict=None):
|
|
# We ignore input action_dict for the purpose of controlling agents
|
|
# Instead, we calculate what the action *should* be
|
|
|
|
self.round += 1
|
|
expert_actions = {}
|
|
|
|
# 1. Update state of all controlled agents to the current timestep (self.round)
|
|
# and compute action from (self.round-1) to (self.round).
|
|
# Wait, usually step() moves T -> T+1.
|
|
# Current state is T. We want to move to T+1.
|
|
# So we need state at T and T+1.
|
|
|
|
# Identify agents that are done (valid=0 at T+1 or T+1 >= length)
|
|
agents_to_remove = []
|
|
|
|
# Update SDC/ego first (otherwise it will stay still with default action=0)
|
|
if self.replay_sdc and self.sdc_vehicle is not None and self.sdc_track is not None:
|
|
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],
|
|
}
|
|
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 = action
|
|
|
|
for agent_id, vehicle in self.controlled_agents.items():
|
|
track = vehicle.expert_track
|
|
# current_step = self.round - 1 # Since we incremented at start
|
|
# But vehicle is currently at state corresponding to self.round - 1.
|
|
# We want to move it to self.round.
|
|
|
|
# Check bounds
|
|
next_step = self.round
|
|
curr_step = self.round - 1
|
|
|
|
if next_step >= len(track['state']['position']):
|
|
agents_to_remove.append(agent_id)
|
|
continue
|
|
|
|
valid = track['state']['valid'][next_step]
|
|
if not valid:
|
|
agents_to_remove.append(agent_id)
|
|
continue
|
|
|
|
# Get states
|
|
curr_pos = track['state']['position'][curr_step]
|
|
next_pos = track['state']['position'][next_step]
|
|
curr_heading = track['state']['heading'][curr_step]
|
|
next_heading = track['state']['heading'][next_step]
|
|
curr_vel = track['state']['velocity'][curr_step]
|
|
next_vel = track['state']['velocity'][next_step]
|
|
|
|
# Prepare state dicts for Inverse Dynamics
|
|
curr_state = {
|
|
'position': curr_pos,
|
|
'heading': curr_heading,
|
|
'velocity': curr_vel
|
|
}
|
|
next_state = {
|
|
'position': next_pos,
|
|
'heading': next_heading,
|
|
'velocity': next_vel
|
|
}
|
|
|
|
# Calculate action
|
|
action, raw_info = self.inverse_dynamics.compute_action(curr_state, next_state, dt=0.1) # Waymo is 10Hz?
|
|
expert_actions[agent_id] = action
|
|
|
|
# Force update vehicle state
|
|
vehicle.set_position(next_pos)
|
|
vehicle.set_heading_theta(next_heading)
|
|
vehicle.set_velocity(next_vel)
|
|
|
|
# Also record this action in the vehicle for later retrieval if needed
|
|
vehicle.last_expert_action = action
|
|
|
|
# Remove finished agents
|
|
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])
|
|
|
|
# Step physics world to update sensors/collision detection
|
|
# We don't need full integration, but we need to update the physics world state
|
|
self.engine.taskMgr.step()
|
|
|
|
# Spawn new agents for this turn
|
|
self._spawn_controlled_agents()
|
|
self._update_background_vehicles()
|
|
|
|
# Get observations
|
|
obs = self._get_all_obs()
|
|
|
|
# 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():
|
|
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
|