275 lines
11 KiB
Python
275 lines
11 KiB
Python
import json
|
|
import os
|
|
from dataclasses import dataclass
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from Env.hbbc_actor_critic import ActorCritic
|
|
|
|
|
|
def _wrap_to_pi(angle: float) -> float:
|
|
return (angle + np.pi) % (2 * np.pi) - np.pi
|
|
|
|
|
|
def _normalize_eps(eps: np.ndarray) -> np.ndarray:
|
|
eps = np.asarray(eps, dtype=np.float32).reshape(-1)
|
|
if eps.shape[0] != 6:
|
|
raise ValueError(f"latent_eps must be 6-dim, got {eps.shape[0]}")
|
|
norm = float(np.linalg.norm(eps))
|
|
if norm < 1e-8:
|
|
eps = np.array([1.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float32)
|
|
else:
|
|
eps = eps / norm
|
|
return np.clip(eps, -1.0, 1.0)
|
|
|
|
|
|
def _normalize_c(latent_c: np.ndarray) -> np.ndarray:
|
|
c = np.asarray(latent_c, dtype=np.float32).reshape(-1)
|
|
if c.shape[0] != 4:
|
|
raise ValueError(f"latent_c must be 4-dim, got {c.shape[0]}")
|
|
idx = int(np.argmax(c))
|
|
one_hot = np.zeros(4, dtype=np.float32)
|
|
one_hot[idx] = 1.0
|
|
return one_hot
|
|
|
|
|
|
def _sample_latent(rng: np.random.RandomState) -> Tuple[np.ndarray, np.ndarray]:
|
|
eps = _normalize_eps(rng.randn(6).astype(np.float32))
|
|
mode = int(rng.randint(0, 4))
|
|
c = np.zeros(4, dtype=np.float32)
|
|
c[mode] = 1.0
|
|
return eps, c
|
|
|
|
|
|
@dataclass
|
|
class VehicleStateCache:
|
|
last_heading_theta: Optional[float] = None
|
|
last_action: Tuple[float, float] = (0.0, 0.0)
|
|
last_speed_km_h: Optional[float] = None
|
|
|
|
|
|
class HBBCModelWrapper:
|
|
_cache: Dict[Tuple[str, str], "HBBCModelWrapper"] = {}
|
|
|
|
def __init__(self, model_path: str, device: str = "cpu"):
|
|
self.model_path = os.path.abspath(model_path)
|
|
self.device = torch.device(device)
|
|
self.model = self._load_model()
|
|
|
|
@classmethod
|
|
def get(cls, model_path: str, device: str = "cpu") -> "HBBCModelWrapper":
|
|
key = (os.path.abspath(model_path), str(torch.device(device)))
|
|
if key not in cls._cache:
|
|
cls._cache[key] = HBBCModelWrapper(model_path=key[0], device=key[1])
|
|
return cls._cache[key]
|
|
|
|
def _load_model(self) -> ActorCritic:
|
|
model = ActorCritic(
|
|
num_actor_obs=18,
|
|
num_critic_obs=18,
|
|
num_actions=2,
|
|
latent_c_dim=4,
|
|
latent_eps_dim=6,
|
|
use_style_latent=True,
|
|
).to(self.device)
|
|
try:
|
|
ckpt = torch.load(self.model_path, map_location=self.device, weights_only=True)
|
|
except Exception:
|
|
ckpt = torch.load(self.model_path, map_location=self.device, weights_only=False)
|
|
state_dict = ckpt["actor_critic"] if isinstance(ckpt, dict) and "actor_critic" in ckpt else ckpt
|
|
missing, unexpected = model.load_state_dict(state_dict, strict=False)
|
|
if missing:
|
|
raise RuntimeError(
|
|
f"HBBC checkpoint missing required keys for {self.model_path}: {missing}"
|
|
)
|
|
if unexpected:
|
|
print(f"[HBBC] ignore extra checkpoint keys: {unexpected[:8]}{'...' if len(unexpected) > 8 else ''}")
|
|
model.eval()
|
|
return model
|
|
|
|
def act_batch(self, obs_batch: np.ndarray) -> np.ndarray:
|
|
obs_batch = np.asarray(obs_batch, dtype=np.float32)
|
|
with torch.no_grad():
|
|
obs_t = torch.from_numpy(obs_batch).to(self.device)
|
|
actions = self.model.act_inference(obs_t).cpu().numpy()
|
|
return np.clip(actions, -1.0, 1.0)
|
|
|
|
|
|
class HBBCLatentManager:
|
|
def __init__(self, mode: str = "per_vehicle_fixed", seed: int = 0, latent_json_path: Optional[str] = None):
|
|
self.mode = mode
|
|
self.rng = np.random.RandomState(seed)
|
|
self.latent_json_path = latent_json_path
|
|
self.manual_object_latent: Dict[str, Dict[str, np.ndarray]] = {}
|
|
self.manual_agent_latent: Dict[str, Dict[str, np.ndarray]] = {}
|
|
self.manual_global_latent: Optional[Tuple[np.ndarray, np.ndarray]] = None
|
|
self.vehicle_latent: Dict[str, Tuple[np.ndarray, np.ndarray]] = {}
|
|
self._episode_latent: Optional[Tuple[np.ndarray, np.ndarray]] = None
|
|
self._load_manual_latent_json()
|
|
|
|
def reset_episode(self):
|
|
self.vehicle_latent.clear()
|
|
self._episode_latent = None
|
|
if self.mode == "per_episode_reset":
|
|
self._episode_latent = _sample_latent(self.rng)
|
|
|
|
def _load_manual_latent_json(self):
|
|
if not self.latent_json_path:
|
|
return
|
|
path = os.path.abspath(self.latent_json_path)
|
|
if not os.path.exists(path):
|
|
print(f"[HBBC] latent json not found: {path}, fallback to random sampling.")
|
|
return
|
|
try:
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
except Exception as e:
|
|
print(f"[HBBC] failed to load latent json ({path}): {e}. fallback to random sampling.")
|
|
return
|
|
|
|
object_section = data.get("object_id", {})
|
|
agent_section = data.get("agent_id", {})
|
|
global_section = data.get("global")
|
|
|
|
if global_section is not None:
|
|
parsed = self._parse_one_latent(global_section, "global")
|
|
if parsed is not None:
|
|
self.manual_global_latent = (parsed["latent_eps"], parsed["latent_c"])
|
|
|
|
for key, value in object_section.items():
|
|
parsed = self._parse_one_latent(value, f"object_id:{key}")
|
|
if parsed is not None:
|
|
self.manual_object_latent[str(key)] = parsed
|
|
for key, value in agent_section.items():
|
|
parsed = self._parse_one_latent(value, f"agent_id:{key}")
|
|
if parsed is not None:
|
|
self.manual_agent_latent[str(key)] = parsed
|
|
|
|
@staticmethod
|
|
def _parse_one_latent(value: dict, name: str) -> Optional[Dict[str, np.ndarray]]:
|
|
if not isinstance(value, dict):
|
|
print(f"[HBBC] invalid latent entry ({name}): expect dict.")
|
|
return None
|
|
try:
|
|
eps = _normalize_eps(value["latent_eps"])
|
|
c = _normalize_c(value["latent_c"])
|
|
return {"latent_eps": eps, "latent_c": c}
|
|
except Exception as e:
|
|
print(f"[HBBC] invalid latent entry ({name}): {e}")
|
|
return None
|
|
|
|
def _lookup_manual(self, object_id: Optional[str], agent_id: Optional[str]) -> Optional[Tuple[np.ndarray, np.ndarray]]:
|
|
if object_id is not None and object_id in self.manual_object_latent:
|
|
e = self.manual_object_latent[object_id]["latent_eps"]
|
|
c = self.manual_object_latent[object_id]["latent_c"]
|
|
return e, c
|
|
if agent_id is not None and agent_id in self.manual_agent_latent:
|
|
e = self.manual_agent_latent[agent_id]["latent_eps"]
|
|
c = self.manual_agent_latent[agent_id]["latent_c"]
|
|
return e, c
|
|
if self.manual_global_latent is not None:
|
|
return self.manual_global_latent
|
|
return None
|
|
|
|
def get_latent(self, vehicle_key: str, object_id: Optional[str], agent_id: Optional[str]) -> Tuple[np.ndarray, np.ndarray]:
|
|
manual = self._lookup_manual(object_id=object_id, agent_id=agent_id)
|
|
if manual is not None:
|
|
return manual
|
|
if self.mode == "per_episode_reset":
|
|
if self._episode_latent is None:
|
|
self._episode_latent = _sample_latent(self.rng)
|
|
return self._episode_latent
|
|
if vehicle_key not in self.vehicle_latent:
|
|
self.vehicle_latent[vehicle_key] = _sample_latent(self.rng)
|
|
return self.vehicle_latent[vehicle_key]
|
|
|
|
|
|
class HBBCBackgroundController:
|
|
def __init__(
|
|
self,
|
|
model_path: str,
|
|
device: str = "cpu",
|
|
latent_mode: str = "per_vehicle_fixed",
|
|
latent_json_path: Optional[str] = None,
|
|
seed: int = 0,
|
|
dt: float = 0.1,
|
|
):
|
|
self.model = HBBCModelWrapper.get(model_path=model_path, device=device)
|
|
self.latent_mgr = HBBCLatentManager(mode=latent_mode, seed=seed, latent_json_path=latent_json_path)
|
|
self.dt = float(dt)
|
|
self.vehicle_state: Dict[str, VehicleStateCache] = {}
|
|
|
|
def reset_episode(self):
|
|
self.latent_mgr.reset_episode()
|
|
self.vehicle_state.clear()
|
|
|
|
def remove_vehicle(self, vehicle_key: str):
|
|
self.vehicle_state.pop(vehicle_key, None)
|
|
self.latent_mgr.vehicle_latent.pop(vehicle_key, None)
|
|
|
|
def _build_base_state(self, vehicle, vehicle_key: str) -> np.ndarray:
|
|
state = self.vehicle_state.get(vehicle_key)
|
|
if state is None:
|
|
state = VehicleStateCache()
|
|
self.vehicle_state[vehicle_key] = state
|
|
|
|
speed_km_h = float(getattr(vehicle, "speed_km_h", 0.0))
|
|
max_speed_km_h = float(getattr(vehicle, "max_speed_km_h", 120.0))
|
|
veh_vel = np.clip((speed_km_h + 1.0) / (max_speed_km_h + 1.0), 0.0, 1.0)
|
|
|
|
heading_theta = float(getattr(vehicle, "heading_theta", 0.0))
|
|
if state.last_heading_theta is None:
|
|
yaw_rate = 0.0
|
|
else:
|
|
yaw_rate = _wrap_to_pi(heading_theta - state.last_heading_theta) / self.dt
|
|
yaw_rate = float(np.clip(yaw_rate, -5.0, 5.0))
|
|
|
|
current_action = getattr(vehicle, "current_action", None)
|
|
if current_action is None:
|
|
last_action_0, last_action_1 = state.last_action
|
|
else:
|
|
try:
|
|
last_action_0, last_action_1 = float(current_action[0]), float(current_action[1])
|
|
except Exception:
|
|
last_action_0, last_action_1 = state.last_action
|
|
|
|
state.last_heading_theta = heading_theta
|
|
state.last_speed_km_h = speed_km_h
|
|
state.last_action = (last_action_0, last_action_1)
|
|
|
|
obs = np.array(
|
|
[
|
|
0.0,
|
|
0.0,
|
|
0.0,
|
|
veh_vel,
|
|
0.0,
|
|
yaw_rate * 0.5,
|
|
last_action_0,
|
|
last_action_1,
|
|
],
|
|
dtype=np.float32,
|
|
)
|
|
return obs
|
|
|
|
def build_obs(self, vehicle, vehicle_key: str, object_id: Optional[str], agent_id: Optional[str]) -> np.ndarray:
|
|
base = self._build_base_state(vehicle, vehicle_key=vehicle_key)
|
|
eps, c = self.latent_mgr.get_latent(vehicle_key=vehicle_key, object_id=object_id, agent_id=agent_id)
|
|
return np.concatenate([base, eps, c], axis=-1).astype(np.float32)
|
|
|
|
def infer_actions(self, batch: List[Tuple[str, object, Optional[str], Optional[str]]]) -> Dict[str, np.ndarray]:
|
|
if not batch:
|
|
return {}
|
|
obs_list = []
|
|
vehicle_ids = []
|
|
for vehicle_key, vehicle, object_id, agent_id in batch:
|
|
obs_list.append(self.build_obs(vehicle, vehicle_key=vehicle_key, object_id=object_id, agent_id=agent_id))
|
|
vehicle_ids.append(vehicle_key)
|
|
actions = self.model.act_batch(np.stack(obs_list, axis=0))
|
|
out = {}
|
|
for idx, key in enumerate(vehicle_ids):
|
|
out[key] = actions[idx].astype(np.float32)
|
|
return out
|