code/tt_diffusion_planner/tests/stubs.py

# 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