| |
| """A stand-in for :class:`tt_diffusion_planner.api.DiffusionPlanner` with no device and no weights, for the host tests |
| of the server and the smoke test. It quacks like the real model (``validate_params``, ``__call__``, ``info``, |
| ``device_info``, ``close``) and runs the bundle's REAL host code around a fake network: the inputs are checked |
| against ``INPUT_SCHEMA`` and pre-processed by ``host.prepare`` (normalization from the v5.0 statistics embedded below), |
| the "network" extrapolates every agent's current state at a constant normalised speed, and ``host.make_output`` |
| post-processes it, so the output depends on the input deterministically and exercises the real post-processing.""" |
| from __future__ import annotations |
|
|
| from typing import Any, Dict, Optional |
|
|
| import numpy as np |
|
|
| from tt_diffusion_planner import io as tio |
| from tt_diffusion_planner.api import DiffusionPlanner, Output |
| from tt_diffusion_planner.host import pipeline as hp |
| from tt_diffusion_planner.reference import config as C |
| from tt_diffusion_planner.reference.weights import Normalization |
|
|
| ETH_P150 = {"dispatch": "eth", "dispatch_requested": "eth", "fallback": None, "grid": "12x10", "grid_x": 12, |
| "grid_y": 10, "cores": 120, "arch": "blackhole", "eth_patch": True, "device_id": 0, |
| "num_command_queues": DiffusionPlanner.DEVICE_DEFAULTS["num_command_queues"]} |
|
|
|
|
| def v5_normalization() -> Normalization: |
| """The v5.0 ``observation_normalizer`` / ``state_normalizer`` (HF diffusion_planner.param.json; SPEC 3.7), so the |
| stub needs no weights download.""" |
| def mv(mean, std): |
| return np.asarray(mean, np.float32), np.asarray(std, np.float32) |
|
|
| lane_std = [20.0] * 8 + [1.0] * 25 |
| obs = {"ego_agent_past": mv([10, 0, 0, 0], [20, 20, 1, 1]), |
| "ego_current_state": mv([10] + [0] * 9, [20, 20, 1, 1, 20, 20, 20, 20, 1, 1]), |
| "neighbor_agents_past": mv([10] + [0] * 10, [20, 20, 1, 1, 20, 20, 20, 20, 1, 1, 1]), |
| "static_objects": mv([10] + [0] * 9, [20, 20, 1, 1, 20, 20, 1, 1, 1, 1]), |
| "lanes": mv([10] + [0] * 32, lane_std), "route_lanes": mv([10] + [0] * 32, lane_std), |
| "lanes_speed_limit": mv([0], [20]), "route_lanes_speed_limit": mv([0], [20]), |
| "polygons": mv([10, 0, 0], [20, 20, 1]), "line_strings": mv([10, 0, 0, 0], [20, 20, 1, 1]), |
| "goal_pose": mv([10, 0, 0, 0], [20, 20, 1, 1])} |
| mean = np.tile(np.float32([10, 0, 0, 0]), (C.MAX_NUM_AGENTS, 1, 1)) |
| std = np.tile(np.float32([20, 20, 1, 1]), (C.MAX_NUM_AGENTS, 1, 1)) |
| return Normalization(obs, mean, std, 5) |
|
|
|
|
| class StubModel: |
| """``StubModel(cfg)`` is a valid ``app.state.ttaw.model_factory`` (``cfg`` = ``config_from_env`` output).""" |
|
|
| device_info: Dict[str, Any] = ETH_P150 |
| validate_params = DiffusionPlanner.validate_params |
| normalization = v5_normalization() |
|
|
| def __init__(self, cfg: Optional[Dict[str, Any]] = None): |
| self.cfg = dict(cfg or {}) |
| self.closed = False |
|
|
| @property |
| def info(self) -> Dict[str, Any]: |
| cls = DiffusionPlanner |
| return {"model": cls.MODEL_NAME, "variant": self.cfg.get("variant") or cls.DEFAULT_VARIANT, |
| "weights": {"repo": self.cfg.get("model_id") or cls.DEFAULT_REPO, |
| "revision": self.cfg.get("revision") or cls.DEFAULT_REVISION}, |
| "device": self.device_info, "input_kind": cls.INPUT_KIND, "point_fields": list(cls.POINT_FIELDS), |
| "camera_order": list(cls.CAMERA_ORDER), "labels": list(cls.LABELS), |
| "runtime_params": {k: v[3] for k, v in cls.RUNTIME_PARAMS.items()}, |
| "input_schema": {k: {"shape": list(s), "dtype": "float32"} for k, (s, _) in cls.INPUT_SCHEMA.items()}} |
|
|
| @staticmethod |
| def fake_network(prep: hp.Prepared) -> Dict[str, np.ndarray]: |
| """Every agent moves along its current heading at 0.25 normalised units (5 m) per second.""" |
| cs = prep.decoder.current_states |
| t = np.arange(C.OUTPUT_T + 1, dtype=np.float32)[None, :, None] * np.float32(0.025) |
| x0 = np.repeat(cs[:, None, :], C.OUTPUT_T + 1, axis=1).astype(np.float32) |
| x0[..., 0:1] += t * cs[:, None, 2:3] |
| x0[..., 1:2] += t * cs[:, None, 3:4] |
| logit = np.array([-4.0, 1.0 + float(np.tanh(cs[0, 0])), 0.5, 0.25, 0.0], np.float32) |
| return {"final_x0": x0, "logit": logit} |
|
|
| def __call__(self, points: Any = None, *, inputs: Any = None, **params: Any) -> Output: |
| p = self.validate_params(params) |
| if points is not None: |
| raise tio.InputError("this model takes only `inputs` (the planner tensors), not ['points']") |
| if inputs is None: |
| raise tio.InputError("this model needs `inputs`: the 15 planner tensors of INPUT_SCHEMA") |
| arrays = tio.load_inputs(inputs) |
| prep = hp.prepare(arrays, self.normalization.observation) |
| raw = self.fake_network(prep) |
| return hp.make_output(raw["final_x0"], raw["logit"], prep, self.normalization, p, |
| model=DiffusionPlanner.MODEL_NAME, timing_ms={"device": 1.0}) |
|
|
| def close(self) -> None: |
| self.closed = True |
|
|
|
|
| def sample_inputs(seed: int = 0, n_neighbors: int = 3) -> Dict[str, np.ndarray]: |
| """A small synthetic planner scene in the raw (pre-normalization) layout of INPUT_SCHEMA.""" |
| rng = np.random.default_rng(seed) |
| raw = {k: np.zeros(s, np.float32) for k, s in C.INPUT_SHAPES.items()} |
| t = np.arange(C.INPUT_T + 1, dtype=np.float32) |
| raw["ego_agent_past"][0, :, 0] = (t - C.INPUT_T) * 0.8 |
| raw["ego_agent_past"][0, :, 2] = 1.0 |
| raw["ego_current_state"][0] = (0, 0, 1, 0, 8.0, 0, 0.2, 0, 0, 0) |
| for i in range(n_neighbors): |
| x0, y0, v = float(rng.uniform(5, 40)), float(rng.choice([-3.5, 0.0, 3.5])), float(rng.uniform(2, 10)) |
| nb = raw["neighbor_agents_past"][0, i] |
| nb[:, 0], nb[:, 1], nb[:, 2] = x0 + (t - C.INPUT_T) * v * 0.1, y0, 1.0 |
| nb[:, 4], nb[:, 6], nb[:, 7], nb[:, 8] = v, 1.8, 4.5, 1.0 |
| for k in range(4): |
| lane = raw["lanes"][0, k] |
| lane[:, 0], lane[:, 1] = np.linspace(-20 + 20 * k, 20 * k, 20), 0.0 |
| lane[:19, 2] = lane[1:, 0] - lane[:-1, 0] |
| lane[:, 5], lane[:, 7], lane[:, 12] = 1.75, -1.75, 1.0 |
| raw["lanes_speed_limit"][0, k, 0] = 13.9 |
| raw["route_lanes"][0, :4] = raw["lanes"][0, :4] |
| raw["route_lanes_speed_limit"][0, :4] = 13.9 |
| raw["goal_pose"][0] = (150.0, 0.0, 1.0, 0.0) |
| raw["ego_shape"][0] = (2.79, 4.89, 1.896) |
| raw["turn_indicators"][:] = 1.0 |
| return raw |
|
|
