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

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

"""Architectural constants of GR00T N1.5 / N1.6 / N1.7 as frozen dataclasses.

Every number here is copied from the research reports under ``docs/research/`` (the section is cited next to the
field or in the class docstring) or from the checkpoints' ``config.json``; nothing is derived from live model code.
Conventions follow ``01-architecture-synthesis.md``: batch 1, ``S`` = LLM sequence length, ``H`` = action horizon,
``V`` = cameras, ``n_img`` = images per sample.

Usage::

    from models.experimental.gr00t.common.configs import get_config, GR00T_VERSIONS, snapshot_dir
    cfg = get_config("n17")
    cfg.dit.n_blocks            # 32
    cfg.canonical_layout.s_pad_llm   # 512
    snapshot_dir("n17")         # ~/.cache/huggingface/hub/models--nvidia--GR00T-N1.7-3B/snapshots/2fc962b9...
"""

from __future__ import annotations

import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, Mapping, Optional, Tuple

GR00T_VERSIONS: Tuple[str, ...] = ("n15", "n16", "n17")

# --------------------------------------------------------------------------------------------------------------------
# Checkpoint identity (weight-loader-spec.md header table + §1.1; 01-architecture-synthesis.md §1.9)
# --------------------------------------------------------------------------------------------------------------------

HF_REPOS: Mapping[str, str] = {
    "n15": "nvidia/GR00T-N1.5-3B",
    "n16": "nvidia/GR00T-N1.6-3B",
    "n17": "nvidia/GR00T-N1.7-3B",
}
HF_SNAPSHOT_SHAS: Mapping[str, str] = {
    "n15": "869830fc749c35f34771aa5209f923ac57e4564e",
    "n16": "d0814e7ecb19202e7c8468b46098b0b7ef3a6d61",
    "n17": "2fc962b973bccdd5d8ce4f67cc63b264d6886495",
}


def hf_hub_cache() -> Path:
    """Directory holding the ``models--*`` snapshot trees (``HF_HUB_CACHE`` > ``HF_HOME/hub`` > ``~/.cache/huggingface/hub``)."""
    if "HF_HUB_CACHE" in os.environ:
        return Path(os.environ["HF_HUB_CACHE"])
    if "HF_HOME" in os.environ:
        return Path(os.environ["HF_HOME"]) / "hub"
    return Path.home() / ".cache" / "huggingface" / "hub"


def snapshot_dir(version: str, hub_cache: Optional[Path] = None) -> Path:
    """Resolve the local HF snapshot directory of ``version`` and verify that it exists (no fallback).

    The pinned revision is the one every golden dump and every research report was produced from
    (``weight-loader-spec.md`` header table; ``golden-tap-map-and-thresholds.md`` §1 provenance row).
    """
    _check_version(version)
    hub = Path(hub_cache) if hub_cache is not None else hf_hub_cache()
    repo_dir = "models--" + HF_REPOS[version].replace("/", "--")
    path = hub / repo_dir / "snapshots" / HF_SNAPSHOT_SHAS[version]
    if not (path / "model.safetensors.index.json").is_file():
        raise FileNotFoundError(
            f"{version}: snapshot {HF_REPOS[version]}@{HF_SNAPSHOT_SHAS[version][:8]} not found at {path} "
            "(expected model.safetensors.index.json)"
        )
    return path


def _check_version(version: str) -> None:
    if version not in GR00T_VERSIONS:
        raise KeyError(f"unknown GR00T version {version!r}; expected one of {GR00T_VERSIONS}")


# --------------------------------------------------------------------------------------------------------------------
# Component configs
# --------------------------------------------------------------------------------------------------------------------


@dataclass(frozen=True)
class CheckpointConfig:
    """Checkpoint facts (weight-loader-spec.md header table + §1.1; 01-architecture-synthesis.md §1.9)."""

    hf_repo: str
    hf_sha: str
    n_shards: int
    n_tensors: int  # keys in model.safetensors.index.json
    n_params: int
    n_bytes: int  # == metadata.total_size; all tensors bf16
    shard_files: Tuple[str, ...]
    # The unit (LLM layer / DiT block) whose tensors are split across two shards -> shards must be resolved per key
    # (weight-loader-spec.md §1.1 consequence (b)).
    straddling_prefix: str
    # Canonical checkpoint key prefixes (weight-loader-spec.md §2 "strip prefix" column).
    vision_prefix: str
    projector_prefix: Optional[str]
    llm_prefix: str  # prefix of `embed_tokens.weight`, `layers.N.*`, `norm.weight`
    lm_head_key: str
    action_head_prefix: str = "action_head."
    dit_prefix: str = "action_head.model."
    vlln_prefix: str = "action_head.vlln."
    vl_self_attention_prefix: Optional[str] = None
    embodiment_prefixes: Tuple[str, ...] = (
        "action_head.state_encoder.",
        "action_head.action_encoder.",
        "action_head.action_decoder.",
    )


@dataclass(frozen=True)
class VisionConfig:
    """Vision encoder (01-architecture-synthesis.md §1.1; arch-n15 §2.1, arch-n16 §1/§4, arch-n17 §2.1-2.2)."""

    family: str  # "siglip_v1" | "siglip2_naflex" | "qwen3vl"
    n_layers: int
    hidden: int
    n_heads: int
    head_dim: int
    mlp_dim: int
    activation: str  # "gelu_tanh" (gelu_pytorch_tanh) in all three towers
    ln_eps: float
    attention_bias: bool  # q/k/v/out Linear carry biases in all three
    qkv_fused_in_checkpoint: bool  # N1.7 stores one `attn.qkv` Linear(1024->3072); SigLIP stores q/k/v separately
    patch_size: int
    temporal_patch: int  # Conv3d temporal kernel (N1.7 = 2, still frame duplicated to t=2)
    patch_in_features: int  # 3*14*14 = 588 (SigLIP) / 3*2*16*16 = 1536 (Qwen3-VL)
    patch_column_order: str  # "(kh,kw,c)" for SigLIP (== pi0 unfold), "(c,t,ph,pw)" for Qwen3-VL (host-layout §1-3)
    pos_table_rows: int  # learned position table rows: 256 (16x16) / 256 / 2304 (48x48)
    pos_grid: int  # side of the learned grid
    pos_interpolation: str  # "none" | "bilinear_16_to_18" | "bilinear_4tap_per_grid"
    image_size: Optional[int]  # fixed square input side; None = dynamic (multiple of 32)
    patches_per_image: Optional[int]  # None = depends on grid
    tokens_per_image: Optional[int]  # LLM tokens per image after merge; None = depends on grid
    merge_size: int  # 1 (none) / 2 (2x2 pixel-unshuffle or merger)
    frames_per_camera: int  # 1 / 1 / 2 (history + current)
    rope: str  # "none" | "2d"
    rope_theta: Optional[float]
    rope_freqs_per_axis: Optional[int]  # N1.7: 16 frequencies per (h, w) axis -> cos/sin [P, 64]
    post_layernorm: bool  # SigLIP `post_layernorm`; Qwen3-VL has the LN inside the merger
    deepstack_layers: Tuple[int, ...]  # blocks whose outputs feed the deepstack mergers (N1.7: 5, 11, 17)
    full_attention_per_image: bool = True  # N1.6 windows are disabled (critique C1, golden-tap-map §5)
    head_dim_pad: int = 0  # ttnn SDPA head-dim padding (static-shape-plan §3.1): 72->96, 64 unchanged
    patch_in_features_pad: int = 0  # 588 -> 608 tile padding (static-shape-plan §3.1)
    dynamic_min_pixels: Optional[int] = None  # N1.7 smart_resize: >= 65,536 px, factor 32
    dynamic_factor: Optional[int] = None
    executed_params: int = 0  # perf-model-v2 §2.2 (pooling head excluded)
    executed_bytes_bf16: int = 0

    @property
    def attention_scale(self) -> float:
        """Softmax scale; must be passed explicitly when heads are padded (static-shape-plan D5)."""
        return self.head_dim**-0.5


