code/tt_diffusion_planner/host/postprocess.py

# SPDX-License-Identifier: Apache-2.0
"""Post-processing of the Autoware node (``postprocessing_utils.cpp``, ``turn_indicator_manager.cpp``,
``diffusion_planner_core.cpp:605-694``), in the ego frame of the input tensors.

The node transforms every predicted pose to the map frame before building messages; the bundle has no map pose, so
it applies the same rules with an identity ego-to-map transform (``base_link`` output). Distances, velocities and
accelerations are invariant to that rigid transform; the published orientation is not exactly (see below).

Rules reproduced, with their numeric types:

- denormalisation: drop t = 0 (unless kept), ``v * std + mean`` per agent and pose dim, float32;
- poses: 4x4 matrices in double from the RAW (unnormalised) network cos / sin; the quaternion is Eigen's conversion
  of that 3x3 block, which does not normalise, so for ``|(cos, sin)| != 1`` (0.84-1.09 in the goldens, SPEC 4.6.6)
  the heading a consumer reads with ``tf2::getYaw`` differs from ``atan2(sin, cos)``; ``yaw`` here is that
  ``tf2::getYaw`` value of the identity-transform quaternion;
- ego trajectory (``get_trajectory_from_poses``): ``time_from_start = 0.1 (i + 1)``; velocity = 3-D distance to the
  previous point (the first point to the ego position) / 0.1, stored as float; forward moving average over
  ``velocity_smoothing_window`` points summed in double; force stop (only when the ego moves: ``vx > DBL_EPSILON``)
  once the smoothed velocity crosses below ``stopping_threshold``: velocity 0 and the pose frozen to the previous
  pose; the last ``window - 1`` points keep the last smoothed velocity; acceleration = forward difference / 0.1 in
  double stored as float, 0 for the last point;
- predicted objects: one 80-point path per emitted neighbour (the first ``n`` rows of the neighbour tensor, sorted
  by distance), window 1, no force stop (only the poses are published);
- turn indicator (``TurnIndicatorManager::evaluate``): hold the last non-KEEP command for ``hold_duration`` (state
  across calls), else add ``keep_offset`` to the KEEP logit, ``p_i = expf(l_i - max)``, ``p_i /= (1e-4f + sum)``,
  argmax (first maximum); KEEP repeats the previous TurnIndicatorsReport, otherwise the command is the index;
- ``~/debug/denoising_steps``: the ego row of every solver iterate, denormalised with t = 0 kept.
"""
from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Dict, Optional, Tuple

import numpy as np

from ..reference import config as C
from .solver import SCALAR_MATH

__all__ = ["denormalize", "quaternion_from_cos_sin", "tf2_yaw", "EgoTrajectory", "trajectory_from_poses",
           "predicted_paths", "TurnIndicatorManager", "TurnDecision", "denoising_steps_ego"]

DBL_EPSILON = float(np.finfo(np.float64).eps)


def denormalize(x: np.ndarray, state_mean: np.ndarray, state_std: np.ndarray, *,
                keep_current_state: bool = False) -> np.ndarray:
    """``denormalize_prediction``: ``x`` ``[321, 81, 4]`` (normalised, t = 0 = current state) -> ``[321, 80, 4]``
    metres (``[321, 81, 4]`` with ``keep_current_state``). ``state_mean`` / ``state_std`` broadcast as
    ``[321 or 1, 1, 4]``."""
    x = np.asarray(x, np.float32)
    if x.shape[-3:] != (C.MAX_NUM_AGENTS, C.OUTPUT_T + 1, C.POSE_DIM):
        raise ValueError(f"unsupported prediction shape {x.shape}")
    p = x if keep_current_state else x[..., 1:, :]
    return (p * np.asarray(state_std, np.float32) + np.asarray(state_mean, np.float32)).astype(np.float32)


