code/models/experimental/gr00t/common/normalization.py

# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

"""Normalization statistics and physical-action decoders of GR00T N1.5 / N1.6 / N1.7 (CPU, numpy + torch).

Sources (every formula cites the report or reference line it was ported from):

* ``host-pipeline-parity.md`` §5 -- decode recipes and the bit-exactness evidence; §4.2 -- state recipes;
  §5 "Correction": N1.7 RELATIVE keys are un-normalised with the raw ``relative_action`` **min/max**, never q01/q99.
* ``arch-n15.md`` §4.3-4.5 -- N1.5 ``experiment_cfg/metadata.json`` schema and the ``FourierGr1ArmsOnlyDataConfig``
  recipe (``Isaac-GR00T-n1.5/gr00t/experiment/data_config.py:150-205``).
* ``Isaac-GR00T{-n1.6,}/gr00t/data/state_action/state_action_processor.py`` (``_compute_normalization_parameters``,
  ``apply_state``, ``unapply_action``) and ``gr00t/data/utils.py`` (``normalize_values_minmax``,
  ``unnormalize_values_minmax``, ``apply_sin_cos_encoding``) -- ported verbatim, including the numpy op order that
  makes the results bit-identical to the goldens.
* ``Isaac-GR00T/gr00t/data/state_action/pose.py:426-456`` (``_rot6d_to_matrix``) and
  ``action_chunking.py:624-658`` (``to_absolute_chunking``: ``T_abs = T_ref @ T_rel``).

Statistics come from the pinned HF snapshots (``configs.snapshot_dir``): N1.5 ``experiment_cfg/metadata.json``;
N1.6 / N1.7 ``processor_config.json`` (modality configs + processor flags) and ``statistics.json``.  No file under
``Isaac-GR00T*`` is imported at runtime.

This module deliberately does not import :mod:`preprocessing` (which imports it); ``decode`` therefore accepts either
an ``Observation`` (duck-typed through its ``.state`` attribute) or a plain ``{group: (T, D)}`` mapping.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path
from typing import Any, Dict, Mapping, Optional, Sequence, Tuple, Union

import numpy as np
import torch

from models.experimental.gr00t.common.configs import GR00T_VERSIONS, get_config, snapshot_dir

# --------------------------------------------------------------------------------------------------------------------
# Specs
# --------------------------------------------------------------------------------------------------------------------


@dataclass(frozen=True)
class StateGroupSpec:
    """One state joint group and how it is normalised before the state encoder.

    ``mode``: ``"sin_cos"`` (``cat[sin(x), cos(x)]``, doubles the width; no statistics), ``"min_max"``
    (``2 (x - lo) / (hi - lo) - 1`` with ``lo/hi`` = min/max or q01/q99, optional clip to [-1, 1]) or ``"mean_std"``.
    """

    key: str
    dim: int  # raw width D
    mode: str
    lo: Optional[np.ndarray] = None  # (D,) min (or q01)  -- min_max only
    hi: Optional[np.ndarray] = None  # (D,) max (or q99)  -- min_max only
    mean: Optional[np.ndarray] = None
    std: Optional[np.ndarray] = None
    clip: bool = False

    @property
    def encoded_dim(self) -> int:
        return 2 * self.dim if self.mode == "sin_cos" else self.dim


@dataclass(frozen=True)
class ActionGroupSpec:
    """One action joint group: representation flags and the min/max used by the decoder.

    ``lo`` / ``hi`` are ``(D,)`` for ABSOLUTE groups and ``(H, D)`` per-step bounds for RELATIVE groups (N1.6 / N1.7
    ``statistics[tag]["relative_action"][key]["min"|"max"]``, ``state_action_processor.py:179-204`` override).
    """

    key: str
    dim: int
    rep: str  # "ABSOLUTE" | "RELATIVE"
    type: str  # "EEF" | "NON_EEF"
    format: str  # "DEFAULT" | "XYZ_ROT6D"
    state_key: str
    lo: np.ndarray
    hi: np.ndarray
    stats_source: str  # "action.min/max" | "action.q01/q99" | "relative_action.min/max"


@dataclass(frozen=True)
class EmbodimentSpec:
    """Everything the host needs to know about one (version, embodiment tag)."""

    version: str
    tag: str
    embodiment_id: int
    video_keys: Tuple[str, ...]
    video_delta_indices: Tuple[int, ...]  # frames per camera = len(...)
    state_delta_indices: Tuple[int, ...]
    state_groups: Tuple[StateGroupSpec, ...]
    action_groups: Tuple[ActionGroupSpec, ...]
    action_horizon: int
    max_state_dim: int
    max_action_dim: int
    # image pipeline knobs (N1.6 / N1.7 processor_config; N1.5 data config constants)
    shortest_image_edge: int
    crop_fraction: float
    letter_box_transform: bool
    formalize_language: bool
    decode_dtype: str  # "torch_fp32" (N1.5) | "numpy_fp64" (N1.6 / N1.7)

    @property
    def frames_per_camera(self) -> int:
        return len(self.video_delta_indices)

    @property
    def n_images(self) -> int:
        return len(self.video_keys) * self.frames_per_camera

    @property
    def state_dim(self) -> int:
        """Encoded state width before zero-padding to ``max_state_dim`` (52 / 58 / 17 for the canonical embodiments)."""
        return sum(g.encoded_dim for g in self.state_groups)

    @property
    def action_dim(self) -> int:
        return sum(g.dim for g in self.action_groups)

    @property
    def state_keys(self) -> Tuple[str, ...]:
        return tuple(g.key for g in self.state_groups)

    @property
    def action_keys(self) -> Tuple[str, ...]:
        return tuple(g.key for g in self.action_groups)


# --------------------------------------------------------------------------------------------------------------------
# Raw file loaders (cached)
# --------------------------------------------------------------------------------------------------------------------


def _read_json(path: Path) -> Any:
    if not path.is_file():
        raise FileNotFoundError(f"{path} does not exist")
    with open(path) as fh:
        return json.load(fh)


@lru_cache(maxsize=None)
def load_n15_metadata() -> Dict[str, Any]:
    """``experiment_cfg/metadata.json`` of the N1.5 snapshot: ``{tag: {statistics, modalities, embodiment_tag}}``
    (arch-n15 §4.5)."""
    return _read_json(snapshot_dir("n15") / "experiment_cfg" / "metadata.json")


@lru_cache(maxsize=None)
def load_processor_kwargs(version: str) -> Dict[str, Any]:
    """``processor_config.json["processor_kwargs"]`` of the N1.6 / N1.7 snapshot."""
    if version not in ("n16", "n17"):
        raise KeyError(f"{version}: processor_config.json exists only for n16 / n17")
    j = _read_json(snapshot_dir(version) / "processor_config.json")
    if "processor_kwargs" not in j:
        raise ValueError(f"{version}: processor_config.json has no 'processor_kwargs'")
    return j["processor_kwargs"]


@lru_cache(maxsize=None)
def load_statistics(version: str) -> Dict[str, Any]:
    """``statistics.json`` of the N1.6 / N1.7 snapshot: ``{tag: {state, action, relative_action}}``."""
    if version not in ("n16", "n17"):
        raise KeyError(f"{version}: statistics.json exists only for n16 / n17")
    return _read_json(snapshot_dir(version) / "statistics.json")


@lru_cache(maxsize=None)
def load_embodiment_id_json(version: str) -> Dict[str, int]:
    """``embodiment_id.json`` of the N1.6 / N1.7 snapshot (tag -> slot id as shipped by NVIDIA)."""
    if version not in ("n16", "n17"):
        raise KeyError(f"{version}: embodiment_id.json exists only for n16 / n17")
    return {str(k): int(v) for k, v in _read_json(snapshot_dir(version) / "embodiment_id.json").items()}


def embodiment_id(version: str, tag: str) -> int:
    """Embodiment tag -> integer slot of the ``[32, in, out]`` encoder / decoder weights (configs ``embodiment_ids``;
    cross-checked against the snapshot's ``embodiment_id.json`` when the tag is listed there)."""
    cfg = get_config(version)
    eid = cfg.embodiment_id(tag)
    if version in ("n16", "n17"):
        shipped = load_embodiment_id_json(version)
        if tag in shipped and shipped[tag] != eid:
            raise ValueError(f"{version}: embodiment_id.json says {tag} -> {shipped[tag]}, configs say {eid}")
    return eid


# --------------------------------------------------------------------------------------------------------------------
# N1.5 data-config recipes (code-defined in the reference, so tabulated here)
# --------------------------------------------------------------------------------------------------------------------

# FourierGr1ArmsOnlyDataConfig (Isaac-GR00T-n1.5/gr00t/experiment/data_config.py:150-205): video ego_view, 4 state keys
# through StateActionSinCosTransform, 4 action keys through StateActionTransform(min_max), 16-step horizon, GR00TTransform
# max_state_dim 64 / max_action_dim 32.  The other shipped tags (oxe_droid: euler->rotation_6d + unnormalised deltas +
# binary gripper; agibot_genie1: min_max on 20/22 dims) have no golden and are intentionally not tabulated: encoding an
# unknown tag raises instead of guessing.
_N15_RECIPES: Mapping[str, Dict[str, Any]] = {
    "gr1": dict(
        video_keys=("ego_view",),
        state_keys=("left_arm", "right_arm", "left_hand", "right_hand"),
        state_mode="sin_cos",
        action_keys=("left_arm", "right_arm", "left_hand", "right_hand"),
        action_mode="min_max",
        action_horizon=16,
        max_state_dim=64,
        max_action_dim=32,
    ),
}


def _n15_spec(tag: str) -> EmbodimentSpec:
    if tag not in _N15_RECIPES:
        raise NotImplementedError(
            f"n15: no data-config recipe tabulated for embodiment {tag!r} (have {sorted(_N15_RECIPES)}); "
            "the N1.5 transforms are code-defined per data config and only gr1 has a golden"
        )
    r = _N15_RECIPES[tag]
    meta = load_n15_metadata()
    if tag not in meta:
        raise KeyError(f"n15: {tag!r} not in experiment_cfg/metadata.json (have {sorted(meta)})")
    stats = meta[tag]["statistics"]
    mods = meta[tag]["modalities"]
    state_groups = []
    for k in r["state_keys"]:
        d = int(mods["state"][k]["shape"][0])
        if r["state_mode"] != "sin_cos":
            raise NotImplementedError(f"n15 state mode {r['state_mode']!r}")
        state_groups.append(StateGroupSpec(key=k, dim=d, mode="sin_cos"))
    action_groups = []
    for k in r["action_keys"]:
        st = stats["action"][k]
        lo = np.asarray(st["min"], dtype=np.float64)
        hi = np.asarray(st["max"], dtype=np.float64)
        d = int(mods["action"][k]["shape"][0])
        if lo.shape != (d,) or hi.shape != (d,):
            raise ValueError(f"n15 {tag}/{k}: stats shape {lo.shape} != modality dim {d}")
        action_groups.append(
            ActionGroupSpec(
                key=k,
                dim=d,
                rep="ABSOLUTE",
                type="NON_EEF",
                format="DEFAULT",
                state_key=k,
                lo=lo,
                hi=hi,
                stats_source="action.min/max",
            )
        )
    return EmbodimentSpec(
        version="n15",
        tag=tag,
        embodiment_id=embodiment_id("n15", tag),
        video_keys=tuple(r["video_keys"]),
        video_delta_indices=(0,),
        state_delta_indices=(0,),
        state_groups=tuple(state_groups),
        action_groups=tuple(action_groups),
        action_horizon=int(r["action_horizon"]),
        max_state_dim=int(r["max_state_dim"]),
        max_action_dim=int(r["max_action_dim"]),
        shortest_image_edge=256,
        crop_fraction=0.95,
        letter_box_transform=False,
        formalize_language=False,
        decode_dtype="torch_fp32",
    )


# --------------------------------------------------------------------------------------------------------------------
# N1.6 / N1.7 specs from processor_config.json + statistics.json
# --------------------------------------------------------------------------------------------------------------------


def _processor_spec(version: str, tag: str) -> EmbodimentSpec:
    kw = load_processor_kwargs(version)
    stats_all = load_statistics(version)
    mcs = kw["modality_configs"]
    if tag not in mcs:
        raise KeyError(f"{version}: embodiment {tag!r} not in processor_config modality_configs (have {sorted(mcs)})")
    if tag not in stats_all:
        raise KeyError(f"{version}: embodiment {tag!r} has no entry in statistics.json")
    mc = mcs[tag]
    stats = stats_all[tag]
    use_percentiles = bool(kw["use_percentiles"])
    # clip_outliers is absent from the N1.7 processor_config -> default True (processing_gr00t_n1d7.py:222).
    clip_outliers = bool(kw["clip_outliers"]) if "clip_outliers" in kw else True
    apply_sincos = bool(kw["apply_sincos_state_encoding"])
    use_relative = bool(kw["use_relative_action"])
    if kw.get("use_mean_std", False):
        raise NotImplementedError(f"{version}: use_mean_std=True is not covered by any golden")

    # ---- state groups (state_action_processor.apply_state strategies 1-3)
    sc = mc["state"]
    sincos_keys = tuple(sc.get("sin_cos_embedding_keys") or ()) if apply_sincos else ()
    meanstd_keys = tuple(sc.get("mean_std_embedding_keys") or ())
    state_groups = []
    for k in sc["modality_keys"]:
        if k not in stats["state"]:
            raise KeyError(f"{version}/{tag}: no state statistics for {k!r}")
        st = stats["state"][k]
        d = len(st["min"])
        if k in sincos_keys:
            state_groups.append(StateGroupSpec(key=k, dim=d, mode="sin_cos"))
        elif k in meanstd_keys:
            state_groups.append(
                StateGroupSpec(
                    key=k,
                    dim=d,
                    mode="mean_std",
                    mean=np.asarray(st["mean"], np.float64),
                    std=np.asarray(st["std"], np.float64),
                )
            )
        else:
            lo_k, hi_k = ("q01", "q99") if use_percentiles else ("min", "max")
            state_groups.append(
                StateGroupSpec(
                    key=k,
                    dim=d,
                    mode="min_max",
                    lo=np.asarray(st[lo_k], np.float64),
                    hi=np.asarray(st[hi_k], np.float64),
                    clip=clip_outliers,
                )
            )

    # ---- action groups (_compute_normalization_parameters + the RELATIVE override)
    ac = mc["action"]
    horizon = len(ac["delta_indices"])
    if horizon > int(kw["max_action_horizon"]):
        raise ValueError(f"{version}/{tag}: horizon {horizon} > max_action_horizon {kw['max_action_horizon']}")
    if ac.get("mean_std_embedding_keys"):
        raise NotImplementedError(f"{version}/{tag}: mean_std action keys are not covered by any golden")
    action_cfgs = ac.get("action_configs")
    if action_cfgs is None:
        action_cfgs = [dict(rep="ABSOLUTE", type="NON_EEF", format="DEFAULT", state_key=None)] * len(
            ac["modality_keys"]
        )
    if len(action_cfgs) != len(ac["modality_keys"]):
        raise ValueError(f"{version}/{tag}: {len(action_cfgs)} action_configs for {len(ac['modality_keys'])} keys")
    action_groups = []
    for k, acfg in zip(ac["modality_keys"], action_cfgs):
        if k not in stats["action"]:
            raise KeyError(f"{version}/{tag}: no action statistics for {k!r}")
        d = len(stats["action"][k]["min"])
        rep, typ, fmt = str(acfg["rep"]), str(acfg["type"]), str(acfg["format"])
        if rep not in ("ABSOLUTE", "RELATIVE") or typ not in ("EEF", "NON_EEF") or fmt not in ("DEFAULT", "XYZ_ROT6D"):
            raise NotImplementedError(f"{version}/{tag}/{k}: action config {acfg}")
        if rep == "RELATIVE" and use_relative:
            rel = stats.get("relative_action", {})
            if k not in rel:
                raise KeyError(f"{version}/{tag}: RELATIVE key {k!r} has no relative_action statistics")
            # host-pipeline-parity §5 correction: the whole relative_action dict replaces the action norm params and
            # unnormalize_values_minmax reads params["min"] / params["max"] -> use_percentiles is ignored here.
            lo = np.asarray(rel[k]["min"], np.float64)
            hi = np.asarray(rel[k]["max"], np.float64)
            if lo.shape != (horizon, d):
                raise ValueError(f"{version}/{tag}/{k}: relative_action min shape {lo.shape} != ({horizon}, {d})")
            src = "relative_action.min/max"
        else:
            lo_k, hi_k = ("q01", "q99") if use_percentiles else ("min", "max")
            lo = np.asarray(stats["action"][k][lo_k], np.float64)
            hi = np.asarray(stats["action"][k][hi_k], np.float64)
            if lo.shape != (d,):
                raise ValueError(f"{version}/{tag}/{k}: action {lo_k} shape {lo.shape} != ({d},)")
            src = f"action.{lo_k}/{hi_k}"
        if typ == "EEF" and fmt != "XYZ_ROT6D":
            raise NotImplementedError(f"{version}/{tag}/{k}: EEF format {fmt}")
        if typ == "EEF" and d != 9:
            raise ValueError(f"{version}/{tag}/{k}: XYZ_ROT6D group must have 9 dims, got {d}")
        action_groups.append(
            ActionGroupSpec(
                key=k,
                dim=d,
                rep=rep,
                type=typ,
                format=fmt,
                state_key=str(acfg.get("state_key") or k),
                lo=lo,
                hi=hi,
                stats_source=src,
            )
        )

    if kw.get("shortest_image_edge") is None or kw.get("crop_fraction") is None:
        raise NotImplementedError(f"{version}: image_crop_size/image_target_size fallback path (no golden)")
    return EmbodimentSpec(
        version=version,
        tag=tag,
        embodiment_id=embodiment_id(version, tag),
        video_keys=tuple(mc["video"]["modality_keys"]),
        video_delta_indices=tuple(int(i) for i in mc["video"]["delta_indices"]),
        state_delta_indices=tuple(int(i) for i in sc["delta_indices"]),
        state_groups=tuple(state_groups),
        action_groups=tuple(action_groups),
        action_horizon=horizon,
        max_state_dim=int(kw["max_state_dim"]),
        max_action_dim=int(kw["max_action_dim"]),
        shortest_image_edge=int(kw["shortest_image_edge"]),
        crop_fraction=float(kw["crop_fraction"]),
        # N1.6's albumentations builder always prepends LetterBoxPad (gr00t_n1d6/image_augmentations.py:398-403);
        # N1.7 only when processor_config says so (gr00t_n1d7/image_augmentations.py:418-419, 480-483).
        letter_box_transform=True if version == "n16" else bool(kw.get("letter_box_transform", False)),
        formalize_language=bool(kw["formalize_language"]),
        decode_dtype="numpy_fp64",
    )


@lru_cache(maxsize=None)
def get_embodiment_spec(version: str, tag: str) -> EmbodimentSpec:
    """Resolve (version, embodiment tag) -> :class:`EmbodimentSpec` from the pinned snapshot files (cached)."""
    if version not in GR00T_VERSIONS:
        raise KeyError(f"unknown GR00T version {version!r}; expected one of {GR00T_VERSIONS}")
    if version == "n15":
        return _n15_spec(tag)
    return _processor_spec(version, tag)


# --------------------------------------------------------------------------------------------------------------------
# Numeric primitives (verbatim ports; op order is load-bearing for bit-exactness)
# --------------------------------------------------------------------------------------------------------------------


def apply_sin_cos_encoding(values: np.ndarray) -> np.ndarray:
    """``[sin(v), cos(v)]`` along the last axis in the input dtype (``gr00t/data/utils.py:25-40``)."""
    return np.concatenate([np.sin(values), np.cos(values)], axis=-1)


def normalize_values_minmax(values: np.ndarray, min_vals: np.ndarray, max_vals: np.ndarray) -> np.ndarray:
    """``gr00t/data/utils.py`` ``normalize_values_minmax``: ``2 (x - lo) / (hi - lo) - 1`` where ``hi != lo``
    (``~np.isclose``), 0 elsewhere.  The first assignment rounds the float64 quotient into ``values.dtype`` before the
    ``2 x - 1`` step is applied in that dtype -- keep both statements as they are (host-pipeline-parity §4.2 N1.7)."""
    normalized = np.zeros_like(values)
    mask = ~np.isclose(max_vals, min_vals)
    normalized[..., mask] = (values[..., mask] - min_vals[..., mask]) / (max_vals[..., mask] - min_vals[..., mask])
    normalized[..., mask] = 2 * normalized[..., mask] - 1
    return normalized


def unnormalize_values_minmax(normalized: np.ndarray, min_vals: np.ndarray, max_vals: np.ndarray) -> np.ndarray:
    """``gr00t/data/utils.py`` ``unnormalize_values_minmax``: clip to [-1, 1], then ``(y + 1) / 2 * (hi - lo) + lo``
    (float64 because the bounds are float64)."""
    range_vals = max_vals - min_vals
    return (np.clip(normalized, -1.0, 1.0) + 1.0) / 2.0 * range_vals + min_vals


def normalize_values_meanstd(values: np.ndarray, mean: np.ndarray, std: np.ndarray) -> np.ndarray:
    """``gr00t/data/utils.py`` ``normalize_values_meanstd``: ``(x - mean) / std`` where ``std != 0``, ``x`` elsewhere."""
    mask = std != 0
    out = np.zeros_like(values)
    out[..., mask] = (values[..., mask] - mean[..., mask]) / std[..., mask]
    out[..., ~mask] = values[..., ~mask]
    return out


def rot6d_to_matrix(rot6d: np.ndarray) -> np.ndarray:
    """Gram-Schmidt 6-D -> 3x3 rotation (``pose.py:426-456`` ``_rot6d_to_matrix``): rows ``a = r[0]/|r[0]|``,
    ``b = (r[1] - (a.b) a) / |...|``, ``c = a x b``."""
    r = np.asarray(rot6d, dtype=np.float64).reshape(2, 3)
    row1 = r[0] / np.linalg.norm(r[0])
    row2 = r[1] - np.dot(row1, r[1]) * row1
    row2 = row2 / np.linalg.norm(row2)
    row3 = np.cross(row1, row2)
    return np.vstack([row1, row2, row3])


def xyz_rot6d_to_homogeneous(xyz_rot6d: np.ndarray) -> np.ndarray:
    """``EndEffectorPose(translation=v[:3], rotation=v[3:], rotation_type="rot6d").homogeneous`` without the scipy
    ``Rotation`` round trip (which perturbs the matrix by ~1e-16, host-pipeline-parity §5)."""
    v = np.asarray(xyz_rot6d, dtype=np.float64)
    if v.shape != (9,):
        raise ValueError(f"xyz_rot6d must have shape (9,), got {v.shape}")
    h = np.eye(4, dtype=np.float64)
    h[:3, :3] = rot6d_to_matrix(v[3:9])
    h[:3, 3] = v[:3]
    return h


def eef_relative_to_absolute(relative: np.ndarray, reference_state: np.ndarray) -> np.ndarray:
    """``EndEffectorActionChunk.to_absolute_chunking`` + ``to(XYZ_ROT6D)`` (``action_chunking.py:624-658``):
    ``T_abs = T_ref @ T_rel`` per step, back to ``[t, R[0], R[1]]``.  ``relative`` (H, 9), ``reference_state`` (9,)."""
    rel = np.asarray(relative, dtype=np.float64)
    if rel.ndim != 2 or rel.shape[1] != 9:
        raise ValueError(f"relative EEF chunk must be (H, 9), got {rel.shape}")
    t_ref = xyz_rot6d_to_homogeneous(reference_state)
    out = np.empty_like(rel)
    for i in range(rel.shape[0]):
        t_abs = t_ref @ xyz_rot6d_to_homogeneous(rel[i])
        out[i, :3] = t_abs[:3, 3]
        out[i, 3:9] = t_abs[:3, :3][:2, :].reshape(-1)
    return out


# --------------------------------------------------------------------------------------------------------------------
# State encoding (used by preprocessing.encode)
# --------------------------------------------------------------------------------------------------------------------


def _state_array(spec: EmbodimentSpec, state: Mapping[str, np.ndarray], key: str, dim: int) -> np.ndarray:
    if key not in state:
        raise KeyError(f"{spec.version}/{spec.tag}: state group {key!r} missing (have {sorted(state)})")
    v = np.asarray(state[key])
    if v.ndim == 1:
        v = v[None, :]
    t_s = len(spec.state_delta_indices)
    if v.ndim != 2 or v.shape != (t_s, dim):
        raise ValueError(f"{spec.version}/{spec.tag}: state {key!r} must be ({t_s}, {dim}), got {v.shape}")
    if v.dtype.kind != "f":
        raise TypeError(f"{spec.version}/{spec.tag}: state {key!r} must be floating point, got {v.dtype}")
    return v


def encode_state(spec: EmbodimentSpec, state: Mapping[str, np.ndarray]) -> Tuple[torch.Tensor, torch.Tensor]:
    """Raw ``{group: (T_s, D)}`` -> ``(state (1, T_s, max_state_dim) fp32, mask (1, T_s, max_state_dim) bool)``.

    * N1.5 (``StateActionSinCosTransform`` -> ``ConcatTransform`` -> ``GR00TTransform._prepare_state``,
      ``transforms.py:240-270``): torch ``cat[sin, cos]`` per key in the input dtype (float64 in the dataset),
      concatenated, zero-padded to 64, cast to fp32 at the end (bit-exact after the cast, host-pipeline-parity §4.1).
    * N1.6 / N1.7 (``StateActionProcessor.apply_state`` + ``processing_gr00t_n1d{6,7}`` padding): numpy per group in
      the input dtype (sin/cos, or min-max with ``~isclose`` mask + clip), concatenated, zero-padded, fp32.
    The mask marks the real dims (N1.5 ``state_mask``; unused by N1.6 / N1.7 but returned for symmetry).
    """
    parts = []
    for g in spec.state_groups:
        v = _state_array(spec, state, g.key, g.dim)
        if spec.version == "n15":
            t = torch.from_numpy(np.ascontiguousarray(v))
            parts.append(torch.cat([torch.sin(t), torch.cos(t)], dim=-1))
            continue
        if g.mode == "sin_cos":
            parts.append(apply_sin_cos_encoding(v))
        elif g.mode == "min_max":
            n = normalize_values_minmax(v, g.lo, g.hi)
            if g.clip:
                n = np.clip(n, -1.0, 1.0)
            parts.append(n)
        elif g.mode == "mean_std":
            parts.append(normalize_values_meanstd(v, g.mean, g.std))
        else:
            raise ValueError(f"unknown state mode {g.mode!r}")
    if spec.version == "n15":
        cat = torch.cat(parts, dim=-1)  # (T_s, D_real) in the input dtype
    else:
        cat = torch.from_numpy(np.concatenate(parts, axis=-1))
    d_real = cat.shape[-1]
    if d_real != spec.state_dim:
        raise ValueError(f"{spec.version}/{spec.tag}: encoded state width {d_real} != spec {spec.state_dim}")
    if d_real > spec.max_state_dim:
        raise ValueError(f"{spec.version}/{spec.tag}: state width {d_real} exceeds max_state_dim {spec.max_state_dim}")
    padded = torch.zeros(cat.shape[0], spec.max_state_dim, dtype=cat.dtype)
    padded[:, :d_real] = cat
    mask = torch.zeros(1, cat.shape[0], spec.max_state_dim, dtype=torch.bool)
    mask[..., :d_real] = True
    return padded.to(torch.float32)[None], mask


# --------------------------------------------------------------------------------------------------------------------
# Decoders
# --------------------------------------------------------------------------------------------------------------------


def _as_state_mapping(obs: Any) -> Mapping[str, np.ndarray]:
    state = getattr(obs, "state", obs)
    if not isinstance(state, Mapping):
        raise TypeError(f"decode expects an Observation or a {{group: array}} mapping, got {type(obs).__name__}")
    return state


def _check_pred(spec: EmbodimentSpec, pred: torch.Tensor) -> torch.Tensor:
    if not isinstance(pred, torch.Tensor):
        raise TypeError(f"action_pred_normalized must be a torch.Tensor, got {type(pred).__name__}")
    if pred.ndim != 3 or pred.shape[0] != 1:
        raise ValueError(f"action_pred_normalized must be (1, H_max, A_max), got {tuple(pred.shape)}")
    if pred.shape[1] < spec.action_horizon or pred.shape[2] < spec.action_dim:
        raise ValueError(
            f"{spec.version}/{spec.tag}: prediction {tuple(pred.shape)} smaller than horizon {spec.action_horizon} x "
            f"dims {spec.action_dim}"
        )
    if pred.shape[2] != spec.max_action_dim:
        raise ValueError(
            f"{spec.version}/{spec.tag}: expected {spec.max_action_dim} action columns, got {pred.shape[2]}"
        )
    return pred.detach().cpu().to(torch.float32)


def n15_decode(pred: torch.Tensor, spec: EmbodimentSpec) -> Dict[str, torch.Tensor]:
    """N1.5 ``ConcatTransform.unapply`` + ``StateActionTransform.unapply`` (``state_action.py:205-208``): per key in
    concat order, ``(x + 1) / 2 * (max - min) + min`` in **torch fp32** with ``torch.tensor(list)`` statistics, no
    clip.  Bit-exact vs the goldens (host-pipeline-parity §5; a float64 numpy variant is 3.7e-7 off)."""
    pred = _check_pred(spec, pred)
    out: Dict[str, torch.Tensor] = {}
    s = 0
    for g in spec.action_groups:
        x = pred[..., : spec.action_horizon, s : s + g.dim]
        s += g.dim
        mn = torch.tensor(g.lo.tolist(), dtype=x.dtype)
        mx = torch.tensor(g.hi.tolist(), dtype=x.dtype)
        out[g.key] = ((x + 1) / 2 * (mx - mn) + mn).contiguous()
    return out


def _processor_decode(
    pred: torch.Tensor, spec: EmbodimentSpec, state: Mapping[str, np.ndarray], return_relative: bool = False
) -> Dict[str, torch.Tensor]:
    """N1.6 / N1.7 ``Gr00tN1d{6,7}Processor.decode_action`` -> ``StateActionProcessor.unapply_action``: slice the
    padded prediction by horizon / dims, clip + per-step (RELATIVE) or per-dim (ABSOLUTE) min-max un-normalisation in
    float64, then RELATIVE -> absolute: NON_EEF ``+ state[key][-1]``, EEF ``T_state @ T_rel`` with Gram-Schmidt rot6d
    (``state_action_processor.py:421-525``).  Output fp32 ``(1, H, D)`` per key (the goldens are saved as fp32)."""
    pred = _check_pred(spec, pred).numpy()
    unnorm: Dict[str, np.ndarray] = {}
    s = 0
    for g in spec.action_groups:
        y = pred[..., : spec.action_horizon, s : s + g.dim]  # (1, H, D) fp32
        s += g.dim
        unnorm[g.key] = unnormalize_values_minmax(y, g.lo, g.hi)  # broadcast (H, D) or (D,) bounds -> float64
    if return_relative:
        return {k: torch.from_numpy(np.ascontiguousarray(v.astype(np.float32))) for k, v in unnorm.items()}
    out: Dict[str, torch.Tensor] = {}
    for g in spec.action_groups:
        val = unnorm[g.key]
        if g.rep == "RELATIVE":
            if g.state_key not in state:
                raise KeyError(f"{spec.version}/{spec.tag}: reference state {g.state_key!r} missing for {g.key!r}")
            ref = np.asarray(state[g.state_key])
            if ref.ndim == 1:
                ref = ref[None, :]
            if ref.ndim != 2 or ref.shape[1] != g.dim:
                raise ValueError(
                    f"{spec.version}/{spec.tag}: state {g.state_key!r} must be (T_s, {g.dim}), got {ref.shape}"
                )
            ref_last = ref[-1]  # "use last timestep of state as reference"
            if g.type == "EEF":
                val = np.stack([eef_relative_to_absolute(val[b], ref_last) for b in range(val.shape[0])])
            else:
                val = val + ref_last
        out[g.key] = torch.from_numpy(np.ascontiguousarray(val.astype(np.float32)))
    return out


def n16_decode(pred: torch.Tensor, spec: EmbodimentSpec, state: Mapping[str, np.ndarray]) -> Dict[str, torch.Tensor]:
    """N1.6: ``(1, 50, 128)`` -> ``{left_arm (1,16,7), right_arm, left_hand (1,16,6), right_hand, waist (1,16,3)}``;
    RELATIVE keys use the per-step ``relative_action`` min/max and add ``state[key][-1]``, ``waist`` (ABSOLUTE) the
    ``action`` min/max (``use_percentiles=False``).  Exact vs the goldens (host-pipeline-parity §5)."""
    if spec.version != "n16":
        raise ValueError(f"n16_decode called with a {spec.version} spec")
    return _processor_decode(pred, spec, state)


def n17_decode(pred: torch.Tensor, spec: EmbodimentSpec, state: Mapping[str, np.ndarray]) -> Dict[str, torch.Tensor]:
    """N1.7: ``(1, 40, 132)`` -> per-key ``(1, 40, D)``; RELATIVE keys: raw ``relative_action`` **min/max** (not
    q01/q99), EEF groups composed as ``T_state @ T_rel`` (rot6d Gram-Schmidt), joints ``+ state[-1]``; ABSOLUTE keys
    (gripper, G1 hands / waist / commands): ``action`` q01/q99.  Joints / gripper exact, eef within 1 fp32 ulp
    (6e-8) of the goldens (host-pipeline-parity §5)."""
    if spec.version != "n17":
        raise ValueError(f"n17_decode called with a {spec.version} spec")
    return _processor_decode(pred, spec, state)


def decode(
    version: str, action_pred_normalized: torch.Tensor, obs: Any, embodiment: Optional[str] = None
) -> Dict[str, torch.Tensor]:
    """Normalised model output ``(1, H_max, A_max)`` -> ``{action key: physical action (1, H, D) fp32}``.

    ``obs`` is the :class:`~preprocessing.Observation` the prediction was made for (its raw ``state`` is the reference
    frame of the RELATIVE keys) or a plain ``{group: (T_s, D)}`` mapping; ``embodiment`` defaults to
    ``obs.embodiment_tag``.  Dispatches to :func:`n15_decode` / :func:`n16_decode` / :func:`n17_decode`.
    """
    tag = embodiment if embodiment is not None else getattr(obs, "embodiment_tag", None)
    if tag is None:
        raise ValueError("decode: embodiment tag missing (pass embodiment= or an Observation with embodiment_tag)")
    spec = get_embodiment_spec(version, tag)
    if version == "n15":
        return n15_decode(action_pred_normalized, spec)
    state = _as_state_mapping(obs)
    if version == "n16":
        return n16_decode(action_pred_normalized, spec, state)
    return n17_decode(action_pred_normalized, spec, state)


def decode_relative(version: str, action_pred_normalized: torch.Tensor, embodiment: str) -> Dict[str, torch.Tensor]:
    """Step 1 only (clip + un-normalise, no state added): the N1.6 golden ``final_action_unnormalized_relative_or_abs``
    tap and the N1.7 relative chunk before SE(3) composition.  Not defined for N1.5 (no relative actions)."""
    if version == "n15":
        raise ValueError("decode_relative: N1.5 actions are absolute; use decode()")
    spec = get_embodiment_spec(version, embodiment)
    return _processor_decode(action_pred_normalized, spec, {}, return_relative=True)


__all__ = [
    "ActionGroupSpec",
    "EmbodimentSpec",
    "StateGroupSpec",
    "apply_sin_cos_encoding",
    "decode",
    "decode_relative",
    "eef_relative_to_absolute",
    "embodiment_id",
    "encode_state",
    "get_embodiment_spec",
    "load_embodiment_id_json",
    "load_n15_metadata",
    "load_processor_kwargs",
    "load_statistics",
    "n15_decode",
    "n16_decode",
    "n17_decode",
    "normalize_values_meanstd",
    "normalize_values_minmax",
    "rot6d_to_matrix",
    "unnormalize_values_minmax",
    "xyz_rot6d_to_homogeneous",
]