# 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", ]