Compare commits
2 Commits
95cc78d940
...
BC_dev
| Author | SHA1 | Date | |
|---|---|---|---|
| 0f9f080e77 | |||
| ceb6648a31 |
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -51,10 +51,8 @@ class BCScenarioEnv(MultiAgentScenarioEnv):
|
||||
return car_birth_info_list, background_vehicles, obj_to_clean
|
||||
|
||||
def _spawn_background_vehicles(self):
|
||||
"""Spawn static background vehicles so they appear in active_agents and thus in obs (same as ExpertReplayEnv)."""
|
||||
"""Spawn all static background vehicles once at reset (no show_time filter; same as ExpertReplayEnv)."""
|
||||
for sid, car in self.background_vehicles.items():
|
||||
if car["show_time"] != self.round:
|
||||
continue
|
||||
bg_id = f"bg_{car['id']}"
|
||||
if bg_id in self.engine.agent_manager.active_agents:
|
||||
continue
|
||||
@@ -70,24 +68,12 @@ class BCScenarioEnv(MultiAgentScenarioEnv):
|
||||
)
|
||||
v.set_velocity([0, 0])
|
||||
self.engine.agent_manager.active_agents[bg_id] = v
|
||||
v.valid_mask = car["valid"]
|
||||
v.start_t = car["show_time"]
|
||||
v.valid_mask = car.get("valid")
|
||||
v.start_t = car.get("show_time")
|
||||
|
||||
def _update_background_vehicles(self):
|
||||
self._spawn_background_vehicles()
|
||||
to_remove = []
|
||||
objects_to_clear = []
|
||||
for aid, v in self.engine.agent_manager.active_agents.items():
|
||||
if not aid.startswith("bg_"):
|
||||
continue
|
||||
if hasattr(v, "valid_mask"):
|
||||
if self.round >= len(v.valid_mask) or not v.valid_mask[self.round]:
|
||||
to_remove.append(aid)
|
||||
objects_to_clear.append(v)
|
||||
for aid in to_remove:
|
||||
self.engine.agent_manager.active_agents.pop(aid, None)
|
||||
if objects_to_clear:
|
||||
self.engine.clear_objects([v.id for v in objects_to_clear])
|
||||
# Static vehicles are spawned once at init and never removed.
|
||||
pass
|
||||
|
||||
def step(self, action_dict):
|
||||
self.round += 1
|
||||
|
||||
@@ -97,7 +97,7 @@ class ExpertReplayEnv(MultiAgentScenarioEnv):
|
||||
# We covered most of it.
|
||||
|
||||
self._spawn_controlled_agents()
|
||||
self._spawn_background_vehicles() # Initial spawn for background
|
||||
self._spawn_all_background_vehicles_at_init()
|
||||
|
||||
# Ensure SDC/ego is moved to the correct initial expert state.
|
||||
if self.replay_sdc:
|
||||
@@ -114,35 +114,18 @@ class ExpertReplayEnv(MultiAgentScenarioEnv):
|
||||
|
||||
return self._get_all_obs()
|
||||
|
||||
def _spawn_background_vehicles(self):
|
||||
# Spawn static/background vehicles
|
||||
# Since they are static, we might just spawn them once if their show_time is 0
|
||||
# But Waymo tracks have valid bits, they might appear/disappear.
|
||||
# For optimization, if they are truly static (never move), we just spawn them when show_time matches.
|
||||
|
||||
# We need to track spawned background vehicles to remove them if they become invalid?
|
||||
# Since we defined them as "static", they probably stay put.
|
||||
# But validity might change (e.g. late spawn).
|
||||
|
||||
# For simplicity in this step, let's just iterate and spawn if time matches
|
||||
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():
|
||||
if car['show_time'] == self.round:
|
||||
# Spawn as a Traffic Vehicle (not PolicyVehicle), or just a static object?
|
||||
# Using DefaultVehicle is fine, but don't add to controlled_agents
|
||||
|
||||
# Check duplication
|
||||
bg_id = f"bg_{car['id']}"
|
||||
# if bg_id in self.engine.obj_to_id: # obj_to_id might not be available in all versions
|
||||
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,
|
||||
@@ -150,51 +133,14 @@ class ExpertReplayEnv(MultiAgentScenarioEnv):
|
||||
position=car['begin'],
|
||||
heading=car['heading']
|
||||
)
|
||||
|
||||
# Set color to grey/dark to indicate background
|
||||
v.set_velocity([0, 0])
|
||||
# Maybe set color? MetaDrive vehicles random color.
|
||||
# v.set_color(...) if supported
|
||||
|
||||
# Register as an active object but NOT controlled agent
|
||||
# The engine manages it.
|
||||
# CRITICAL: We need it in self.engine.agent_manager.active_agents for Observation?
|
||||
# If we want it to be seen by Lidar/Observation, it needs to be an "agent" or "traffic".
|
||||
# DefaultVehicle spawned this way is just an object.
|
||||
# We should add it to traffic manager? Or just leave it as object?
|
||||
# MultiAgentScenarioEnv._get_all_obs iterates self.engine.agent_manager.active_agents
|
||||
|
||||
# If we want it in observation, we must add it to active_agents OR iterate over all objects.
|
||||
# Adding to active_agents is easier for compatibility.
|
||||
self.engine.agent_manager.active_agents[bg_id] = v
|
||||
|
||||
# Store valid mask to remove it later if needed?
|
||||
v.valid_mask = car['valid']
|
||||
v.valid_mask = car.get('valid')
|
||||
v.start_t = car['show_time']
|
||||
|
||||
def _update_background_vehicles(self):
|
||||
# Remove background vehicles if they become invalid
|
||||
# Or spawn new ones
|
||||
self._spawn_background_vehicles()
|
||||
|
||||
# Check validity for existing
|
||||
to_remove = []
|
||||
objects_to_clear = []
|
||||
|
||||
for aid, v in self.engine.agent_manager.active_agents.items():
|
||||
if aid.startswith("bg_"):
|
||||
# Check validity
|
||||
if hasattr(v, 'valid_mask'):
|
||||
curr_step = self.round
|
||||
if curr_step >= len(v.valid_mask) or not v.valid_mask[curr_step]:
|
||||
to_remove.append(aid)
|
||||
objects_to_clear.append(v)
|
||||
|
||||
for aid in to_remove:
|
||||
self.engine.agent_manager.active_agents.pop(aid, None)
|
||||
|
||||
if objects_to_clear:
|
||||
self.engine.clear_objects([v.id for v in objects_to_clear])
|
||||
# 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:
|
||||
|
||||
@@ -64,6 +64,11 @@ class MultiAgentScenarioEnv(ScenarioEnv):
|
||||
self.round = 0
|
||||
super().__init__(config)
|
||||
|
||||
@property
|
||||
def num_controlled_in_scenario(self) -> int:
|
||||
"""整个场景中受控车轨迹总数(car_birth_info_list 长度),会在不同 show_time 陆续 spawn。"""
|
||||
return len(getattr(self, "car_birth_info_list", []))
|
||||
|
||||
def reset(self, seed: Union[None, int] = None):
|
||||
self.round = 0
|
||||
if self.logger is None:
|
||||
@@ -76,6 +81,11 @@ class MultiAgentScenarioEnv(ScenarioEnv):
|
||||
if self.engine is None:
|
||||
raise ValueError("Broken MetaDrive instance.")
|
||||
|
||||
# 注意:_build_birth_lists_from_traffic() 在 engine.reset() 之前执行,读的是当前 engine 的
|
||||
# current_traffic_data 与 map_manager.current_map。若复用同一 env 连续 reset(0)、reset(1),
|
||||
# MetaDrive 可能已按 seed 更新了 traffic 为 scenario 1,但 map 仍为 scenario 0(在 engine.reset() 才切图),
|
||||
# 导致 is_on_lane( scenario_1 车位, scenario_0 地图 ) 全为 False → 全部 off_lane → 0 受控车。
|
||||
# 因此多场景时应“每个 scenario 单独建 env”(start_scenario_index=i, num_scenarios=1)再 reset(seed=i)。
|
||||
self.background_vehicles = getattr(self, "background_vehicles", {})
|
||||
self.car_birth_info_list, self.background_vehicles, _obj_to_clean = self._build_birth_lists_from_traffic()
|
||||
for scenario_id in _obj_to_clean:
|
||||
|
||||
84
Env/utils.py
84
Env/utils.py
@@ -5,6 +5,86 @@ import random
|
||||
from metadrive.type import MetaDriveType
|
||||
|
||||
|
||||
def _static_obbs_overlap(car_a, car_b):
|
||||
"""
|
||||
Check if two static vehicle OBBs overlap (2D SAT).
|
||||
car_a, car_b: dicts with "begin" (x, y), "heading" (rad), "length", "width".
|
||||
begin is center; half-extents are length/2, width/2.
|
||||
"""
|
||||
def _get_corners(car):
|
||||
cx, cy = car["begin"][0], car["begin"][1]
|
||||
h = float(car["heading"])
|
||||
L2 = float(car["length"]) / 2.0
|
||||
W2 = float(car["width"]) / 2.0
|
||||
ux, uy = np.cos(h), np.sin(h)
|
||||
vx, vy = -np.sin(h), np.cos(h)
|
||||
return np.array([
|
||||
[cx + L2 * ux + W2 * vx, cy + L2 * uy + W2 * vy],
|
||||
[cx + L2 * ux - W2 * vx, cy + L2 * uy - W2 * vy],
|
||||
[cx - L2 * ux - W2 * vx, cy - L2 * uy - W2 * vy],
|
||||
[cx - L2 * ux + W2 * vx, cy - L2 * uy + W2 * vy],
|
||||
])
|
||||
|
||||
def _get_axes(car):
|
||||
h = float(car["heading"])
|
||||
return [
|
||||
np.array([np.cos(h), np.sin(h)]),
|
||||
np.array([-np.sin(h), np.cos(h)]),
|
||||
]
|
||||
|
||||
corners_a = _get_corners(car_a)
|
||||
corners_b = _get_corners(car_b)
|
||||
axes = _get_axes(car_a) + _get_axes(car_b)
|
||||
|
||||
for axis in axes:
|
||||
proj_a = corners_a @ axis
|
||||
proj_b = corners_b @ axis
|
||||
min_a, max_a = proj_a.min(), proj_a.max()
|
||||
min_b, max_b = proj_b.min(), proj_b.max()
|
||||
if max_a < min_b or max_b < min_a:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _deduplicate_background_by_collision(background_vehicles):
|
||||
"""
|
||||
Merge static tracks that collide (same physical vehicle). Build collision graph,
|
||||
find connected components, keep one representative per component (min show_time, then scenario_id).
|
||||
"""
|
||||
if not background_vehicles:
|
||||
return background_vehicles
|
||||
items = list(background_vehicles.items())
|
||||
n = len(items)
|
||||
# Build adjacency by index
|
||||
parent = list(range(n))
|
||||
|
||||
def find(i):
|
||||
if parent[i] != i:
|
||||
parent[i] = find(parent[i])
|
||||
return parent[i]
|
||||
|
||||
def union(i, j):
|
||||
pi, pj = find(i), find(j)
|
||||
if pi != pj:
|
||||
parent[pi] = pj
|
||||
|
||||
for i in range(n):
|
||||
for j in range(i + 1, n):
|
||||
if _static_obbs_overlap(items[i][1], items[j][1]):
|
||||
union(i, j)
|
||||
|
||||
# Representative per component: index with min (show_time, scenario_id)
|
||||
comp_rep = {}
|
||||
for i in range(n):
|
||||
r = find(i)
|
||||
sid, car = items[i][0], items[i][1]
|
||||
key = (car.get("show_time", 0), sid)
|
||||
if r not in comp_rep or key < comp_rep[r][0]:
|
||||
comp_rep[r] = (key, sid, car)
|
||||
|
||||
return {sid: car for (_, sid, car) in comp_rep.values()}
|
||||
|
||||
|
||||
def is_on_lane(pos, map_manager, threshold=2.0):
|
||||
"""Check if a position is on a valid lane (within lateral tolerance)."""
|
||||
if map_manager is None or map_manager.current_map is None:
|
||||
@@ -31,6 +111,7 @@ def filter_traffic_tracks_to_birth_lists(
|
||||
static_displacement_threshold=5.0,
|
||||
static_speed_threshold=1.0,
|
||||
return_stats=False,
|
||||
deduplicate_static_by_collision=True,
|
||||
):
|
||||
"""
|
||||
Filter traffic tracks into controlled (car_birth_info_list) and background lists.
|
||||
@@ -127,6 +208,9 @@ def filter_traffic_tracks_to_birth_lists(
|
||||
"width": track["state"]["width"][first_show],
|
||||
})
|
||||
|
||||
if deduplicate_static_by_collision and background_vehicles:
|
||||
background_vehicles = _deduplicate_background_by_collision(background_vehicles)
|
||||
|
||||
if return_stats:
|
||||
stats = {
|
||||
"n_total": n_total,
|
||||
|
||||
@@ -49,6 +49,35 @@ def load_expert_pkl(expert_data_path):
|
||||
return obs_data, act_data
|
||||
|
||||
|
||||
def get_expert_scenario_ids(expert_data_path, max_ids=10):
|
||||
"""
|
||||
从专家 pkl 中收集出现过的 scenario_id(这些场景在采集时曾有受控车)。
|
||||
用于 eval 时只在这些场景上评估,保证 eval 有受控车。
|
||||
返回排序后的 list,最多 max_ids 个。
|
||||
"""
|
||||
if os.path.isdir(expert_data_path):
|
||||
pkl_files = glob.glob(os.path.join(expert_data_path, "*.pkl"))
|
||||
elif os.path.exists(expert_data_path):
|
||||
pkl_files = [expert_data_path]
|
||||
else:
|
||||
return []
|
||||
|
||||
seen = set()
|
||||
for pkl_file in pkl_files:
|
||||
try:
|
||||
with open(pkl_file, "rb") as f:
|
||||
data = pickle.load(f)
|
||||
if isinstance(data, list):
|
||||
for traj in data:
|
||||
if "scenario_id" in traj:
|
||||
seen.add(traj["scenario_id"])
|
||||
# dict 格式通常没有 per-trajectory scenario_id,跳过
|
||||
except Exception:
|
||||
continue
|
||||
out = sorted(seen)[:max_ids]
|
||||
return out
|
||||
|
||||
|
||||
class MAGAILExpertDataset(Dataset):
|
||||
def __init__(self, data_dir, transform=None):
|
||||
"""
|
||||
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -30,6 +30,13 @@ def _run_replay(args):
|
||||
max_available = len(summary_lookup) - args.start_index
|
||||
num_to_run = min(args.num_scenarios, max_available)
|
||||
|
||||
print(f"Initializing ExpertReplayEnv with data from {data_path}...")
|
||||
|
||||
try:
|
||||
for i in range(args.start_index, args.start_index + num_to_run):
|
||||
print(f"\n--- Playing Scenario {i} ---")
|
||||
# Each scenario uses a fresh env (start_scenario_index=i, num_scenarios=1) so the second
|
||||
# scenario and beyond are fully cleaned and loaded like a single-scenario run.
|
||||
env_config = {
|
||||
"data_directory": data_path,
|
||||
"is_multi_agent": True,
|
||||
@@ -38,24 +45,19 @@ def _run_replay(args):
|
||||
"use_render": True,
|
||||
"sequential_seed": True,
|
||||
"reactive_traffic": False,
|
||||
"start_scenario_index": args.start_index,
|
||||
"num_scenarios": -1,
|
||||
"start_scenario_index": i,
|
||||
"num_scenarios": 1,
|
||||
"log_level": 40,
|
||||
}
|
||||
|
||||
print(f"Initializing ExpertReplayEnv with data from {data_path}...")
|
||||
env = ExpertReplayEnv(config=env_config)
|
||||
|
||||
try:
|
||||
for i in range(args.start_index, args.start_index + num_to_run):
|
||||
print(f"\n--- Playing Scenario {i} ---")
|
||||
try:
|
||||
obs = env.reset(seed=i)
|
||||
except Exception as e:
|
||||
print(f"Error resetting scenario {i}: {e}")
|
||||
env.close()
|
||||
continue
|
||||
|
||||
print(f"Scenario loaded. Controlled agents: {len(env.controlled_agents)}")
|
||||
print(f"Scenario loaded. Controlled agents (current): {len(env.controlled_agents)}, total in scenario: {env.num_controlled_in_scenario}")
|
||||
|
||||
for step in range(args.horizon):
|
||||
obs, rewards, dones, infos = env.step(None)
|
||||
@@ -67,6 +69,7 @@ def _run_replay(args):
|
||||
if dones["__all__"]:
|
||||
print(f"Scenario {i} finished at step {step}")
|
||||
break
|
||||
env.close()
|
||||
except KeyboardInterrupt:
|
||||
print("Interrupted by user")
|
||||
except Exception as e:
|
||||
@@ -74,7 +77,6 @@ def _run_replay(args):
|
||||
traceback.print_exc()
|
||||
print(f"Global error: {e}")
|
||||
finally:
|
||||
env.close()
|
||||
print("Environment closed.")
|
||||
|
||||
|
||||
@@ -176,7 +178,7 @@ def _run_policy(args):
|
||||
pass
|
||||
continue
|
||||
|
||||
print(f"Scenario loaded. Controlled agents: {len(obs_dict)}")
|
||||
print(f"Scenario loaded. Controlled agents (current): {len(obs_dict)}, total in scenario: {env.num_controlled_in_scenario}")
|
||||
if len(obs_dict) == 0:
|
||||
print(f"Scenario {i} has no controlled agents (all filtered out). Skipping.")
|
||||
continue
|
||||
|
||||
66
train_bc.py
66
train_bc.py
@@ -15,11 +15,12 @@ from torch.utils.tensorboard import SummaryWriter
|
||||
from Algorithm.policy import StateIndependentPolicy
|
||||
from Algorithm.bc import train_bc_epoch, eval_bc_epoch
|
||||
from Env.bc_env import BCScenarioEnv
|
||||
from dataset.loader import load_expert_pkl
|
||||
from dataset.loader import load_expert_pkl, get_expert_scenario_ids
|
||||
|
||||
|
||||
def evaluate_policy(policy, args, device):
|
||||
"""在 BCScenarioEnv 中评估策略,跑若干 episode,返回平均 reward。"""
|
||||
"""在 BCScenarioEnv 中评估策略:仅使用专家数据中出现过的 scenario_id,保证 eval 有受控车。
|
||||
输出与 replay 对齐:agents (current)=reset 时受控车数,total in scenario=该场景受控轨迹总数(car_birth_info_list 长度)。"""
|
||||
waymo_data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
|
||||
data_dir = os.path.join(waymo_data_dir, "exp_filtered")
|
||||
if not os.path.exists(data_dir):
|
||||
@@ -28,28 +29,47 @@ def evaluate_policy(policy, args, device):
|
||||
print(f"[ERROR] Could not find scenario data in {waymo_data_dir}. Evaluation skipped.")
|
||||
return 0.0
|
||||
|
||||
scenario_ids = get_expert_scenario_ids(args.expert_data_path, max_ids=5)
|
||||
if not scenario_ids:
|
||||
print("[WARN] No scenario_id in expert pkl, falling back to scenarios [0,1,2]. Eval may have 0 controlled agents.")
|
||||
scenario_ids = [0, 1, 2]
|
||||
|
||||
total_rewards = []
|
||||
horizon = 200
|
||||
|
||||
for idx, scenario_id in enumerate(scenario_ids):
|
||||
env_config = {
|
||||
"data_directory": data_dir,
|
||||
"is_multi_agent": True,
|
||||
"num_controlled_agents": 3,
|
||||
"num_controlled_agents": 100,
|
||||
"use_render": False,
|
||||
"sequential_seed": True,
|
||||
"horizon": 200,
|
||||
"horizon": horizon,
|
||||
"start_scenario_index": scenario_id,
|
||||
"num_scenarios": 1,
|
||||
}
|
||||
env = BCScenarioEnv(env_config, agent2policy=None)
|
||||
total_rewards = []
|
||||
|
||||
try:
|
||||
for i in range(3):
|
||||
obs_dict = env.reset(seed=i)
|
||||
episode_reward = 0
|
||||
dones = {"__all__": False}
|
||||
obs_dict = env.reset(seed=scenario_id)
|
||||
except Exception as e:
|
||||
print(f" Eval Episode {idx} (scenario {scenario_id}): reset failed: {e}")
|
||||
env.close()
|
||||
continue
|
||||
|
||||
n_controlled = len(env.controlled_agents)
|
||||
n_total_in_scenario = getattr(env, "num_controlled_in_scenario", n_controlled)
|
||||
if n_controlled == 0:
|
||||
print(
|
||||
f" Eval Episode {idx} (scenario {scenario_id}): 0 controlled agents (total in scenario: {n_total_in_scenario}), skip."
|
||||
)
|
||||
env.close()
|
||||
continue
|
||||
|
||||
episode_reward = 0.0
|
||||
step_count = 0
|
||||
horizon = 200
|
||||
while not dones["__all__"]:
|
||||
dones = {"__all__": False}
|
||||
while not dones["__all__"] and step_count < horizon:
|
||||
step_count += 1
|
||||
if step_count >= horizon:
|
||||
break
|
||||
if not obs_dict:
|
||||
obs_dict, _, dones, _ = env.step({})
|
||||
continue
|
||||
@@ -62,18 +82,20 @@ def evaluate_policy(policy, args, device):
|
||||
action_dict = {aid: act for aid, act in zip(agent_ids, actions)}
|
||||
obs_dict, rewards, dones, _ = env.step(action_dict)
|
||||
episode_reward += sum(rewards.values())
|
||||
|
||||
total_rewards.append(episode_reward)
|
||||
print(f" Eval Episode {i}: Total Reward {episode_reward:.2f}")
|
||||
print(
|
||||
f" Eval Episode {idx} (scenario {scenario_id}): Total Reward {episode_reward:.2f}, steps {step_count}, "
|
||||
f"agents (current): {n_controlled}, total in scenario: {n_total_in_scenario}"
|
||||
)
|
||||
env.close()
|
||||
|
||||
if not total_rewards:
|
||||
print(" No valid eval episodes (all skipped or failed).")
|
||||
return 0.0
|
||||
avg_reward = float(np.mean(total_rewards))
|
||||
print(f" Average Evaluation Reward: {avg_reward:.2f}")
|
||||
return avg_reward
|
||||
except Exception as e:
|
||||
print(f"Evaluation failed: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return 0.0
|
||||
finally:
|
||||
env.close()
|
||||
|
||||
|
||||
def main(args):
|
||||
|
||||
Reference in New Issue
Block a user