# SPDX-License-Identifier: Apache-2.0 """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 # [321, 4] normalised 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