@dataclass(frozen=True)
class ProjectorConfig:
    """Vision -> LLM projector (01-architecture-synthesis.md §1.2; arch-n15 §2.2, arch-n16 §4, arch-n17 §2.2)."""

    kind: str  # "linear" | "unshuffle_ln_mlp" | "qwen3vl_merger"
    in_dim: int  # width entering the first Linear (1152 / 4608 / 4096)
    hidden_dim: Optional[int]  # None for the single Linear
    out_dim: int  # 2048 = LLM hidden
    ln_dim: Optional[int]  # LayerNorm width (None = no LN): N1.6 4608 (post-concat); N1.7 1024 (pre-concat)
    ln_eps: Optional[float]
    ln_position: Optional[str]  # "post_concat" | "pre_concat"
    activation: Optional[str]  # "gelu_erf" (nn.GELU()) or None -- never tanh here (synthesis §3 N9)
    merge_channel_order: Optional[str]  # N1.6 "c*4+dy*2+dx" (F.pixel_unshuffle); N1.7 "(mh,mw,c)" row-major concat
    image_token_id: int  # slot id overwritten by the projector rows (151669 <IMG_CONTEXT> / 151655 <|image_pad|>)
    splice: str  # "overwrite" (N1.5/N1.6 indexed assignment) | "masked_scatter" (N1.7)
    n_deepstack: int = 0  # N1.7: 3 extra mergers (post-concat LN(4096)) added to the residual after LLM layers 0..2
    deepstack_ln_dim: Optional[int] = None
    params: int = 0
    params_deepstack_each: int = 0


@dataclass(frozen=True)
class LLMConfig:
    """Truncated Qwen3 decoder (01-architecture-synthesis.md §1.3; arch-n15 §2.4, arch-n16 §5, arch-n17 §1.4/§2.3)."""

    family: str  # "qwen3_1p7b" | "qwen3vl_2b_text"
    n_layers_total: int  # 28 in the source model
    n_layers_executed: int  # `select_layer`: 12 / 16 / 16
    hidden: int
    n_heads: int
    n_kv_heads: int
    head_dim: int
    intermediate: int  # SwiGLU width
    rms_eps: float
    qk_norm: bool  # per-head RMSNorm(128) on q and k before RoPE
    attention_bias: bool
    rope_mode: str  # "1d" | "mrope_interleaved"
    rope_theta: float
    mrope_section: Optional[Tuple[int, int, int]]  # N1.7 (24, 20, 20): T owns 24 half-dims, H 20, W 20
    rope_layout: str  # "rotate_half" (HF) in all three
    vocab: int  # embed_tokens rows
    tie_word_embeddings: bool
    final_norm_applied: bool  # backbone feature = post final RMSNorm (N1.5/N1.6) or raw layer-15 residual (N1.7)
    deepstack_add_after_layers: Tuple[int, ...]  # N1.7: (0, 1, 2)
    causal: bool = True
    massive_activation_channels: Tuple[int, ...] = ()  # residual channels reaching ~1e4 at token 0 (ref-env §5.4)
    executed_layer_params: int = 0  # perf-model-v2 §2.2
    executed_layer_bytes_bf16: int = 0


@dataclass(frozen=True)
class VLAdapterConfig:
    """`vlln` + optional 4-block VL self-attention (01-architecture-synthesis.md §1.4; arch-n15 §3.5, arch-n17 §3.2)."""

    dim: int  # 2048
    vlln_eps: float  # LayerNorm(2048, eps 1e-5, affine)
    n_blocks: int  # 4 / 0 / 4
    n_heads: int
    head_dim: int
    ff_dim: int
    ln_eps: float
    activation: str  # "gelu_tanh"
    attention_bias: bool
    masked: bool  # reference applies no key mask at B=1
    reference_dtype: str  # "fp32" (N1.5 autocast promotes LN, residual fp32) | "bf16"
    params: int = 0


@dataclass(frozen=True)
class DiTConfig:
    """AdaLN cross/self-interleaved DiT (01-architecture-synthesis.md §1.5; critique C6; weight-loader §4.6)."""

    class_name: str  # "cross_attention_dit.DiT" | "dit.AlternateVLDiT"
    n_blocks: int
    width: int
    n_heads: int
    head_dim: int
    ff_dim: int
    activation: str  # "gelu_tanh"
    attention_bias: bool
    cross_attention_dim: int  # 2048 (to_k/to_v: 2048 -> 1536)
    cross_block_parity: str  # "even" blocks cross-attend, odd self-attend (C S C S ...)
    cross_key_set: str  # "all" (N1.5, no mask) | "alternate_text_image" (idx%4==0 -> non-image, idx%4==2 -> image)
    adaln_chunks: int  # 2
    adaln_order: Tuple[str, str]  # ("scale", "shift") for norm1.linear -- attention branch only, ungated
    adaln_ln_eps: float  # 1e-5, non-affine
    norm3_affine: bool  # False: plain LN(1536, eps 1e-5) before the FF, no weights in any checkpoint (synthesis §3 N8)
    norm3_eps: float
    out_norm_eps: float  # 1e-6, non-affine
    proj_out_1_order: Tuple[str, str]  # ("shift", "scale") -- opposite order to the blocks
    out_dim: int  # proj_out_2: 1536 -> 1024
    timestep_embed_dim: int  # Timesteps(256, flip_sin_to_cos=True, downscale_freq_shift=1)
    timestep_flip_sin_to_cos: bool
    timestep_downscale_freq_shift: int
    timestep_max_period: float
    timestep_mlp_hidden: int  # Linear(256->1536) -> SiLU -> Linear(1536->1536); `temb` shared by all blocks
    n_query_tokens: int  # 49 = state 1 + future 32 + action 16 | 51 | 41
    n_future_tokens: int  # 32 (N1.5 learned FLARE tokens) | 0
    head_dim_pad: int = 64  # 48 -> 64 at load (static-shape-plan D5)
    m_pad: int = 64  # query rows padded to 64 (static-shape-plan D4)
    params_blocks: int = 0
    params_timestep_encoder: int = 0
    params_proj_out: int = 0

    @property
    def attention_scale(self) -> float:
        """48**-0.5; must be passed explicitly after padding to 64 (ttnn-blocks-dit §3.1)."""
        return self.head_dim**-0.5

    @property
    def cross_block_indices(self) -> Tuple[int, ...]:
        return tuple(range(0, self.n_blocks, 2))

    @property
    def self_block_indices(self) -> Tuple[int, ...]:
        return tuple(range(1, self.n_blocks, 2))

    def cross_key_subset(self, block_idx: int) -> str:
        """Key set of cross block ``block_idx``: "all" | "text" (non-image) | "image" (arch-n16 §7.4, arch-n17 §3.3)."""
        if block_idx % 2 != 0:
            raise ValueError(f"block {block_idx} is a self-attention block")
        if self.cross_key_set == "all":
            return "all"
        return "text" if block_idx % 4 == 0 else "image"

    @property
    def velocity_rows(self) -> slice:
        """Rows of the decoder output that form the velocity: the last H rows (synthesis §1.6 'action decoder')."""
        return slice(self.n_query_tokens - (self.n_query_tokens - 1 - self.n_future_tokens), None)