def quaternion_from_cos_sin(cos_yaw: np.ndarray, sin_yaw: np.ndarray) -> np.ndarray:
    """Eigen ``Quaterniond(rotation_matrix)`` of ``[[c, -s, 0], [s, c, 0], [0, 0, 1]]`` (no normalisation), in double:
    ``[..., 4]`` = (x, y, z, w). Eigen's branch: trace = 2c + 1 > 0 -> w = sqrt(trace + 1) / 2, z = (m10 - m01) / (4 w);
    else the largest diagonal is m22 -> z = sqrt(m22 - m00 - m11 + 1) / 2, w = (m10 - m01) / (4 z)."""
    c = np.asarray(cos_yaw, np.float64)
    s = np.asarray(sin_yaw, np.float64)
    q = np.zeros(c.shape + (4,), np.float64)
    trace = c + c + 1.0
    pos = trace > 0.0
    with np.errstate(divide="ignore", invalid="ignore"):
        t = np.sqrt(trace + 1.0)
        w_pos, z_pos = 0.5 * t, (s - (-s)) * (0.5 / t)
        # trace <= 0: i = 2 (m22 = 1 > m00 = c, since c <= -0.5), j = 0, k = 1
        t2 = np.sqrt(1.0 - c - c + 1.0)
        z_neg, w_neg = 0.5 * t2, (s - (-s)) * (0.5 / t2)
    q[..., 2] = np.where(pos, z_pos, z_neg)
    q[..., 3] = np.where(pos, w_pos, w_neg)
    return q


def tf2_yaw(q: np.ndarray) -> np.ndarray:
    """``tf2::getYaw`` (scale-invariant): ``atan2(2 (x y + w z), w^2 + x^2 - y^2 - z^2)``."""
    x, y, z, w = (q[..., i] for i in range(4))
    return np.arctan2(2.0 * (x * y + w * z), w * w + x * x - y * y - z * z)


@dataclass
class EgoTrajectory:
    """One trajectory as the node publishes it (in the ego frame here): positions [N, 3] (double), the raw cos / sin
    of each pose [N, 2] (after force-stop freezing), quaternions [N, 4] (x, y, z, w), yaw [N] (``tf2::getYaw``),
    ``velocity`` / ``acceleration`` [N] float32, ``time_from_start`` [N] seconds, and whether force stop fired."""

    position: np.ndarray
    cos_sin: np.ndarray
    quaternion: np.ndarray
    yaw: np.ndarray
    velocity: np.ndarray
    acceleration: np.ndarray
    time_from_start: np.ndarray
    force_stop: bool

    def as_columns(self) -> np.ndarray:
        """``[N, 7]`` float32: x, y, yaw, cos, sin, velocity, acceleration (the ``Trajectory.poses`` layout)."""
        return np.stack([self.position[:, 0], self.position[:, 1], self.yaw, self.cos_sin[:, 0], self.cos_sin[:, 1],
                         self.velocity, self.acceleration], axis=-1).astype(np.float32)



def _host_fast() -> bool:
    """``DIFFUSION_PLANNER_HOST_FAST`` (OPT round 5 item 3), read once: the vectorised host functions (bit-exact)
    instead of the first port's ``*_ref`` versions."""
    from ..tt.config import KNOBS

    return bool(KNOBS.read().HOST_FAST)


HOST_FAST = _host_fast()


def trajectory_from_poses_ref(poses_xycs: np.ndarray, base_position: Tuple[float, float, float] = (0.0, 0.0, 0.0), *,
                          velocity_smoothing_window: int = C.VELOCITY_SMOOTHING_WINDOW,
                          enable_force_stop: bool = True,
                          stopping_threshold: float = C.STOPPING_THRESHOLD) -> EgoTrajectory:
    """``get_trajectory_from_poses`` (``postprocessing_utils.cpp:362-455``) for poses ``[N, 4]`` = denormalised
    (x, y, cos, sin) float32 of one agent, positions at z = ``base_position[2]`` (identity transform)."""
    p = np.asarray(poses_xycs, np.float32)
    n = p.shape[0]
    dt = C.TRAJECTORY_DT
    pos = np.zeros((n, 3), np.float64)
    pos[:, 0] = p[:, 0].astype(np.float64)
    pos[:, 1] = p[:, 1].astype(np.float64)
    pos[:, 2] = float(base_position[2])
    cs = p[:, 2:4].astype(np.float64).copy()
    quat = quaternion_from_cos_sin(cs[:, 0], cs[:, 1])
    vel = np.zeros(n, np.float32)
    prev = np.asarray(base_position, np.float64)
    for i in range(n):
        d = math.hypot(pos[i, 0] - prev[0], pos[i, 1] - prev[1], pos[i, 2] - prev[2])
        vel[i] = np.float32(d / dt)
        prev = pos[i]
    w = int(velocity_smoothing_window)
    if n <= w:
        raise ValueError("velocity_smoothing_window must be smaller than number of points")
    thr = np.float32(stopping_threshold)
    force_stop = False
    for i in range(0, n - w + 1):
        acc = 0.0
        for k in range(w):
            acc += float(vel[i + k])
        vel[i] = np.float32(acc / float(w))
        if enable_force_stop and i > 0 and abs(vel[i - 1]) > thr and abs(vel[i]) < thr:
            force_stop = True
        if i > 0 and force_stop:
            vel[i] = np.float32(0.0)
            pos[i], cs[i], quat[i] = pos[i - 1], cs[i - 1], quat[i - 1]
    last = vel[n - w]
    for i in range(n - w + 1, n):
        vel[i] = last
        if force_stop:
            vel[i] = np.float32(0.0)
            pos[i], cs[i], quat[i] = pos[i - 1], cs[i - 1], quat[i - 1]
    accel = np.zeros(n, np.float32)
    for i in range(n - 1):
        accel[i] = np.float32((float(vel[i + 1]) - float(vel[i])) / dt)
    tfs = dt * (np.arange(n, dtype=np.float64) + 1.0)
    return EgoTrajectory(pos, cs, quat, tf2_yaw(quat), vel, accel, tfs, force_stop)


