195 lines
9.9 KiB
Python
195 lines
9.9 KiB
Python
"""
|
|
Single-agent BC evaluation environment: only ego (SDC) is controlled by the policy;
|
|
other vehicles are replayed from expert trajectories (same as data collection).
|
|
"""
|
|
import numpy as np
|
|
from Env.expert_replay_env import ExpertReplayEnv
|
|
from Env.hbbc_background_policy import HBBCBackgroundController
|
|
|
|
|
|
class BCEgoReplayEnv(ExpertReplayEnv):
|
|
"""
|
|
For single-agent BC evaluation: controlled_agents exposes only SDC (default_agent).
|
|
Other vehicles are still spawned and replayed by expert; internally we keep them
|
|
in _replay_agents so step() can update them.
|
|
"""
|
|
|
|
def reset(self, seed=None):
|
|
obs = super().reset(seed=seed)
|
|
self.enable_hbbc_background = bool(self.config.get("enable_hbbc_background", False))
|
|
self.hbbc_controller = None
|
|
self._hbbc_runtime_logged = False
|
|
if self.enable_hbbc_background:
|
|
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()
|
|
# Expose only SDC as the controlled agent for the evaluator
|
|
self._replay_agents = dict(self.controlled_agents)
|
|
if self.replay_sdc and self.sdc_vehicle is not None:
|
|
self.controlled_agents = {self.sdc_agent_id: self.sdc_vehicle}
|
|
self.controlled_agent_ids = [self.sdc_agent_id]
|
|
else:
|
|
self.controlled_agents = {}
|
|
self.controlled_agent_ids = []
|
|
return self._get_all_obs()
|
|
|
|
def _get_all_obs(self):
|
|
"""Return only ego (SDC) observation so evaluator has a single agent."""
|
|
if not self.controlled_agents or self.sdc_vehicle is None:
|
|
return {}
|
|
obs = self._obs_for_vehicle(self.sdc_vehicle, exclude_agent_id=self.sdc_agent_id)
|
|
return {self.sdc_agent_id: obs}
|
|
|
|
def step(self, action_dict=None):
|
|
self.round += 1
|
|
expert_actions = {}
|
|
agents_to_remove = []
|
|
|
|
# SDC: use policy action if provided, else expert replay
|
|
if self.replay_sdc and self.sdc_vehicle is not None and self.sdc_track is not None:
|
|
policy_action = None
|
|
if action_dict and self.sdc_agent_id in action_dict:
|
|
policy_action = np.asarray(action_dict[self.sdc_agent_id], dtype=np.float64)
|
|
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],
|
|
}
|
|
if policy_action is not None:
|
|
next_state = self.inverse_dynamics.apply_action(curr_state, policy_action, dt=0.1)
|
|
expert_actions[self.sdc_agent_id] = policy_action
|
|
else:
|
|
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 = expert_actions[self.sdc_agent_id]
|
|
|
|
# Replay other vehicles: restore full controlled_agents for internal logic
|
|
self.controlled_agents = dict(self._replay_agents)
|
|
self.controlled_agent_ids = list(self.controlled_agents.keys())
|
|
hbbc_batch = []
|
|
hbbc_curr_states = {}
|
|
for agent_id, vehicle in self.controlled_agents.items():
|
|
track = vehicle.expert_track
|
|
next_step = self.round
|
|
if next_step >= len(track["state"]["position"]):
|
|
agents_to_remove.append(agent_id)
|
|
continue
|
|
if not track["state"]["valid"][next_step]:
|
|
agents_to_remove.append(agent_id)
|
|
continue
|
|
if self.enable_hbbc_background and self.hbbc_controller is not None:
|
|
# HBBC autonomous rollout: use vehicle's own previous-step state
|
|
curr_state = {
|
|
"position": np.asarray(vehicle.position, dtype=np.float64),
|
|
"heading": float(vehicle.heading_theta),
|
|
"velocity": np.asarray(vehicle.velocity, dtype=np.float64),
|
|
}
|
|
object_id = str(getattr(vehicle, "original_id", agent_id))
|
|
hbbc_batch.append((agent_id, vehicle, object_id, agent_id))
|
|
hbbc_curr_states[agent_id] = curr_state
|
|
else:
|
|
curr_step = self.round - 1
|
|
curr_state = {
|
|
"position": track["state"]["position"][curr_step],
|
|
"heading": track["state"]["heading"][curr_step],
|
|
"velocity": track["state"]["velocity"][curr_step],
|
|
}
|
|
next_state = {
|
|
"position": track["state"]["position"][next_step],
|
|
"heading": track["state"]["heading"][next_step],
|
|
"velocity": track["state"]["velocity"][next_step],
|
|
}
|
|
action, _ = self.inverse_dynamics.compute_action(curr_state, next_state, dt=0.1)
|
|
expert_actions[agent_id] = action
|
|
vehicle.set_position(next_state["position"])
|
|
vehicle.set_heading_theta(next_state["heading"])
|
|
vehicle.set_velocity(next_state["velocity"])
|
|
vehicle.last_expert_action = action
|
|
|
|
if hbbc_batch and self.hbbc_controller is not None:
|
|
hbbc_actions = self.hbbc_controller.infer_actions(hbbc_batch)
|
|
if not self._hbbc_runtime_logged:
|
|
print(f"[HBBC] background policy active, current dynamic agents: {len(hbbc_batch)}")
|
|
self._hbbc_runtime_logged = True
|
|
for agent_id, _, _, _ in hbbc_batch:
|
|
curr_state = hbbc_curr_states[agent_id]
|
|
action = hbbc_actions[agent_id]
|
|
next_state = self.inverse_dynamics.apply_action(curr_state, action, dt=0.1)
|
|
expert_actions[agent_id] = action
|
|
vehicle = self.controlled_agents[agent_id]
|
|
vehicle.set_position(next_state["position"])
|
|
vehicle.set_heading_theta(next_state["heading"])
|
|
vehicle.set_velocity(next_state["velocity"])
|
|
try:
|
|
vehicle.last_current_action.append(action)
|
|
except Exception:
|
|
pass
|
|
vehicle.last_expert_action = action
|
|
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])
|
|
if self.hbbc_controller is not None:
|
|
self.hbbc_controller.remove_vehicle(agent_id)
|
|
self.engine.taskMgr.step()
|
|
self._spawn_controlled_agents()
|
|
self._update_background_vehicles()
|
|
self._replay_agents = dict(self.controlled_agents)
|
|
# Expose only SDC again
|
|
if self.replay_sdc and self.sdc_vehicle is not None:
|
|
self.controlled_agents = {self.sdc_agent_id: self.sdc_vehicle}
|
|
self.controlled_agent_ids = [self.sdc_agent_id]
|
|
else:
|
|
self.controlled_agents = {}
|
|
self.controlled_agent_ids = []
|
|
|
|
obs = self._get_all_obs()
|
|
rewards = {}
|
|
infos = {aid: {"expert_action": expert_actions.get(aid, np.zeros(2))} for aid in self.controlled_agents}
|
|
if self.sdc_agent_id in self.controlled_agents and self.sdc_vehicle is not None:
|
|
speed_coef = float(self.config.get("reward_speed_coef", 0.05))
|
|
collision_distance = float(self.config.get("collision_distance", 6.0))
|
|
collision_penalty = float(self.config.get("collision_penalty", 100.0))
|
|
speed = float(np.linalg.norm(self.sdc_vehicle.velocity))
|
|
r_speed = speed_coef * speed
|
|
min_dist = float("inf")
|
|
for other_id, other_vehicle in self.engine.agent_manager.active_agents.items():
|
|
if other_id == self.sdc_agent_id:
|
|
continue
|
|
try:
|
|
d = float(np.linalg.norm(self.sdc_vehicle.position - other_vehicle.position))
|
|
min_dist = min(min_dist, d)
|
|
except Exception:
|
|
continue
|
|
near_collision = min_dist < collision_distance
|
|
r_collision = -collision_penalty if near_collision else 0.0
|
|
rewards[self.sdc_agent_id] = r_speed + r_collision
|
|
infos[self.sdc_agent_id].update(
|
|
near_collision=near_collision,
|
|
min_dist=min_dist if np.isfinite(min_dist) else None,
|
|
r_speed=r_speed,
|
|
r_collision=r_collision,
|
|
)
|
|
dones = {aid: False for aid in self.controlled_agents}
|
|
dones["__all__"] = self.round >= self.config["horizon"] or (len(self._replay_agents) == 0 and self.round > 190)
|
|
return obs, rewards, dones, infos
|