@dataclass(frozen=True)
class EncoderConfig:
    """Embodiment-conditioned MLPs (01-architecture-synthesis.md §1.6; weight-loader-spec.md §4.7, arch-n1x §3.1)."""

    max_num_embodiments: int  # 32 -> W [32, in, out], b [32, out]; y = x @ W[e] + b[e] ([in,out], no transpose)
    max_state_dim: int
    max_action_dim: int
    action_horizon: int
    hidden_size: int  # 1024 (state encoder hidden, decoder hidden, DiT output width)
    input_embedding_dim: int  # 1536 (DiT width)
    state_encoder_dims: Tuple[int, int, int]  # (max_state_dim, 1024, 1536); ReLU between
    action_encoder_w1: Tuple[int, int]  # (max_action_dim, 1536)
    action_encoder_w2: Tuple[int, int]  # (3072, 1536) on cat[a, tau]; swish after
    action_encoder_w3: Tuple[int, int]  # (1536, 1536)
    tau_dim: int  # 1536 sinusoid: half 768, freqs f_i = exp(-i ln(1e4)/768), enc = [sin(t f), cos(t f)] (arch-n15 §3.3)
    tau_order: Tuple[str, str]
    tau_max_period: float
    action_decoder_dims: Tuple[int, int, int]  # (1024, 1024, max_action_dim); ReLU between
    position_embedding_rows: int  # Embedding(1024, 1536); rows 0..H-1 added to the action tokens
    state_dim_pad: int  # tile-friendly padding of the state input (weight-loader §4.7): 64 / 128 / 160
    action_dim_pad: int  # 32 / 128 / 160
    used_slot_bytes_bf16: int = 0  # perf-model-v2 §2.2


@dataclass(frozen=True)
class SamplerConfig:
    """4-step Euler flow matching (01-architecture-synthesis.md §1.6; arch-n16 §7.6; golden-tap-map §2.1)."""

    num_inference_timesteps: int  # 4
    num_timestep_buckets: int  # 1000
    dt: float  # 0.25
    timestep_buckets: Tuple[int, ...]  # int(k/4 * 1000) = 0, 250, 500, 750
    noise_shape: Tuple[int, int, int]  # [1, H, max_action_dim]; all dims noisy incl. pads
    euler_dtype: str  # "fp32_actions_bf16_velocity" (N1.5) | "bf16" (N1.6/N1.7)
    noise_dtype: str  # RNG dtype in the reference: "fp32" | "bf16"
    velocity_rows_from_end: int  # decoder rows -H:

    def timestep_bucket(self, k: int) -> int:
        """Discretised timestep fed to the encoders: ``int(k / num_inference_timesteps * num_timestep_buckets)``."""
        if not 0 <= k < self.num_inference_timesteps:
            raise ValueError(f"step {k} out of range")
        return int(k / self.num_inference_timesteps * self.num_timestep_buckets)


@dataclass(frozen=True)
class PreprocessingConfig:
    """Host-side image/state preprocessing constants (01-architecture-synthesis.md §1.7; host-pipeline-parity §4.2)."""

    image_pipeline: str  # short recipe name, see host-pipeline-parity §4.2
    crop_fraction: float  # 0.95 in all three
    shortest_edge: int  # 256 source frames
    crop_size: int  # int(256 * 0.95) = 243
    resize_interpolation: str  # "torchvision_bilinear_antialias" | "pil_bicubic" | "qwen2vl_bicubic_antialias"
    output_size: Optional[int]  # 224 / 252 / None (dynamic)
    smart_resize_factor: Optional[int]  # None / 28 / 32
    smart_resize_min_pixels: Optional[int]  # None / 4*28*28 / 65536
    smart_resize_max_pixels: Optional[int]
    normalize_offset: float  # (u8 - 127.5) / 127.5 in fp32 (NOT ((x/255)-0.5)/0.5; differs by 5.9e-8)
    normalize_scale: float
    uint8_truncation_before_processor: bool  # N1.5 casts (x*255).uint8 (truncation) before the Eagle tiler
    frames_per_camera: int
    image_order: str  # N1.7 "frame-major": [cam1@t-15, cam2@t-15, cam1@t, cam2@t]
    formalize_language: bool  # lower-case + strip [^\w\s]
    instruction_as_repr_list: bool  # N1.5: the template sees repr([instruction]) -> tokens `['` (677) and `']` (660)
    state_normalization: str  # "sincos_or_minmax_per_config" | "minmax_clip_or_sincos_per_group" | "q01q99_clip"
    action_decode: str  # host-pipeline-parity §5 recipe name
    host_cost_ms: float  # host-pipeline-parity §6 (pinned cores)


@dataclass(frozen=True)
class PromptConfig:
    """Prompt template and token ids (01-architecture-synthesis.md §1.7; static-shape-plan §4.1; arch-n1x)."""

    tokenizer: str  # "qwen2_bpe_eagle2" | "qwen2_bpe_eagle3" | "qwen2_bpe_cosmos_reason2"
    vocab_entries: int  # tokenizer entries (151680 / 151680 / 151669); embed table rows are in LLMConfig.vocab
    pad_id: int  # <|endoftext|> 151643
    eos_id: int  # <|im_end|> 151645
    im_start_id: int  # 151644
    im_end_id: int  # 151645
    newline_id: int  # 198
    image_token_id: int  # <IMG_CONTEXT> 151669 | <|image_pad|> 151655
    image_open_id: Optional[int]  # <img> 151670 | <|vision_start|> 151652
    image_close_id: Optional[int]  # </img> 151671 | <|vision_end|> 151653
    add_bos: bool  # False in all three
    layout: str  # "images_then_text_with_generation_prompt" | "text_then_images" | "images_then_text_no_system"
    system_prompt: bool
    generation_prompt: bool
    header_ids: Tuple[int, ...]  # fixed leading tokens
    per_image_prefix_ids: Tuple[int, ...]  # tokens before the image run (k=1); `<image k>` digit varies with k
    per_image_prefix_len: int  # 6 (N1.5 `<image k><img>`), 6 (N1.6 `<image k>` 5 + `<img>`), 1 (N1.7 vision_start)
    per_image_suffix_len: int  # `</img>` 1 | `<|vision_end|>` 1
    suffix_ids: Tuple[int, ...]  # fixed trailing tokens
    text_extra_ids: Tuple[int, ...]  # N1.5 bracket tokens (677, 660) that wrap the instruction
    fixed_tokens_formula: str  # human-readable S formula
    # N1.6 `<image k>` renders as 5 tokens `27 1805 220 <digit> 29`; N1.5 as `<image k><img>` = 6 tokens.

    def seq_len(self, n_images: int, n_text: int, tokens_per_image: int) -> int:
        """``S`` for ``n_images`` images and ``n_text`` instruction BPE tokens (static-shape-plan §4.1)."""
        per_image = self.per_image_prefix_len + tokens_per_image + self.per_image_suffix_len
        return len(self.header_ids) + n_images * per_image + n_text + len(self.text_extra_ids) + len(self.suffix_ids)