def trajectory_from_poses(poses_xycs: np.ndarray, base_position: Tuple[float, float, float] = (0.0, 0.0, 0.0), *,
                          velocity_smoothing_window: int = C.VELOCITY_SMOOTHING_WINDOW,
                          enable_force_stop: bool = True,
                          stopping_threshold: float = C.STOPPING_THRESHOLD) -> EgoTrajectory:
    """``get_trajectory_from_poses`` (``postprocessing_utils.cpp:362-455``) for poses ``[N, 4]`` = denormalised
    (x, y, cos, sin) float32 of one agent, positions at z = ``base_position[2]`` (identity transform)."""
    if not HOST_FAST:
        return trajectory_from_poses_ref(poses_xycs, base_position, velocity_smoothing_window=velocity_smoothing_window,
                                         enable_force_stop=enable_force_stop, stopping_threshold=stopping_threshold)
    p = np.asarray(poses_xycs, np.float32)
    n = p.shape[0]
    dt = C.TRAJECTORY_DT
    pos = np.zeros((n, 3), np.float64)
    pos[:, 0] = p[:, 0].astype(np.float64)
    pos[:, 1] = p[:, 1].astype(np.float64)
    pos[:, 2] = float(base_position[2])
    cs = p[:, 2:4].astype(np.float64).copy()
    quat = quaternion_from_cos_sin(cs[:, 0], cs[:, 1])
    # the node's loops, vectorised with the same float64 / float32 operations in the same order (bit-exact)
    prev = np.asarray(base_position, np.float64)
    px, py, pz = float(prev[0]), float(prev[1]), float(prev[2])
    hyp = math.hypot
    dist = []
    for x, y, z in pos.tolist():            # math.hypot (3 args) kept: np.hypot rounds differently
        dist.append(hyp(x - px, y - py, z - pz))
        px, py, pz = x, y, z
    vel = (np.asarray(dist, np.float64) / dt).astype(np.float32)
    w = int(velocity_smoothing_window)
    if n <= w:
        raise ValueError("velocity_smoothing_window must be smaller than number of points")
    thr = np.float32(stopping_threshold)
    m = n - w + 1
    v64 = vel.astype(np.float64)
    acc = np.zeros(m, np.float64)
    for k in range(w):                      # acc = ((0 + v[i]) + v[i+1]) + ... per window, as the C++ loop
        acc = acc + v64[k:k + m]
    sm = (acc / float(w)).astype(np.float32)  # window i reads only original values (it writes vel[i] last)
    vel[:m] = sm
    force_stop = False
    stop = -1
    if enable_force_stop and m > 1:
        hit = np.flatnonzero((np.abs(sm[:-1]) > thr) & (np.abs(sm[1:]) < thr))
        if hit.size:
            stop, force_stop = int(hit[0]) + 1, True
    if force_stop:                          # from the first stop on: zero velocity, the pose frozen
        vel[stop:] = np.float32(0.0)
        pos[stop:], cs[stop:], quat[stop:] = pos[stop - 1], cs[stop - 1], quat[stop - 1]
    else:
        vel[m:] = vel[m - 1]
    accel = np.zeros(n, np.float32)
    accel[:n - 1] = ((vel[1:].astype(np.float64) - vel[:-1].astype(np.float64)) / dt).astype(np.float32)
    tfs = dt * (np.arange(n, dtype=np.float64) + 1.0)
    return EgoTrajectory(pos, cs, quat, tf2_yaw(quat), vel, accel, tfs, force_stop)


