# SPDX-License-Identifier: Apache-2.0 """Golden tensors of the CPU reference: per-module taps and final outputs of one planning scene. Two products per scene (``code/scripts/ref_golden.py`` writes both): - the **full goldens** (``/.npz``, 10-25 MB, kept OUT of the bundle, by default under ``research/diffusion-planner/goldens``): the raw inputs, the host features, every encoder tap (mixer taps compacted to the valid entities, with their row indices), the encoding, for each of the 11 decoder evaluations its input ``x`` (all 321 agents, for teacher forcing), its time, its block taps and output (valid agents), the solver iterates, ``final_x0``, the turn pool and logits, the post-processed outputs, and the port's rewrite constants (adaLN tables, pad-relative island constants); - the **small goldens** (``tests/goldens/.outputs.npz``, ~0.1 MB, shipped): ``final_x0`` of the valid agents, the logits, the ego trajectory columns, the turn command and the ego row of every solver iterate. Tap names follow ``reference.model.TAP_NAMES``; arrays of valid entities carry ``.rows`` (indices into the full tensor). Device tests compare TT replay outputs against these, never TT against TT. """ from __future__ import annotations import datetime as _dt from pathlib import Path from typing import Any, Dict, Optional import numpy as np from . import config as C from ..host import pipeline as hp from ..host.postprocess import denoising_steps_ego from ..host.solver import solver_plan from ..ttaw.golden import TapRegistry, save_goldens from . import rewrites as R from .pipeline import ReferencePlanner __all__ = ["scene_goldens", "small_goldens", "lite_goldens", "write_scene", "SMALL_KEYS"] SMALL_KEYS = ("final_x0", "final_x0.rows", "logit", "trajectory", "turn_command", "denoising_ego") MIXER_CATS = ("ego", "neighbor", "lane", "route", "polygon", "line_string") def _rows(valid: np.ndarray) -> np.ndarray: return np.flatnonzero(np.asarray(valid, bool)).astype(np.int32) def scene_goldens(ref: ReferencePlanner, raw: Any, *, params: Optional[Dict[str, Any]] = None) -> Dict[str, np.ndarray]: """Run the reference on one scene and return the full golden dict (numpy).""" taps = TapRegistry() res = ref.run(raw, taps=taps, keep_eval_io=True) prep = res.prepared f = prep.features g: Dict[str, np.ndarray] = {} for k, v in prep.raw.items(): g[f"in.{k}"] = v # host features (what the device consumes) for name in ("ego", "neighbor", "neighbor_type", "static", "lane", "lane_attr", "lane_speed", "lane_has_speed", "route", "route_attr", "route_speed", "route_has_speed", "polygon", "line_string", "goal", "ego_shape", "turn", "token_valid", "key_valid", "pos"): g[f"host.{name}"] = np.asarray(getattr(f, name)) for cat, v in f.valid.items(): g[f"host.valid.{cat}"] = np.asarray(v, bool) g["host.agent_valid"] = prep.decoder.agent_valid g["host.current_states"] = prep.decoder.current_states t = taps.to_dict() # encoder taps: mixer internals compacted to the valid entities for cat in MIXER_CATS: rows = _rows(f.valid[cat]) for suffix in ("pre", "mixer"): g[f"enc.{cat}.{suffix}"] = t[f"enc.{cat}.{suffix}"][rows] g[f"enc.{cat}.{suffix}.rows"] = rows for name, _ in C.TOKEN_LAYOUT: g[f"enc.{name}"] = t[f"enc.{name}"] g["enc.tokens"] = t["enc.tokens"] for i in range(C.FUSION_DEPTH): g[f"enc.fusion.{i}"] = t[f"enc.fusion.{i}"] g["enc.encoding"] = res.encoding # decoder evaluations: inputs for teacher forcing, outputs and block taps on the valid agents arows = _rows(prep.decoder.agent_valid) g["dec.rows"] = arows g["dec.t"] = np.asarray(res.eval_times, np.float32) g["dec.x_in"] = np.stack(res.eval_inputs).astype(np.float32) # [11, 321, 81, 4] g["dec.out"] = np.stack(res.eval_outputs)[:, arows].astype(np.float32) for k in range(len(res.eval_times)): g[f"dec.{k}.temb"] = t[f"dec.{k}.temb"][0] g[f"dec.{k}.x"] = t[f"dec.{k}.x"][arows] for i in range(C.DIT_DEPTH): g[f"dec.{k}.block{i}"] = t[f"dec.{k}.block{i}"][arows] g["solver.iterates"] = np.stack(res.denoising_steps).astype(np.float32)[:, arows] g["solver.timesteps"] = np.asarray(res.denoising_timesteps, np.float32) g["final_x0"] = res.final_x0 g["turn.pool"] = t["turn.pool"] g["turn.logit"] = res.logit # post-processed outputs (the API's Trajectory) p = {k: spec[3] for k, spec in hp.RUNTIME_PARAMS.items()} p.update(params or {}) out = hp.make_output(res.final_x0, res.logit, prep, ref.normalization, p, denoising_steps=res.denoising_steps) g["out.trajectory"] = out.poses g["out.predicted_agents"] = out.predicted_agents g["out.turn_command"] = np.asarray(out.turn_indicator["command"], np.int32) g["out.denoising_ego"] = denoising_steps_ego(np.stack(res.denoising_steps), *ref.normalization.state()) # the port's rewrite constants (built in float64, rounded once) plan = solver_plan(C.DPM_SOLVER_STEPS) tab = R.adaln_tables(ref.weights.params, plan.eval_times) g["port.adaln.temb"] = tab.temb for i, blk in enumerate(tab.blocks): for k, v in blk.items(): g[f"port.adaln.block{i}.{k}"] = v g["port.adaln.final_gamma"], g["port.adaln.final_beta"] = tab.final_gamma, tab.final_beta for cat in ("ego", "neighbor"): cst = R.island_constants(ref.weights.params, cat) for k in ("gelu_b1", "c0", "t1_pad", "g_pad", "t2_pad", "w_t1_valid"): g[f"port.island.{cat}.{k}"] = getattr(cst, k) g["port.solver"] = np.asarray([[u.order, u.t, u.a, u.b, u.c, u.r0] for u in plan.updates], np.float64) return g LITE_PREFIXES = ("host.valid.", "host.token_valid", "host.agent_valid", "enc.", "dec.rows", "dec.t", "final_x0", "turn.", "out.", "solver.timesteps") def lite_goldens(g: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: """The compact per-scene goldens of the public-data instants (~1 MB): validity, the encoder category outputs and the encoding on valid rows, final x0 of the valid agents, logits and the post-processed outputs (no inputs: they stay in ``public_data/inputs``; no mixer or per-evaluation decoder taps).""" out: Dict[str, np.ndarray] = {} tok = np.flatnonzero(g["host.token_valid"]) for k, v in g.items(): if not k.startswith(LITE_PREFIXES) or k.startswith(("enc.fusion.", "enc.tokens")): continue if k.endswith((".pre", ".mixer", ".pre.rows", ".mixer.rows")): continue out[k] = v for name, _ in C.TOKEN_LAYOUT: rows = np.flatnonzero(g[f"host.valid.{name}"]) out[f"enc.{name}"], out[f"enc.{name}.rows"] = g[f"enc.{name}"][rows], rows.astype(np.int32) out["enc.encoding"], out["enc.encoding.rows"] = g["enc.encoding"][tok], tok.astype(np.int32) out["final_x0"] = g["final_x0"][g["dec.rows"]] return out def small_goldens(g: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: """The shipped subset of a full golden dict (see the module docstring).""" rows = g["dec.rows"] return {"final_x0": g["final_x0"][rows], "final_x0.rows": rows, "logit": g["turn.logit"], "trajectory": g["out.trajectory"], "turn_command": g["out.turn_command"], "denoising_ego": g["out.denoising_ego"]} def write_scene(ref: ReferencePlanner, raw: Any, scene: str, full_dir: Optional[Path], small_dir: Optional[Path], meta: Optional[Dict[str, Any]] = None, *, lite: bool = False) -> Dict[str, Any]: g = scene_goldens(ref, raw) info = {"scene": scene, "created": _dt.datetime.now(_dt.timezone.utc).isoformat(timespec="seconds"), "weights_sha256": ref.weights.sha256, "solver_steps": C.DPM_SOLVER_STEPS, "valid_counts": {k: int(np.asarray(v).sum()) for k, v in g.items() if k.startswith("host.valid.")}, "producer": "tt_diffusion_planner.reference (fp32 CPU, torch)", **(meta or {})} paths = {} if full_dir is not None: tensors = lite_goldens(g) if lite else g info["lite"] = bool(lite) paths["full"] = str(save_goldens(Path(full_dir) / f"{scene}.npz", tensors, info, compress=True)) if small_dir is not None: paths["small"] = str(save_goldens(Path(small_dir) / f"{scene}.outputs.npz", small_goldens(g), info, compress=True)) return {"info": info, "paths": paths, "goldens": g}