@dataclass(frozen=True)
class LayoutConfig:
    """One static shape key (static-shape-plan.md §2, §4.2, §6): all device tensor shapes are fixed per layout."""

    name: str
    embodiment_tag: str
    embodiment_id: int
    n_cameras: int
    n_images: int  # cameras x frames
    grid_thw: Optional[Tuple[int, int, int]]  # N1.7 ViT grid per image (t, h, w) in patches; None for SigLIP
    patches_per_image: int
    patch_rows_pad: int  # per image, multiple of 32 (324 -> 352)
    tokens_per_image: int
    s_golden: Optional[int]  # observed S of the golden sample (None if no golden for this layout)
    s_typical: str  # observed / expected range as text
    s_pad_llm: int  # multiple of 128 (patched tt_transformers MLP or tt_dit encoder)
    s_pad_llm_unpatched: int  # multiple of 512 when >= 512 for the unpatched tt_transformers MLP
    s_pad_vl: int  # rows for vlln / VL-SA / cross K/V (== s_pad_llm in the first implementation)
    l_max: int  # text-token budget at s_pad_llm (incl. N1.5's 2 bracket tokens)
    text_padding: str  # "right" (N1.5/N1.7) | "text_slot" (N1.6: pads between text and images, explicit mask)
    image_slots: Tuple[Tuple[int, int], ...]  # fixed half-open slot ranges of the image tokens under the padding policy
    non_image_slots: Optional[int]  # count of non-image slots at s_pad_llm
    dit_m_logical: int  # 49 / 51 / 41
    dit_m_pad: int  # 64
    cross_image_keys: Optional[Tuple[int, int]]  # (real, padded) image keys for idx%4==2 blocks (None for N1.5)
    cross_text_keys: Optional[Tuple[int, int]]  # (real slots, padded) non-image keys for idx%4==0 blocks
    cross_all_keys: Optional[Tuple[int, int]]  # N1.5: (S, s_pad_vl) all tokens
    state_shape: Tuple[int, int, int]
    noise_shape: Tuple[int, int, int]
    valid_slice: Tuple[int, int]  # (horizon_valid, dims_valid) of the padded action tensors
    action_groups: Tuple[str, ...]
    rope_table: str  # "constant_arange" | "per_call_rows" | "constant_mrope_per_layout"
    notes: str = ""


@dataclass(frozen=True)
class GR00TConfig:
    """Everything the port needs to know about one GR00T version (01-architecture-synthesis.md §1)."""

    version: str
    display_name: str
    checkpoint: CheckpointConfig
    vision: VisionConfig
    projector: ProjectorConfig
    llm: LLMConfig
    vl_adapter: VLAdapterConfig
    dit: DiTConfig
    encoders: EncoderConfig
    sampler: SamplerConfig
    preprocessing: PreprocessingConfig
    prompt: PromptConfig
    embodiment_ids: Mapping[str, int]  # tag -> slot id (base checkpoint)
    canonical_embodiment: str
    layouts: Tuple[LayoutConfig, ...]
    canonical_layout_name: str
    backbone_feature_dim: int = 2048
    reference_code: str = ""
    golden_dir: str = ""  # relative to the golden root

    @property
    def canonical_embodiment_id(self) -> int:
        return self.embodiment_ids[self.canonical_embodiment]

    @property
    def canonical_layout(self) -> LayoutConfig:
        return self.layout(self.canonical_layout_name)

    def layout(self, name: str) -> LayoutConfig:
        for lay in self.layouts:
            if lay.name == name:
                return lay
        raise KeyError(f"{self.version}: unknown layout {name!r}; have {[l.name for l in self.layouts]}")

    def embodiment_id(self, tag: str) -> int:
        if tag not in self.embodiment_ids:
            raise KeyError(f"{self.version}: unknown embodiment tag {tag!r}; have {sorted(self.embodiment_ids)}")
        return self.embodiment_ids[tag]


# --------------------------------------------------------------------------------------------------------------------
# Shared constants
# --------------------------------------------------------------------------------------------------------------------

# Qwen2 BPE special ids shared by the three tokenizers (arch-n15 §4.2, arch-n16 §1, arch-n17 §1.3).
_PAD_ID = 151643
_IM_START = 151644
_IM_END = 151645
_NEWLINE = 198
# `<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n` (arch-n15 §4.2 table, arch-n16 §4.2).
_SYSTEM_USER_HEADER: Tuple[int, ...] = (
    151644,
    8948,
    198,
    2610,
    525,
    264,
    10950,
    17847,
    13,
    151645,
    198,
    151644,
    872,
    198,
)

_SAMPLER_BUCKETS: Tuple[int, ...] = (0, 250, 500, 750)

_QWEN3_LLM_COMMON = dict(
    n_layers_total=28,
    hidden=2048,
    n_heads=16,
    n_kv_heads=8,
    head_dim=128,
    intermediate=6144,
    rms_eps=1e-6,
    qk_norm=True,
    attention_bias=False,
    rope_layout="rotate_half",
    tie_word_embeddings=True,
)

_DIT_COMMON = dict(
    width=1536,
    n_heads=32,
    head_dim=48,
    ff_dim=6144,
    activation="gelu_tanh",
    attention_bias=True,
    cross_attention_dim=2048,
    cross_block_parity="even",
    adaln_chunks=2,
    adaln_order=("scale", "shift"),
    adaln_ln_eps=1e-5,
    norm3_affine=False,
    norm3_eps=1e-5,
    out_norm_eps=1e-6,
    proj_out_1_order=("shift", "scale"),
    out_dim=1024,
    timestep_embed_dim=256,
    timestep_flip_sin_to_cos=True,
    timestep_downscale_freq_shift=1,
    timestep_max_period=10000.0,
    timestep_mlp_hidden=1536,
    params_timestep_encoder=2_755_584,
    params_proj_out=6_295_552,
)

_SIGLIP_COMMON = dict(
    n_layers=27,
    hidden=1152,
    n_heads=16,
    head_dim=72,
    mlp_dim=4304,
    activation="gelu_tanh",
    ln_eps=1e-6,
    attention_bias=True,
    qkv_fused_in_checkpoint=False,
    patch_size=14,
    temporal_patch=1,
    patch_in_features=588,
    patch_column_order="(kh,kw,c)",
    pos_table_rows=256,
    pos_grid=16,
    merge_size=1,
    frames_per_camera=1,
    rope="none",
    rope_theta=None,
    rope_freqs_per_axis=None,
    post_layernorm=True,
    deepstack_layers=(),
    head_dim_pad=96,
    patch_in_features_pad=608,
    executed_params=412_442_352,
    executed_bytes_bf16=824_884_704,
)

# --------------------------------------------------------------------------------------------------------------------
# N1.5
# --------------------------------------------------------------------------------------------------------------------

_N15_EMBODIMENTS: Dict[str, int] = {  # arch-n15 §4.4 (gr00t/data/embodiment_tags.py:19-47)
    "gr1": 24,
    "oxe_droid": 17,
    "agibot_genie1": 26,
    "new_embodiment": 31,
}
_N15_GROUPS = ("left_arm", "right_arm", "left_hand", "right_hand")