def predicted_paths(denorm: np.ndarray, n_neighbors: int) -> np.ndarray:
    """``create_predicted_objects`` poses for the first ``n_neighbors`` neighbours: ``[n, 80, 5]`` float32 =
    x, y, yaw (``tf2::getYaw``), cos, sin. ``denorm`` is ``[321, 80, 4]`` (agent 0 = ego)."""
    nb = np.asarray(denorm, np.float32)[1:1 + int(n_neighbors)]
    q = quaternion_from_cos_sin(nb[..., 2], nb[..., 3])
    return np.stack([nb[..., 0], nb[..., 1], tf2_yaw(q).astype(np.float32), nb[..., 2], nb[..., 3]],
                    axis=-1).astype(np.float32)


def denoising_steps_ego(iterates: np.ndarray, state_mean: np.ndarray, state_std: np.ndarray) -> np.ndarray:
    """``create_denoising_steps_message`` (batch 1): the ego row of each iterate, denormalised with t = 0 kept:
    ``[steps, 81, 4]``. ``iterates`` is ``[steps, 321, 81, 4]`` or just the ego rows ``[steps, 1, 81, 4]`` (what the
    device plan reads back); the float32 arithmetic is the same element by element (``x * std + mean``)."""
    it = np.asarray(iterates, np.float32)
    if it.ndim != 4 or it.shape[1] not in (1, C.MAX_NUM_AGENTS) or it.shape[2:] != (C.OUTPUT_T + 1, C.POSE_DIM):
        raise ValueError(f"unsupported iterates shape {it.shape}")
    shape = (C.MAX_NUM_AGENTS, 1, C.POSE_DIM)
    mean0 = np.broadcast_to(np.asarray(state_mean, np.float32), shape)[0]
    std0 = np.broadcast_to(np.asarray(state_std, np.float32), shape)[0]
    return (it[:, 0] * std0 + mean0).astype(np.float32)


@dataclass
class TurnDecision:
    command: int
    logits: Tuple[float, ...]
    probabilities: Tuple[float, ...]
    keep_selected: bool
    held: bool = False

    def to_dict(self) -> Dict[str, object]:
        name = C.TURN_INDICATOR_COMMAND_NAMES.get(self.command, str(self.command))
        return {"command": int(self.command), "command_name": name, "keep_selected": bool(self.keep_selected),
                "held": bool(self.held), "logits": [float(v) for v in self.logits],
                "probabilities": [float(v) for v in self.probabilities]}


@dataclass
class TurnIndicatorManager:
    """``TurnIndicatorManager`` (stateful: the last non-KEEP command and its stamp). A fresh manager has no held
    command, which is what a single stateless request sees."""

    hold_duration_s: float = C.TURN_INDICATOR_HOLD_DURATION_S
    keep_offset: float = C.TURN_INDICATOR_KEEP_OFFSET
    _last_command: int = field(default=0, repr=False)
    _last_stamp_s: Optional[float] = field(default=None, repr=False)

    def evaluate(self, logit: np.ndarray, stamp_s: float = 0.0,
                 prev_report: int = C.TURN_INDICATORS_REPORT_DISABLE) -> TurnDecision:
        lg = np.asarray(logit, np.float32).reshape(-1).copy()
        raw = tuple(float(v) for v in lg)
        if lg.size == 0:
            return TurnDecision(1, raw, (), False)  # TurnIndicatorsCommand::DISABLE
        if self._last_stamp_s is not None and self._last_stamp_s > 0 and stamp_s <= self._last_stamp_s + float(
                self.hold_duration_s):
            return TurnDecision(self._last_command, raw, (), False, held=True)
        lg[C.TURN_INDICATOR_OUTPUT_KEEP] = np.float32(lg[C.TURN_INDICATOR_OUTPUT_KEEP] + np.float32(self.keep_offset))
        mx = lg.max()
        prob = np.array([SCALAR_MATH.exp(np.float32(v - mx)) for v in lg], np.float32)
        total = np.float32(0.0001)
        for v in prob:
            total = np.float32(total + v)
        prob = (prob / total).astype(np.float32)
        idx = int(np.argmax(prob))  # std::max_element: the first maximum
        keep = idx == C.TURN_INDICATOR_OUTPUT_KEEP
        command = (int(prev_report) & 0xFF) if keep else idx
        if not keep:
            self._last_command, self._last_stamp_s = command, float(stamp_s)
        return TurnDecision(command, raw, tuple(float(v) for v in prob), keep)