N15_CONFIG = GR00TConfig(
    version="n15",
    display_name="GR00T-N1.5-3B",
    reference_code="Isaac-GR00T-n1.5 (n1.5-release, 4af2b622)",
    golden_dir="n15",
    checkpoint=CheckpointConfig(
        hf_repo=HF_REPOS["n15"],
        hf_sha=HF_SNAPSHOT_SHAS["n15"],
        n_shards=3,
        n_tensors=899,
        n_params=2_724_163_520,
        n_bytes=5_448_327_040,
        shard_files=tuple(f"model-0000{i}-of-00003.safetensors" for i in (1, 2, 3)),
        straddling_prefix="backbone.eagle_model.language_model.model.layers.5.",
        vision_prefix="backbone.eagle_model.vision_model.",
        projector_prefix="backbone.eagle_model.mlp1.",
        llm_prefix="backbone.eagle_model.language_model.model.",
        lm_head_key="backbone.eagle_model.language_model.lm_head.weight",
        vl_self_attention_prefix="action_head.vl_self_attention.",
    ),
    vision=VisionConfig(
        family="siglip_v1",
        pos_interpolation="none",
        image_size=224,
        patches_per_image=256,
        tokens_per_image=256,
        **_SIGLIP_COMMON,
    ),
    projector=ProjectorConfig(
        kind="linear",
        in_dim=1152,
        hidden_dim=None,
        out_dim=2048,
        ln_dim=None,
        ln_eps=None,
        ln_position=None,
        activation=None,
        merge_channel_order=None,
        image_token_id=151669,
        splice="overwrite",
        params=2_361_344,
    ),
    llm=LLMConfig(
        family="qwen3_1p7b",
        n_layers_executed=12,
        rope_mode="1d",
        rope_theta=1e6,
        mrope_section=None,
        vocab=151_680,
        final_norm_applied=True,
        deepstack_add_after_layers=(),
        massive_activation_channels=(1793, 1999),
        executed_layer_params=604_032_000,
        executed_layer_bytes_bf16=1_208_064_000,
        **_QWEN3_LLM_COMMON,
    ),
    vl_adapter=VLAdapterConfig(
        dim=2048,
        vlln_eps=1e-5,
        n_blocks=4,
        n_heads=32,
        head_dim=64,
        ff_dim=8192,
        ln_eps=1e-5,
        activation="gelu_tanh",
        attention_bias=True,
        masked=False,
        reference_dtype="fp32",
        params=201_433_088,
    ),
    dit=DiTConfig(
        class_name="cross_attention_dit.DiT",
        n_blocks=16,
        cross_key_set="all",
        n_query_tokens=49,
        n_future_tokens=32,
        params_blocks=541_335_552,
        **_DIT_COMMON,
    ),
    encoders=EncoderConfig(
        max_num_embodiments=32,
        max_state_dim=64,
        max_action_dim=32,
        action_horizon=16,
        hidden_size=1024,
        input_embedding_dim=1536,
        state_encoder_dims=(64, 1024, 1536),
        action_encoder_w1=(32, 1536),
        action_encoder_w2=(3072, 1536),
        action_encoder_w3=(1536, 1536),
        tau_dim=1536,
        tau_order=("sin", "cos"),
        tau_max_period=10000.0,
        action_decoder_dims=(1024, 1024, 32),
        position_embedding_rows=1024,
        state_dim_pad=64,
        action_dim_pad=32,
        used_slot_bytes_bf16=19_700_000,
    ),
    sampler=SamplerConfig(
        num_inference_timesteps=4,
        num_timestep_buckets=1000,
        dt=0.25,
        timestep_buckets=_SAMPLER_BUCKETS,
        noise_shape=(1, 16, 32),
        euler_dtype="fp32_actions_bf16_velocity",
        noise_dtype="fp32",
        velocity_rows_from_end=16,
    ),
    preprocessing=PreprocessingConfig(
        image_pipeline="x/255 -> CenterCrop(243) -> Resize(224, bilinear, antialias) -> (x*255).uint8 -> Eagle2.5 "
        "tiler (1 tile @224, no thumbnail) -> (u8-127.5)/127.5",
        crop_fraction=0.95,
        shortest_edge=256,
        crop_size=243,
        resize_interpolation="torchvision_bilinear_antialias",
        output_size=224,
        smart_resize_factor=None,
        smart_resize_min_pixels=None,
        smart_resize_max_pixels=None,
        normalize_offset=127.5,
        normalize_scale=127.5,
        uint8_truncation_before_processor=True,
        frames_per_camera=1,
        image_order="camera",
        formalize_language=False,
        instruction_as_repr_list=True,
        state_normalization="sincos_or_minmax_per_config",
        action_decode="minmax_no_clip_torch_fp32",
        host_cost_ms=0.8,
    ),
    prompt=PromptConfig(
        tokenizer="qwen2_bpe_eagle2",
        vocab_entries=151_680,
        pad_id=_PAD_ID,
        eos_id=_IM_END,
        im_start_id=_IM_START,
        im_end_id=_IM_END,
        newline_id=_NEWLINE,
        image_token_id=151669,
        image_open_id=151670,
        image_close_id=151671,
        add_bos=False,
        layout="images_then_text_with_generation_prompt",
        system_prompt=True,
        generation_prompt=True,
        header_ids=_SYSTEM_USER_HEADER,
        per_image_prefix_ids=(27, 1805, 220, 16, 29, 151670),  # `<image 1><img>` (arch-n15 §4.2)
        per_image_prefix_len=6,
        per_image_suffix_len=1,
        suffix_ids=(151645, 198, 151644, 77091, 198),  # `<|im_end|>\n<|im_start|>assistant\n`
        text_extra_ids=(677, 660),  # `['` and `']` of repr([instruction]) (critique C4)
        fixed_tokens_formula="S = 14 + 263*V + 2 + n_text + 5",
    ),
    embodiment_ids=_N15_EMBODIMENTS,
    canonical_embodiment="gr1",
    canonical_layout_name="gr1",
    layouts=(
        LayoutConfig(
            name="gr1",
            embodiment_tag="gr1",
            embodiment_id=24,
            n_cameras=1,
            n_images=1,
            grid_thw=None,
            patches_per_image=256,
            patch_rows_pad=256,
            tokens_per_image=256,
            s_golden=296,
            s_typical="296-297",
            s_pad_llm=384,
            s_pad_llm_unpatched=384,
            s_pad_vl=384,
            l_max=102,
            text_padding="right",
            image_slots=((20, 276),),
            non_image_slots=None,
            dit_m_logical=49,
            dit_m_pad=64,
            cross_image_keys=None,
            cross_text_keys=None,
            cross_all_keys=(296, 384),
            state_shape=(1, 1, 64),
            noise_shape=(1, 16, 32),
            valid_slice=(16, 26),
            action_groups=_N15_GROUPS,
            rope_table="constant_arange",
            notes="S_pad_vl 320 (10 tiles) is the tighter alternative (static-shape-plan §4.2/§5).",
        ),
        LayoutConfig(
            name="v3",
            embodiment_tag="oxe_droid",
            embodiment_id=17,
            n_cameras=3,
            n_images=3,
            grid_thw=None,
            patches_per_image=256,
            patch_rows_pad=256,
            tokens_per_image=256,
            s_golden=None,
            s_typical="813-817",
            s_pad_llm=896,
            s_pad_llm_unpatched=1024,
            s_pad_vl=896,
            l_max=88,
            text_padding="right",
            image_slots=((20, 276), (283, 539), (546, 802)),
            non_image_slots=None,
            dit_m_logical=49,
            dit_m_pad=64,
            cross_image_keys=None,
            cross_text_keys=None,
            cross_all_keys=(816, 896),
            state_shape=(1, 1, 64),
            noise_shape=(1, 16, 32),
            valid_slice=(16, 26),
            action_groups=_N15_GROUPS,
            rope_table="constant_arange",
            notes="DROID (id 17) / AgiBot (id 26), no golden; S_pad_vl 832 alternative.",
        ),
    ),
)

# --------------------------------------------------------------------------------------------------------------------
# N1.6
# --------------------------------------------------------------------------------------------------------------------

_N16_EMBODIMENTS: Dict[str, int] = {  # arch-n16 §3.4 (processing_gr00t_n1d6.py:33-45 merged with embodiment_id.json)
    "gr1": 20,
    "robocasa_panda_omron": 13,
    "behavior_r1_pro": 24,
    "unitree_g1": 8,
    "oxe_google": 0,
    "oxe_widowx": 1,
    "libero_panda": 2,
    "oxe_droid": 16,
    "new_embodiment": 10,
}
_N16_GROUPS = ("left_arm", "right_arm", "left_hand", "right_hand", "waist")

N16_CONFIG = GR00TConfig(
    version="n16",
    display_name="GR00T-N1.6-3B",
    reference_code="Isaac-GR00T-n1.6 (n1.6.1-release, 5dc80c4a)",
    golden_dir="n16",
    checkpoint=CheckpointConfig(
        hf_repo=HF_REPOS["n16"],
        hf_sha=HF_SNAPSHOT_SHAS["n16"],
        n_shards=2,
        n_tensors=1106,
        n_params=3_286_608_832,
        n_bytes=6_573_217_664,
        shard_files=("model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"),
        straddling_prefix="action_head.model.transformer_blocks.18.",
        vision_prefix="backbone.model.vision_model.",
        projector_prefix="backbone.model.mlp1.",
        llm_prefix="backbone.model.language_model.model.",
        lm_head_key="backbone.model.language_model.lm_head.weight",
        vl_self_attention_prefix=None,
    ),
    vision=VisionConfig(
        family="siglip2_naflex",
        pos_interpolation="bilinear_16_to_18",
        image_size=252,
        patches_per_image=324,
        tokens_per_image=81,
        **{**_SIGLIP_COMMON, "merge_size": 2},
    ),
    projector=ProjectorConfig(
        kind="unshuffle_ln_mlp",
        in_dim=4608,
        hidden_dim=2048,
        out_dim=2048,
        ln_dim=4608,
        ln_eps=1e-5,
        ln_position="post_concat",
        activation="gelu_erf",
        merge_channel_order="c*4+dy*2+dx",
        image_token_id=151669,
        splice="overwrite",
        params=13_644_800,
    ),
    llm=LLMConfig(
        family="qwen3_1p7b",
        n_layers_executed=16,
        rope_mode="1d",
        rope_theta=1e6,
        mrope_section=None,
        vocab=151_680,
        final_norm_applied=True,
        deepstack_add_after_layers=(),
        massive_activation_channels=(),
        executed_layer_params=805_376_000,
        executed_layer_bytes_bf16=1_610_752_000,
        **_QWEN3_LLM_COMMON,
    ),
    vl_adapter=VLAdapterConfig(
        dim=2048,
        vlln_eps=1e-5,
        n_blocks=0,
        n_heads=32,
        head_dim=64,
        ff_dim=8192,
        ln_eps=1e-5,
        activation="gelu_tanh",
        attention_bias=True,
        masked=False,
        reference_dtype="bf16",
        params=0,
    ),
    dit=DiTConfig(
        class_name="dit.AlternateVLDiT",
        n_blocks=32,
        cross_key_set="alternate_text_image",
        n_query_tokens=51,
        n_future_tokens=0,
        params_blocks=1_082_671_104,
        **_DIT_COMMON,
    ),
    encoders=EncoderConfig(
        max_num_embodiments=32,
        max_state_dim=128,
        max_action_dim=128,
        action_horizon=50,
        hidden_size=1024,
        input_embedding_dim=1536,
        state_encoder_dims=(128, 1024, 1536),
        action_encoder_w1=(128, 1536),
        action_encoder_w2=(3072, 1536),
        action_encoder_w3=(1536, 1536),
        tau_dim=1536,
        tau_order=("sin", "cos"),
        tau_max_period=10000.0,
        action_decoder_dims=(1024, 1024, 128),
        position_embedding_rows=1024,
        state_dim_pad=128,
        action_dim_pad=128,
        used_slot_bytes_bf16=20_300_000,
    ),
    sampler=SamplerConfig(
        num_inference_timesteps=4,
        num_timestep_buckets=1000,
        dt=0.25,
        timestep_buckets=_SAMPLER_BUCKETS,
        noise_shape=(1, 50, 128),
        euler_dtype="bf16",
        noise_dtype="bf16",
        velocity_rows_from_end=50,
    ),
    preprocessing=PreprocessingConfig(
        image_pipeline="LetterBoxPad -> SmallestMaxSize(256, INTER_AREA) -> FractionalCenterCrop(0.95) -> "
        "SmallestMaxSize(256, INTER_AREA) -> smart_resize(factor 28) -> PIL BICUBIC 252x252 -> (u8-127.5)/127.5",
        crop_fraction=0.95,
        shortest_edge=256,
        crop_size=243,
        resize_interpolation="pil_bicubic",
        output_size=252,
        smart_resize_factor=28,
        smart_resize_min_pixels=4 * 28 * 28,
        smart_resize_max_pixels=4096 * 28 * 28,
        normalize_offset=127.5,
        normalize_scale=127.5,
        uint8_truncation_before_processor=False,
        frames_per_camera=1,
        image_order="camera",
        formalize_language=True,
        instruction_as_repr_list=False,
        state_normalization="minmax_clip_or_sincos_per_group",
        action_decode="clip_relative_minmax_per_step_plus_last_state_waist_absolute",
        host_cost_ms=2.0,
    ),
    prompt=PromptConfig(
        tokenizer="qwen2_bpe_eagle3",
        vocab_entries=151_680,
        pad_id=_PAD_ID,
        eos_id=_IM_END,
        im_start_id=_IM_START,
        im_end_id=_IM_END,
        newline_id=_NEWLINE,
        image_token_id=151669,
        image_open_id=151670,
        image_close_id=151671,
        add_bos=False,
        layout="text_then_images",
        system_prompt=True,
        generation_prompt=False,
        header_ids=_SYSTEM_USER_HEADER,
        per_image_prefix_ids=(27, 1805, 220, 16, 29, 151670),  # `<image 1>` (5) + `<img>` (arch-n16 §4.2)
        per_image_prefix_len=6,
        per_image_suffix_len=1,
        suffix_ids=(151645, 198),  # `<|im_end|>\n`
        text_extra_ids=(),
        fixed_tokens_formula="S = 16 + 88*V + n_text",
    ),
    embodiment_ids=_N16_EMBODIMENTS,
    canonical_embodiment="gr1",
    canonical_layout_name="gr1",
    layouts=(
        LayoutConfig(
            name="gr1",
            embodiment_tag="gr1",
            embodiment_id=20,
            n_cameras=1,
            n_images=1,
            grid_thw=None,
            patches_per_image=324,
            patch_rows_pad=352,
            tokens_per_image=81,
            s_golden=116,
            s_typical="116-117",
            s_pad_llm=128,
            s_pad_llm_unpatched=128,
            s_pad_vl=128,
            l_max=24,
            text_padding="text_slot",
            image_slots=((44, 125),),
            non_image_slots=47,
            dit_m_logical=51,
            dit_m_pad=64,
            cross_image_keys=(81, 96),
            cross_text_keys=(47, 64),
            cross_all_keys=None,
            state_shape=(1, 1, 128),
            noise_shape=(1, 50, 128),
            valid_slice=(16, 29),
            action_groups=_N16_GROUPS,
            rope_table="per_call_rows",
            notes="With L_max=24: `<image 1><img>` [38,44), IMG [44,125), `</img>` 125, `<|im_end|>\\n` 126-127.",
        ),
        LayoutConfig(
            name="gr1_long_text",
            embodiment_tag="gr1",
            embodiment_id=20,
            n_cameras=1,
            n_images=1,
            grid_thw=None,
            patches_per_image=324,
            patch_rows_pad=352,
            tokens_per_image=81,
            s_golden=None,
            s_typical="up to 232",
            s_pad_llm=256,
            s_pad_llm_unpatched=256,
            s_pad_vl=256,
            l_max=152,
            text_padding="text_slot",
            image_slots=((172, 253),),
            non_image_slots=175,
            dit_m_logical=51,
            dit_m_pad=64,
            cross_image_keys=(81, 96),
            cross_text_keys=(175, 192),
            cross_all_keys=None,
            state_shape=(1, 1, 128),
            noise_shape=(1, 50, 128),
            valid_slice=(16, 29),
            action_groups=_N16_GROUPS,
            rope_table="per_call_rows",
        ),
        LayoutConfig(
            name="v3",
            embodiment_tag="robocasa_panda_omron",
            embodiment_id=13,
            n_cameras=3,
            n_images=3,
            grid_thw=None,
            patches_per_image=324,
            patch_rows_pad=352,
            tokens_per_image=81,
            s_golden=None,
            s_typical="291",
            s_pad_llm=384,
            s_pad_llm_unpatched=384,
            s_pad_vl=384,
            l_max=104,
            text_padding="text_slot",
            image_slots=((124, 205), (212, 293), (300, 381)),
            non_image_slots=141,
            dit_m_logical=51,
            dit_m_pad=64,
            cross_image_keys=(243, 256),
            cross_text_keys=(141, 160),
            cross_all_keys=None,
            state_shape=(1, 1, 128),
            noise_shape=(1, 50, 128),
            valid_slice=(16, 29),
            action_groups=_N16_GROUPS,
            rope_table="per_call_rows",
            notes="robocasa_panda_omron (13) / behavior_r1_pro (24); no golden.",
        ),
    ),
)

# --------------------------------------------------------------------------------------------------------------------
# N1.7
# --------------------------------------------------------------------------------------------------------------------

_N17_EMBODIMENTS: Dict[str, int] = {  # arch-n17 §4.3/§4.4 (processing_gr00t_n1d7.py:69-118 + embodiment_id.json)
    "oxe_droid_relative_eef_relative_joint": 24,
    "real_g1_relative_eef_relative_joints": 25,
    "real_r1_pro_sharpa_relative_eef": 26,
    "real_r1_pro_sharpa_relative_eef_human": 26,
    "real_r1_pro_sharpa_relative_eef_maxinsights": 26,
    "real_r1_pro_sharpa_relative_eef_mecka": 26,
    "xdof_relative_eef_relative_joint": 27,
    "xdof_relative_eef_relative_joint_subtask": 27,
    "new_embodiment": 10,
}
_N17_DROID_GROUPS = ("eef_9d", "gripper_position", "joint_position")
_N17_G1_GROUPS = (
    "left_wrist_eef_9d",
    "right_wrist_eef_9d",
    "left_hand",
    "right_hand",
    "left_arm",
    "right_arm",
    "waist",
    "base_height_command",
    "navigate_command",
)

N17_CONFIG = GR00TConfig(
    version="n17",
    display_name="GR00T-N1.7-3B",
    reference_code="Isaac-GR00T (main, 51d4c89f)",
    golden_dir="n17",
    checkpoint=CheckpointConfig(
        hf_repo=HF_REPOS["n17"],
        hf_sha=HF_SNAPSHOT_SHAS["n17"],
        n_shards=2,
        n_tensors=1031,
        n_params=3_455_180_928,
        n_bytes=6_910_361_856,
        shard_files=("model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"),
        straddling_prefix="action_head.model.transformer_blocks.19.",
        vision_prefix="backbone.model.model.visual.",
        projector_prefix=None,  # the merger lives inside `visual.` (weight-loader §2.3 rule 3)
        llm_prefix="backbone.model.model.language_model.",
        lm_head_key="backbone.model.lm_head.weight",
        vl_self_attention_prefix="action_head.vl_self_attention.",
    ),
    vision=VisionConfig(
        family="qwen3vl",
        n_layers=24,
        hidden=1024,
        n_heads=16,
        head_dim=64,
        mlp_dim=4096,
        activation="gelu_tanh",
        ln_eps=1e-6,
        attention_bias=True,
        qkv_fused_in_checkpoint=True,
        patch_size=16,
        temporal_patch=2,
        patch_in_features=1536,
        patch_column_order="(c,t,ph,pw)",
        pos_table_rows=2304,
        pos_grid=48,
        pos_interpolation="bilinear_4tap_per_grid",
        image_size=None,
        patches_per_image=None,
        tokens_per_image=None,
        merge_size=2,
        frames_per_camera=2,
        rope="2d",
        rope_theta=1e4,
        rope_freqs_per_axis=16,
        post_layernorm=False,
        deepstack_layers=(5, 11, 17),
        head_dim_pad=64,
        patch_in_features_pad=1536,
        dynamic_min_pixels=65_536,
        dynamic_factor=32,
        executed_params=406_957_056,
        executed_bytes_bf16=813_914_112,
    ),
    projector=ProjectorConfig(
        kind="qwen3vl_merger",
        in_dim=4096,
        hidden_dim=4096,
        out_dim=2048,
        ln_dim=1024,
        ln_eps=1e-6,
        ln_position="pre_concat",
        activation="gelu_erf",
        merge_channel_order="(mh,mw,c)",
        image_token_id=151655,
        splice="masked_scatter",
        n_deepstack=3,
        deepstack_ln_dim=4096,
        params=25_174_016,
        params_deepstack_each=25_180_160,
    ),
    llm=LLMConfig(
        family="qwen3vl_2b_text",
        n_layers_executed=16,
        rope_mode="mrope_interleaved",
        rope_theta=5e6,
        mrope_section=(24, 20, 20),
        vocab=151_936,
        final_norm_applied=False,
        deepstack_add_after_layers=(0, 1, 2),
        massive_activation_channels=(1793,),
        executed_layer_params=805_376_000,
        executed_layer_bytes_bf16=1_610_752_000,
        **_QWEN3_LLM_COMMON,
    ),
    vl_adapter=VLAdapterConfig(
        dim=2048,
        vlln_eps=1e-5,
        n_blocks=4,
        n_heads=32,
        head_dim=64,
        ff_dim=8192,
        ln_eps=1e-5,
        activation="gelu_tanh",
        attention_bias=True,
        masked=False,
        reference_dtype="bf16",
        params=201_433_088,
    ),
    dit=DiTConfig(
        class_name="dit.AlternateVLDiT",
        n_blocks=32,
        cross_key_set="alternate_text_image",
        n_query_tokens=41,
        n_future_tokens=0,
        params_blocks=1_082_671_104,
        **_DIT_COMMON,
    ),
    encoders=EncoderConfig(
        max_num_embodiments=32,
        max_state_dim=132,
        max_action_dim=132,
        action_horizon=40,
        hidden_size=1024,
        input_embedding_dim=1536,
        state_encoder_dims=(132, 1024, 1536),
        action_encoder_w1=(132, 1536),
        action_encoder_w2=(3072, 1536),
        action_encoder_w3=(1536, 1536),
        tau_dim=1536,
        tau_order=("sin", "cos"),
        tau_max_period=10000.0,
        action_decoder_dims=(1024, 1024, 132),
        position_embedding_rows=1024,
        state_dim_pad=160,
        action_dim_pad=160,
        used_slot_bytes_bf16=20_400_000,
    ),
    sampler=SamplerConfig(
        num_inference_timesteps=4,
        num_timestep_buckets=1000,
        dt=0.25,
        timestep_buckets=_SAMPLER_BUCKETS,
        noise_shape=(1, 40, 132),
        euler_dtype="bf16",
        noise_dtype="bf16",
        velocity_rows_from_end=40,
    ),
    preprocessing=PreprocessingConfig(
        image_pipeline="SmallestMaxSize(256, INTER_AREA) -> FractionalCenterCrop(0.95) -> SmallestMaxSize(256) -> "
        "Qwen2VLImageProcessor: smart_resize(factor 32, min 65536 px) bicubic+antialias -> (u8-127.5)/127.5 -> "
        "T=2 duplicate -> patchify (c,t,ph,pw) -> [n_patches, 1536]",
        crop_fraction=0.95,
        shortest_edge=256,
        crop_size=243,
        resize_interpolation="qwen2vl_bicubic_antialias",
        output_size=None,
        smart_resize_factor=32,
        smart_resize_min_pixels=65_536,
        smart_resize_max_pixels=16_777_216,
        normalize_offset=127.5,
        normalize_scale=127.5,
        uint8_truncation_before_processor=False,
        frames_per_camera=2,
        image_order="frame_major",
        formalize_language=True,
        instruction_as_repr_list=False,
        state_normalization="q01q99_clip",
        action_decode="clip_relative_minmax_per_step_rot6d_se3_compose_plus_last_state_absolute_q01q99",
        host_cost_ms=10.0,
    ),
    prompt=PromptConfig(
        tokenizer="qwen2_bpe_cosmos_reason2",
        vocab_entries=151_669,
        pad_id=_PAD_ID,
        eos_id=_IM_END,
        im_start_id=_IM_START,
        im_end_id=_IM_END,
        newline_id=_NEWLINE,
        image_token_id=151655,
        image_open_id=151652,
        image_close_id=151653,
        add_bos=False,
        layout="images_then_text_no_system",
        system_prompt=False,
        generation_prompt=False,
        header_ids=(151644, 872, 198),  # `<|im_start|>user\n` (arch-n17 §1.3)
        per_image_prefix_ids=(151652,),  # `<|vision_start|>`
        per_image_prefix_len=1,
        per_image_suffix_len=1,
        suffix_ids=(151645, 198),  # `<|im_end|>\n`
        text_extra_ids=(),
        fixed_tokens_formula="S = 5 + n_img*(T_img + 2) + n_text",
    ),
    embodiment_ids=_N17_EMBODIMENTS,
    canonical_embodiment="oxe_droid_relative_eef_relative_joint",
    canonical_layout_name="droid",
    layouts=(
        LayoutConfig(
            name="droid",
            embodiment_tag="oxe_droid_relative_eef_relative_joint",
            embodiment_id=24,
            n_cameras=2,
            n_images=4,
            grid_thw=(1, 16, 28),
            patches_per_image=448,
            patch_rows_pad=448,
            tokens_per_image=112,
            s_golden=469,
            s_typical="469-473",
            s_pad_llm=512,
            s_pad_llm_unpatched=512,
            s_pad_vl=512,
            l_max=51,
            text_padding="right",
            image_slots=((4, 116), (118, 230), (232, 344), (346, 458)),
            non_image_slots=64,
            dit_m_logical=41,
            dit_m_pad=64,
            cross_image_keys=(448, 448),
            cross_text_keys=(64, 64),
            cross_all_keys=None,
            state_shape=(1, 1, 132),
            noise_shape=(1, 40, 132),
            valid_slice=(40, 17),
            action_groups=_N17_DROID_GROUPS,
            rope_table="constant_mrope_per_layout",
            notes="320x180 -> 455x256 -> 448x256 (16:9); mRoPE slot j >= 459 -> position j - 392; S_pad_vl 480 alt.",
        ),
        LayoutConfig(
            name="g1_4x3",
            embodiment_tag="real_g1_relative_eef_relative_joints",
            embodiment_id=25,
            n_cameras=1,
            n_images=2,
            grid_thw=(1, 16, 22),
            patches_per_image=352,
            patch_rows_pad=352,
            tokens_per_image=88,
            s_golden=None,
            s_typical="~196",
            s_pad_llm=256,
            s_pad_llm_unpatched=256,
            s_pad_vl=256,
            l_max=71,
            text_padding="right",
            image_slots=((4, 92), (94, 182)),
            non_image_slots=80,
            dit_m_logical=41,
            dit_m_pad=64,
            cross_image_keys=(176, 192),
            cross_text_keys=(80, 96),
            cross_all_keys=None,
            state_shape=(1, 1, 132),
            noise_shape=(1, 40, 132),
            valid_slice=(40, 53),
            action_groups=_N17_G1_GROUPS,
            rope_table="constant_mrope_per_layout",
            notes="4:3 aspect assumed from key names (UNVERIFIED, synthesis §6 #11); the only G1 golden is 16:9 synthetic.",
        ),
        LayoutConfig(
            name="g1_synthetic_16x9",
            embodiment_tag="real_g1_relative_eef_relative_joints",
            embodiment_id=25,
            n_cameras=1,
            n_images=2,
            grid_thw=(1, 16, 28),
            patches_per_image=448,
            patch_rows_pad=448,
            tokens_per_image=112,
            s_golden=243,
            s_typical="243",
            s_pad_llm=256,
            s_pad_llm_unpatched=256,
            s_pad_vl=256,
            l_max=23,
            text_padding="right",
            image_slots=((4, 116), (118, 230)),
            non_image_slots=32,
            dit_m_logical=41,
            dit_m_pad=64,
            cross_image_keys=(224, 224),
            cross_text_keys=(32, 32),
            cross_all_keys=None,
            state_shape=(1, 1, 132),
            noise_shape=(1, 40, 132),
            valid_slice=(40, 53),
            action_groups=_N17_G1_GROUPS,
            rope_table="constant_mrope_per_layout",
            notes="golden/n17/g1_synthetic_*: DROID frames with the G1 embodiment id (golden-tap-map §8 item 6).",
        ),
        LayoutConfig(
            name="cam3_4x3",
            embodiment_tag="xdof_relative_eef_relative_joint",
            embodiment_id=27,
            n_cameras=3,
            n_images=6,
            grid_thw=(1, 16, 22),
            patches_per_image=352,
            patch_rows_pad=352,
            tokens_per_image=88,
            s_golden=None,
            s_typical="~556",
            s_pad_llm=640,
            s_pad_llm_unpatched=1024,
            s_pad_vl=640,
            l_max=95,
            text_padding="right",
            image_slots=((4, 92), (94, 182), (184, 272), (274, 362), (364, 452), (454, 542)),
            non_image_slots=112,
            dit_m_logical=41,
            dit_m_pad=64,
            cross_image_keys=(528, 544),
            cross_text_keys=(112, 128),
            cross_all_keys=None,
            state_shape=(1, 1, 132),
            noise_shape=(1, 40, 132),
            valid_slice=(40, 14),
            action_groups=(),
            rope_table="constant_mrope_per_layout",
            notes="XDOF (27) / R1-Pro Sharpa (26); valid dims 14 (XDOF) or 62 (R1-Pro); no golden.",
        ),
    ),
)

_CONFIGS: Mapping[str, GR00TConfig] = {"n15": N15_CONFIG, "n16": N16_CONFIG, "n17": N17_CONFIG}


def get_config(version: str) -> GR00TConfig:
    """Return the frozen config of ``version`` ('n15' | 'n16' | 'n17')."""
    _check_version(version)
    return _CONFIGS[version]


def all_configs() -> Tuple[GR00TConfig, ...]:
    return tuple(_CONFIGS[v] for v in GR00T_VERSIONS)


# Tile facts shared by every static-shape decision (static-shape-plan.md §1).
TILE = 32
LLM_SEQ_PAD_MULTIPLE = 128
LLM_SEQ_PAD_MULTIPLE_UNPATCHED_MLP = 512
DIT_QUERY_ROWS_PAD = 64
DIT_HEAD_DIM_PAD = 64
SIGLIP_HEAD_DIM_PAD = 96
# Blackhole p150a as measured (00-device-facts.md).
DEVICE_COMPUTE_GRID = (11, 10)
DEVICE_L1_BYTES_PER_CORE = 1536 * 1024
DEVICE_DRAM_CHANNELS = 8