vllm-omni/tests/diffusion/models/pi05/test_pi05_units.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""CPU unit tests for the π0.5 VLA model.

Deliberately lightweight: config plumbing and the checkpoint-boundary rule, the
state discretization path, relative actions, and the weight-load remap. Runs on
CPU with no weights and no lerobot. The full LeRobot parity check lives in
``tests/diffusion/models/pi05/test_pi05_parity.py``.

Most assertions here guard *silent* failure modes — cases where the wrong
behaviour still produces a well-shaped, finite action chunk:

  * normalization must run before state discretization;
  * the normalization mode comes from the checkpoint, since π0.5 defaults to
    quantile where π0 defaults to ``mean_std``;
  * relative-action checkpoints must be honoured, and an unresolvable
    ``relative_exclude_joints`` must raise rather than silently make every
    dimension relative;
  * MEM / RTC checkpoints must be rejected rather than served incorrectly.

    pytest tests/diffusion/models/pi05/test_pi05_units.py -v
"""

from __future__ import annotations

import logging
from typing import Any

import numpy as np
import pytest
import torch

from vllm_omni.diffusion.models.pi05 import modeling_pi05
from vllm_omni.diffusion.models.pi05.config import (
    Pi05Config,
    UnsupportedCheckpointCapabilityError,
    load_lerobot_norm_stats,
    resolve_excluded_action_indices,
)
from vllm_omni.diffusion.models.pi05.modeling_pi05 import (
    GemmaVariantConfig,
    Pi05AdaRMSNorm,
    Pi05ForActionPrediction,
    create_sinusoidal_pos_embedding,
    get_gemma_config,
    make_att_2d_masks,
)
from vllm_omni.diffusion.models.pi05.processor_pi05 import (
    Pi05ImageProcessor,
    Pi05Processor,
    Pi05RelativeActions,
    _assemble_model_inputs,
    apply_norm,
    as_state_vector,
    build_norm_stats,
    build_pi05_prompt,
    discretize_state,
    resize_with_pad,
)

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


# ----------------------------------------------------------------------------
# Pi05Config dataclass + checkpoint resolver
# ----------------------------------------------------------------------------
_LEROBOT_CFG: dict[str, Any] = {
    "type": "pi05",
    "paligemma_variant": "gemma_2b",
    "action_expert_variant": "gemma_300m",
    "chunk_size": 50,
    "max_action_dim": 32,
    "max_state_dim": 32,
    "num_inference_steps": 10,
    "image_resolution": [224, 224],
    "tokenizer_max_length": 200,
    "dtype": "float32",
    "input_features": {
        "observation.images.base_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
        "observation.images.left_wrist_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
        "observation.images.right_wrist_0_rgb": {"type": "VISUAL", "shape": [3, 224, 224]},
        "observation.state": {"type": "STATE", "shape": [32]},
    },
    "output_features": {"action": {"type": "ACTION", "shape": [32]}},
}

_EXPECTED_CAMERA_ORDER = [
    "observation.images.base_0_rgb",
    "observation.images.left_wrist_0_rgb",
    "observation.images.right_wrist_0_rgb",
]

_ACTION_NAMES = ["joint_0", "joint_1", "joint_2", "joint_3", "joint_4", "joint_5", "gripper"]


def test_config_parses_lerobot_pi05_json():
    """``_LEROBOT_CFG`` mirrors the runtime keys of a real checkpoint."""
    c = Pi05Config.from_model_config(_LEROBOT_CFG)
    assert c.tokenizer_max_length == 200
    assert c.state_num_bins == 256
    assert c.image_feature_keys == _EXPECTED_CAMERA_ORDER
    assert c.image_resolution == (224, 224)


def test_config_rejects_wrong_model_type():
    with pytest.raises(ValueError, match="type='pi05'"):
        Pi05Config.from_model_config(dict(_LEROBOT_CFG, type="pi0"))


def test_config_rejects_unsupported_partial_action_chunk():
    with pytest.raises(UnsupportedCheckpointCapabilityError, match="n_action_steps"):
        Pi05Config.from_model_config(dict(_LEROBOT_CFG, n_action_steps=10))


@pytest.mark.parametrize(
    "key,value",
    [
        ("output_features", {"action": {"type": "ACTION", "shape": [33]}}),
        ("input_features", {**_LEROBOT_CFG["input_features"], "observation.state": {"type": "STATE", "shape": [33]}}),
    ],
    ids=["action", "state"],
)
def test_config_rejects_a_schema_wider_than_the_model(key, value):
    """``max_action_dim`` / ``max_state_dim`` are the widths the weights were
    built for, so a wider checkpoint schema cannot be served."""
    with pytest.raises(ValueError, match="over max_state_dim"):
        Pi05Config.from_model_config(dict(_LEROBOT_CFG, **{key: value}))


@pytest.mark.parametrize(
    "key,value,attribute",
    [
        ("output_features", {"action": {"type": "ACTION", "shape": [7]}}, "action_dim"),
        (
            "input_features",
            {**_LEROBOT_CFG["input_features"], "observation.state": {"type": "STATE", "shape": [7]}},
            "state_dim",
        ),
    ],
    ids=["action", "state"],
)
def test_config_derives_the_real_width_from_the_schema(key, value, attribute):
    """The declared width decides how many values reach the prompt and how many
    columns come back, so a 7-joint checkpoint must not be padded out to 32."""
    config = Pi05Config.from_model_config(dict(_LEROBOT_CFG, **{key: value}))
    assert getattr(config, attribute) == 7


def test_config_rejects_inconsistent_openpi_metadata():
    with pytest.raises(ValueError, match="policy_server_config.action_dim"):
        Pi05Config.from_model_config(
            dict(
                _LEROBOT_CFG,
                output_features={"action": {"type": "ACTION", "shape": [7]}},
                policy_server_config={"action_dim": 32},
            )
        )


def test_config_non_square_resolution_raises():
    with pytest.raises(ValueError):
        Pi05Config(image_resolution=(224, 256))


@pytest.mark.parametrize("key", ["max_cameras", "num_inference_steps"])
def test_config_rejects_counts_below_one(key):
    """``build_model_inputs`` emits one image slot per camera and the Euler loop
    steps ``-1 / num_inference_steps``; neither is defined at zero."""
    with pytest.raises(ValueError, match=f"{key} must be at least 1"):
        Pi05Config.from_model_config(dict(_LEROBOT_CFG, **{key: 0}))


@pytest.mark.parametrize("key", ["state", "action"])
def test_config_rejects_norm_stats_that_do_not_match_the_declared_width(key):
    """A checkpoint's statistics and its feature schema come from different
    files, so they can disagree."""
    stats = {key: {"mode": "quantile", "q01": [0.0] * 2, "q99": [1.0] * 2}}
    with pytest.raises(ValueError, match=rf"norm_stats\[{key}\]\[q01\] has 2 entries, expected 32"):
        Pi05Config.from_model_config(dict(_LEROBOT_CFG, norm_stats=stats))


# ----------------------------------------------------------------------------
# LeRobot normalization-stats sidecar
# ----------------------------------------------------------------------------
# LeRobot keeps stats out of config.json: policy_preprocessor.json holds the
# structure and one safetensors file per stateful step holds the numbers.
def _normalized(raw, stats):
    """The serving path's state normalization, without the pipeline."""
    return apply_norm(torch.as_tensor(raw, dtype=torch.float32), build_norm_stats({"state": stats}, "state")).numpy()


def _write_lerobot_checkpoint(
    tmp_path,
    *,
    norm_map: dict[str, str] | None = None,
    stats: dict[str, dict[str, list[float]]] | None = None,
    write_state_file: bool = True,
) -> str:
    """Write a minimal LeRobot-shaped checkpoint dir; return its path."""
    import json as _json

    import safetensors.torch as _st

    tmp_path.mkdir(parents=True, exist_ok=True)
    state_file = "policy_preprocessor_step_3_normalizer_processor.safetensors"
    step: dict = {
        "registry_name": "normalizer_processor",
        "config": {
            "eps": 1e-08,
            "features": {},
            "norm_map": norm_map if norm_map is not None else {"STATE": "QUANTILES", "ACTION": "QUANTILES"},
        },
    }
    if write_state_file:
        step["state_file"] = state_file
        flat = {
            f"{feature}.{stat}": torch.tensor(values, dtype=torch.float32)
            for feature, entry in (stats or {}).items()
            for stat, values in entry.items()
        }
        _st.save_file(flat, str(tmp_path / state_file))

    (tmp_path / "policy_preprocessor.json").write_text(
        _json.dumps({"name": "policy_preprocessor", "steps": [step]}), encoding="utf-8"
    )
    (tmp_path / "config.json").write_text(_json.dumps(_config_for(stats)), encoding="utf-8")
    return str(tmp_path)


def _config_for(stats: dict | None) -> dict:
    """``_LEROBOT_CFG`` with the state/action widths the stats actually cover."""
    if not stats:
        return _LEROBOT_CFG
    config = dict(_LEROBOT_CFG)
    state = next(iter(stats.get("observation.state", {}).values()), None)
    action = next(iter(stats.get("action", {}).values()), None)
    if state is not None:
        features = dict(config["input_features"])
        features["observation.state"] = {"type": "STATE", "shape": [len(state)]}
        config["input_features"] = features
    if action is not None:
        config["output_features"] = {"action": {"type": "ACTION", "shape": [len(action)]}}
    return config


# A LeRobot state_dict carries every stat compute_stats emits, not just the two
# the norm_map selects.
_FULL_DATASET_STATS = {
    "observation.state": {
        "mean": [10.0, 20.0],
        "std": [1.0, 2.0],
        "min": [-5.0, -5.0],
        "max": [5.0, 5.0],
        "q01": [-1.0, -2.0],
        "q99": [1.0, 2.0],
    },
    "action": {
        "mean": [0.5, 0.5],
        "std": [0.1, 0.1],
        "min": [-1.0, -1.0],
        "max": [1.0, 1.0],
        "q01": [-0.25, -0.5],
        "q99": [0.25, 0.5],
    },
}


@pytest.mark.parametrize(
    "declared,expected",
    [
        ("QUANTILES", {"mode": "quantile", "q01": [-1.0, -2.0], "q99": [1.0, 2.0]}),
        ("MEAN_STD", {"mode": "mean_std", "mean": [10.0, 20.0], "std": [1.0, 2.0]}),
        ("IDENTITY", None),
    ],
)
def test_load_lerobot_norm_stats_takes_the_mode_from_norm_map(tmp_path, declared, expected):
    """The regression guard. ``_FULL_DATASET_STATS`` carries all six statistics,
    as a real state_dict does, so an implementation that sniffed the present keys
    would answer mean_std for every case here. Only norm_map distinguishes them.
    """
    path = _write_lerobot_checkpoint(
        tmp_path, norm_map={"STATE": declared, "ACTION": "QUANTILES"}, stats=_FULL_DATASET_STATS
    )
    stats = load_lerobot_norm_stats(path)

    if expected is None:
        assert "state" not in stats
    else:
        assert stats["state"] == expected
    assert stats["action"]["mode"] == "quantile"


@pytest.mark.parametrize("shape", ["no_sidecar", "no_state_file"])
def test_load_lerobot_norm_stats_absent_stats_are_none(tmp_path, shape):
    """``lerobot/pi05_base`` is the second shape: a normalizer step with no
    stats attached. Neither is an error — the client is then expected to send an
    already-normalized state."""
    if shape == "no_sidecar":
        path = str(tmp_path)
    else:
        path = _write_lerobot_checkpoint(tmp_path, write_state_file=False)

    assert load_lerobot_norm_stats(path) is None


def test_load_lerobot_norm_stats_missing_declared_state_file_raises(tmp_path):
    path = _write_lerobot_checkpoint(tmp_path, stats=_FULL_DATASET_STATS)
    (tmp_path / "policy_preprocessor_step_3_normalizer_processor.safetensors").unlink()

    with pytest.raises(FileNotFoundError, match="normalizer state"):
        load_lerobot_norm_stats(path)


def test_pipeline_rejects_missing_model_weights_before_initialization(tmp_path):
    from vllm_omni.diffusion.models.pi05.pipeline_pi05 import Pi05Pipeline

    pipeline = object.__new__(Pi05Pipeline)
    pipeline.model_dir = str(tmp_path)

    with pytest.raises(FileNotFoundError, match="model.safetensors"):
        pipeline._initialize_model()


def test_build_config_merges_the_deploy_yaml_over_the_checkpoint(tmp_path):
    """The deploy yaml carries the runtime knobs and the checkpoint carries the
    feature schema and the stats. Both have to land on one config."""
    from types import SimpleNamespace

    from vllm_omni.diffusion.models.pi05.pipeline_pi05 import Pi05Pipeline

    pipeline = object.__new__(Pi05Pipeline)
    pipeline.model_dir = _write_lerobot_checkpoint(tmp_path, stats=_FULL_DATASET_STATS)

    config = pipeline._build_config(SimpleNamespace(model_config={"type": "pi05", "num_inference_steps": 4}))

    assert config.num_inference_steps == 4
    assert (config.state_dim, config.action_dim) == (2, 2)
    assert config.image_feature_keys == _EXPECTED_CAMERA_ORDER
    assert config.norm_stats["state"]["mode"] == "quantile"


class _WideActionModel:
    """Stands in for the 3.6B model, which is the one collaborator here that is
    genuinely expensive to build. It returns the padded ``max_action_dim`` width
    the real model emits, so the assertion is about the pipeline's cropping."""

    def __init__(self, chunk_size: int, max_action_dim: int):
        self._shape = (1, chunk_size, max_action_dim)

    def sample_actions(self, **kwargs):
        del kwargs
        return torch.ones(self._shape)


def test_pipeline_crops_actions_to_checkpoint_output_schema(monkeypatch):
    from vllm_omni.diffusion.models.pi05 import pipeline_pi05
    from vllm_omni.diffusion.request import OmniDiffusionRequest
    from vllm_omni.inputs.data import OmniDiffusionSamplingParams

    config = Pi05Config.from_model_config(
        dict(_LEROBOT_CFG, output_features={"action": {"type": "ACTION", "shape": [7]}})
    )
    pipeline = object.__new__(pipeline_pi05.Pi05Pipeline)
    pipeline.config = config
    pipeline.tokenizer = object()
    pipeline._device = torch.device("cpu")
    pipeline.processor = Pi05Processor(config, object(), torch.device("cpu"))
    pipeline.model = _WideActionModel(config.chunk_size, config.max_action_dim)
    monkeypatch.setattr(
        pipeline.processor,
        "build_model_inputs",
        lambda robot_obs: ([torch.empty(0)], [torch.empty(0)], torch.empty(0), torch.empty(0)),
    )
    request = OmniDiffusionRequest(
        prompt="",
        sampling_params=OmniDiffusionSamplingParams(
            extra_args={"robot_obs": {"state": np.zeros(config.max_state_dim)}},
        ),
        request_id="crop-test",
    )

    result = pipeline.forward(request)

    assert result.output["actions"].shape == (config.chunk_size, 7)


def test_load_lerobot_norm_stats_unknown_mode_raises(tmp_path):
    """Fail loud: an unreproducible mode served anyway is silently wrong."""
    path = _write_lerobot_checkpoint(
        tmp_path, norm_map={"STATE": "SOMETHING_NEW", "ACTION": "QUANTILES"}, stats=_FULL_DATASET_STATS
    )
    with pytest.raises(ValueError, match="SOMETHING_NEW"):
        load_lerobot_norm_stats(path)


def test_from_pretrained_backfills_sidecar_stats(tmp_path):
    """End to end: the stats reach norm_stats without any deploy yaml."""
    path = _write_lerobot_checkpoint(tmp_path, stats=_FULL_DATASET_STATS)
    config = Pi05Config.from_pretrained(path)

    assert config.norm_stats["state"]["mode"] == "quantile"
    assert config.norm_stats["action"]["mode"] == "quantile"


def test_from_pretrained_does_not_override_explicit_norm_stats(tmp_path):
    """config.json wins over the sidecar, matching LeRobot's override rule."""
    import json as _json

    path = _write_lerobot_checkpoint(tmp_path, stats=_FULL_DATASET_STATS)
    explicit = {"state": {"mode": "min_max", "min": [0.0] * 32, "max": [1.0] * 32}}
    (tmp_path / "config.json").write_text(_json.dumps({**_LEROBOT_CFG, "norm_stats": explicit}), encoding="utf-8")

    assert Pi05Config.from_pretrained(path).norm_stats == explicit


# ----------------------------------------------------------------------------
# Checkpoint boundary: declared-but-unconsumed capabilities must raise
# ----------------------------------------------------------------------------
@pytest.mark.parametrize(
    "key,value",
    [
        ("use_visual_memory", True),
        ("use_proprioceptive_memory", True),
        ("rtc_config", {"mode": "trained", "inference_delay": 2}),
        ("n_obs_steps", 6),
    ],
)
def test_config_rejects_unsupported_capability(key, value):
    """MEM and RTC change what a correct action chunk looks like and are not
    visible in the weights. Serving such a checkpoint must fail loudly."""
    with pytest.raises(UnsupportedCheckpointCapabilityError, match=key):
        Pi05Config.from_model_config(dict(_LEROBOT_CFG, **{key: value}))


# ----------------------------------------------------------------------------
# Relative actions
# ----------------------------------------------------------------------------
def test_relative_actions_without_action_names_raises():
    """``action_feature_names`` is populated from dataset metadata at training
    time. There is no dataset at serving time, so a relative-action checkpoint
    that omits it cannot resolve ``relative_exclude_joints`` and must fail."""
    with pytest.raises(UnsupportedCheckpointCapabilityError, match="action_feature_names"):
        Pi05Config.from_model_config(dict(_LEROBOT_CFG, use_relative_actions=True, action_feature_names=None))


def test_relative_actions_empty_exclude_list_needs_no_action_names():
    c = Pi05Config.from_model_config(dict(_LEROBOT_CFG, use_relative_actions=True, relative_exclude_joints=[]))
    assert c.use_relative_actions


def test_resolve_excluded_action_indices_unresolvable_raises():
    with pytest.raises(UnsupportedCheckpointCapabilityError):
        resolve_excluded_action_indices(["no_such_joint"], _ACTION_NAMES)


def _relative_step(enabled=True):
    return Pi05RelativeActions(
        enabled=enabled,
        exclude_joints=["gripper"],
        action_names=_ACTION_NAMES,
        max_action_dim=32,
    )


def test_relative_mask_excludes_gripper_and_padding():
    rel = _relative_step()
    assert rel.relative_mask[0]
    assert not rel.relative_mask[6], "gripper must stay absolute"
    assert not rel.relative_mask[len(_ACTION_NAMES) :].any(), "padded dims carry no signal"
    assert rel.num_relative_dims == 6


def test_relative_absolute_round_trip_is_exact():
    """The two directions are one object precisely so they cannot disagree."""
    rel = _relative_step()
    g = torch.Generator().manual_seed(0)
    actions = torch.randn(1, 50, 32, generator=g)
    state = torch.randn(32, generator=g)
    assert torch.allclose(rel.to_absolute(rel.to_relative(actions, state), state), actions, atol=1e-6)


def test_relative_rejects_mismatched_action_dim():
    rel = _relative_step()
    with pytest.raises(ValueError, match="max_action_dim"):
        rel.to_absolute(torch.zeros(1, 4, 8), torch.zeros(32))


# ----------------------------------------------------------------------------
# State discretization — π0.5's defining input path
# ----------------------------------------------------------------------------
def test_discretize_underflows_to_minus_one_and_saturates_at_the_top():
    """Out-of-range input is asymmetric, and that asymmetry is LeRobot's.

    ``np.digitize`` returns 0 for anything below the first bin edge, so the
    ``- 1`` makes a below-range state land in bin ``-1``; above-range input
    saturates at ``num_bins - 1`` on its own. The checkpoint was trained with
    ``" -1"`` in the state prompt for under-range dims, so clipping that to 0
    would change the tokens the model sees — LeRobot parity catches it.
    """
    out = discretize_state(np.array([-9.0, 9.0], dtype=np.float32), num_bins=256)
    assert out.tolist() == [-1, 255]


def test_normalize_state_does_not_clip_out_of_range():
    """The parity-critical contract: normalization is affine and nothing else.

    Clamping to [-1, 1] here would hide an under-range dimension from
    ``discretize_state``, turning LeRobot's ``-1`` bin into ``0`` and changing
    the state tokens in the prompt. LeRobot applies
    ``2.0 * (x - q01) / (q99 - q01) - 1.0`` with no clip.
    """
    stats = {"q01": [0.0] * 32, "q99": [10.0] * 32}
    out = _normalized([-10.0, 20.0] + [5.0] * 30, stats)
    assert out[0] == pytest.approx(-3.0)  # would be -1.0 if clipped
    assert out[1] == pytest.approx(3.0)  # would be 1.0 if clipped
    # and the under-range dim must reach the -1 bin end to end
    assert discretize_state(out, num_bins=256)[0] == -1


def test_prompt_serializes_the_declared_state_width():
    """LeRobot's tokenizer step discretizes the state at its real width and
    never pads to ``max_state_dim``."""
    prompt = build_pi05_prompt(task="x", normalized_state=np.zeros(7, dtype=np.float32))
    bins = prompt.split("State: ")[1].split(";")[0].split()
    assert len(bins) == 7


def test_prompt_template_matches_lerobot():
    prompt = build_pi05_prompt(
        task="  pick_up the red\nblock ",
        normalized_state=np.zeros(4, dtype=np.float32),
    )
    assert prompt.startswith("Task: pick up the red block, State: ")
    assert prompt.endswith(";\nAction: ")


# ----------------------------------------------------------------------------
# Input contract: fixed camera slots x 256 tokens + constant 200 text tokens
# ----------------------------------------------------------------------------
@pytest.mark.parametrize("views,expected", [(1, 456), (2, 712), (3, 968)])
def test_prefix_length_is_fixed_by_max_cameras(views, expected):
    """``max_cameras × 256`` image tokens plus a constant ``tokenizer_max_length``.

    The text segment is padded to its maximum and the camera slots are padded to
    ``max_cameras``, so the prefix shape does not vary with the request — only
    the attention mask does.
    """
    keys = _EXPECTED_CAMERA_ORDER[:views]
    config = Pi05Config(max_cameras=views, tokenizer_max_length=200, image_feature_keys=keys)
    observation = {
        "state": np.zeros(32, dtype=np.float32),
        "images": {key: np.zeros((4, 4, 3), dtype=np.uint8) for key in keys},
    }

    images, masks, lang_tokens, _ = _assemble_model_inputs(
        observation, config, _FakeTokenizer(), torch.device("cpu"), None
    )

    assert len(images) == views
    assert all(bool(mask) for mask in masks)
    assert lang_tokens.shape[-1] == 200
    assert 256 * len(images) + lang_tokens.shape[-1] == expected


# ----------------------------------------------------------------------------
# Shared math helpers (same contracts as π0)
# ----------------------------------------------------------------------------
def test_gemma_variant_dims():
    vlm = get_gemma_config("gemma_2b")
    assert (vlm.width, vlm.depth, vlm.mlp_dim, vlm.num_kv_heads, vlm.head_dim) == (2048, 18, 16384, 1, 256)
    expert = get_gemma_config("gemma_300m")
    assert (expert.width, expert.depth, expert.mlp_dim) == (1024, 18, 4096)


def test_sinusoidal_embedding_shape_and_oddity():
    out = create_sinusoidal_pos_embedding(torch.tensor([0.0, 1.0]), 64)
    assert out.shape == (2, 64)
    with pytest.raises(ValueError):
        create_sinusoidal_pos_embedding(torch.tensor([0.0]), 63)


def test_make_att_2d_masks_respects_padding():
    pad = torch.tensor([[True, True, False]])
    att = torch.tensor([[0, 0, 0]])
    out = make_att_2d_masks(pad, att)
    assert not out[0, :, 2].any()


def test_build_norm_stats_rejects_unknown_mode():
    with pytest.raises(ValueError, match="Unsupported normalization mode"):
        build_norm_stats({"action": {"mode": "qunatile"}}, "action")


@pytest.mark.parametrize(
    "entry",
    [
        {"mode": "mean_std", "mean": [0.0]},
        {"mode": "min_max", "min": [0.0]},
        {"mode": "quantile", "q01": [0.0]},
    ],
)
def test_build_norm_stats_rejects_incomplete_stats(entry):
    with pytest.raises(ValueError, match="requires"):
        build_norm_stats({"action": entry}, "action")


def test_normalize_quantile_round_trip_and_padded_tail():
    stats = build_norm_stats({"action": {"q01": [0.0, 0.0], "q99": [2.0, 4.0]}}, "action")
    x = torch.tensor([[0.5, -0.5, 7.0]])  # third entry is padding beyond the stats
    back = apply_norm(apply_norm(x, stats), stats, inverse=True)
    assert torch.allclose(back, x, atol=1e-5)
    assert back[0, 2] == 7.0, "padded tail must pass through untouched"


# ----------------------------------------------------------------------------
# Image processor
# ----------------------------------------------------------------------------
@pytest.mark.parametrize("value,expected", [(0.0, -1.0), (0.25, -0.5), (1.0, 1.0)])
@pytest.mark.parametrize("container", ["numpy", "torch"])
def test_float_image_uses_lerobot_zero_one_domain(value, expected, container):
    if container == "numpy":
        image = np.full((4, 4, 3), value, dtype=np.float32)
    else:
        image = torch.full((3, 4, 4), value, dtype=torch.float32)

    out = Pi05ImageProcessor(image_size=4).preprocess_single(image)

    assert out.shape == (1, 3, 4, 4)
    assert torch.allclose(out, torch.full_like(out, expected))


@pytest.mark.parametrize("bad", [-0.01, 1.01, np.nan, np.inf])
@pytest.mark.parametrize("container", ["numpy", "torch"])
def test_float_image_rejects_values_outside_lerobot_domain(bad, container):
    if container == "numpy":
        image = np.full((4, 4, 3), bad, dtype=np.float32)
    else:
        image = torch.full((3, 4, 4), bad, dtype=torch.float32)

    with pytest.raises(ValueError):
        Pi05ImageProcessor(image_size=4).preprocess_single(image)


def test_resize_with_pad_pads_with_minus_one():
    out = resize_with_pad(torch.zeros(1, 3, 100, 200), 224, 224)
    assert out.shape == (1, 3, 224, 224)
    assert out[0, 0, 0, 0] == -1.0


class _FakeTokenizer:
    def __call__(self, text, **kwargs):
        del text
        length = kwargs["max_length"]
        return {"input_ids": [0] * length, "attention_mask": [1] * length}


@pytest.mark.parametrize("explicit_none", [False, True])
def test_build_model_inputs_preserves_missing_middle_camera_slot(explicit_none):
    config = Pi05Config(image_feature_keys=_EXPECTED_CAMERA_ORDER, max_cameras=3)
    images = {
        _EXPECTED_CAMERA_ORDER[0]: np.zeros((4, 4, 3), dtype=np.uint8),
        _EXPECTED_CAMERA_ORDER[2]: np.full((4, 4, 3), 255, dtype=np.uint8),
    }
    if explicit_none:
        images[_EXPECTED_CAMERA_ORDER[1]] = None
    observation = {
        "images": images,
        "state": np.zeros(32, dtype=np.float32),
    }

    slots, masks, _, _ = _assemble_model_inputs(observation, config, _FakeTokenizer(), torch.device("cpu"), None)

    assert [bool(mask.item()) for mask in masks] == [True, False, True]
    torch.testing.assert_close(slots[0], torch.full_like(slots[0], -1.0))
    torch.testing.assert_close(slots[1], torch.full_like(slots[1], -1.0))
    torch.testing.assert_close(slots[2], torch.full_like(slots[2], 1.0))


def test_build_model_inputs_requires_configured_camera():
    config = Pi05Config(image_feature_keys=_EXPECTED_CAMERA_ORDER)
    observation = {
        "unknown_camera": np.zeros((4, 4, 3), dtype=np.uint8),
        "state": np.zeros(32, dtype=np.float32),
    }
    with pytest.raises(ValueError, match="configured camera"):
        _assemble_model_inputs(observation, config, _FakeTokenizer(), torch.device("cpu"), None)


# ----------------------------------------------------------------------------
# Model behaviour + weight-load remap (tiny action dims; full Gemma backbone)
# ----------------------------------------------------------------------------
# Module *structure* (no ``state_proj``, ``time_mlp_{in,out}`` widths, AdaRMS on
# the expert norms but not the prefix) is not asserted here: ``load_weights``
# audits every model parameter against the checkpoint and raises on a missing,
# unmatched or π0-shaped key, so a structural drift fails loudly at load time
# rather than silently. ``test_pi0_shaped_keys_are_not_silently_loaded`` guards
# that audit, and ``test_pi05_parity.py`` is the numerical oracle.
# Both Gemma towers keep their real widths but drop from 18 layers to 2 for
# every test in this section. Nothing here depends on depth: the weight remap
# matches by name, and the AR mask, prefix length and ODE determinism are
# shape-driven. Full-scale numerical correctness is ``test_pi05_parity.py``'s
# job, against the real checkpoint.
#
# Two dimensions must NOT be shrunk, both because something outside these
# configs is sized against them:
#
# * ``width`` — the vision tower's ``multi_modal_projector`` projects image
#   embeddings to the PaliGemma text width so the two can be concatenated.
#   Narrowing the text tower alone makes that concat fail.
# * the vision tower itself — So400m/14 on 224x224 is what yields exactly 256
#   tokens per view, which ``test_prefix_length_and_valid_len_for_each_view_count``
#   exists to verify.
#
# Depths are equal across the two towers on purpose: the action expert runs
# layer-by-layer alongside the VLM, so a mismatch would not be a valid model.
_SHRUNK_GEMMA = {
    "gemma_2b": GemmaVariantConfig(2048, 2, 1024, 8, 1, 256),
    "gemma_300m": GemmaVariantConfig(1024, 2, 512, 8, 1, 256),
}


@pytest.fixture(scope="module", autouse=True)
def _shrink_gemma_backbones():
    """Patch the module-level lookup ``Pi05ForActionPrediction.__init__`` calls."""
    patch = pytest.MonkeyPatch()
    patch.setattr(modeling_pi05, "get_gemma_config", lambda variant: _SHRUNK_GEMMA[variant])
    yield
    patch.undo()


def _tiny_pi05_model():
    """A fresh model. Use for tests that load weights or otherwise mutate it."""
    return Pi05ForActionPrediction(Pi05Config(max_action_dim=8, max_state_dim=8, chunk_size=4, n_action_steps=4))


@pytest.fixture(scope="module")
def tiny_model():
    """Shared read-only model, so the tests that only inspect it pay for one
    construction between them rather than one each."""
    return _tiny_pi05_model()


@pytest.fixture(autouse=True)
def _silence_pi05_loader():
    logging.getLogger("vllm_omni.diffusion.models.pi05.modeling_pi05").setLevel(logging.CRITICAL)
    yield


@pytest.mark.slow
def test_suffix_has_no_state_token(tiny_model):
    """π0's suffix is ``[state, actions...]`` with AR mask ``[1, 1, 0...]``;
    π0.5's is actions only with ``[1, 0...]``."""
    embs, pad, att, cond = tiny_model.embed_suffix(torch.randn(1, 4, 8), torch.tensor([1.0]))
    assert embs.shape[1] == 4
    assert att[0].tolist() == [1, 0, 0, 0]
    assert cond.shape == (1, tiny_model.expert_width)


# Not marked slow: constructs one small norm layer, no backbone.
def test_adarms_norm_returns_gate_only_when_conditioned():
    norm = Pi05AdaRMSNorm(8, cond_dim=4)
    out, gate = norm(torch.randn(1, 3, 8), torch.randn(1, 4))
    assert gate is not None and gate.shape == (1, 1, 8)

    plain = Pi05AdaRMSNorm(8, cond_dim=None)
    out, gate = plain(torch.randn(1, 3, 8))
    assert gate is None


@pytest.mark.slow
def test_lm_head_remap_to_embed_tokens():
    model = _tiny_pi05_model()
    target = "paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"
    payload = torch.full(dict(model.named_parameters())[target].shape, 0.5)
    model.load_weights([("model.paligemma_with_expert.paligemma.lm_head.weight", payload)], strict=False)
    assert torch.allclose(dict(model.named_parameters())[target], payload)


@pytest.mark.slow
def test_pi0_shaped_keys_are_not_silently_loaded():
    """A π0 checkpoint pointed at the π0.5 class would leave the AdaRMS expert
    randomly initialized. ``state_proj`` is the tell."""
    model = _tiny_pi05_model()
    with pytest.raises(RuntimeError, match="π0-shaped"):
        model.load_weights([("model.state_proj.weight", torch.zeros(1024, 8))])


@pytest.mark.slow
@pytest.mark.parametrize("num_real_views", [1, 2, 3])
def test_prefix_length_is_fixed_and_valid_len_tracks_masks(num_real_views, tiny_model):
    """The deployed model has three slots; only their validity varies."""
    live_text = 120
    images = [torch.zeros(1, 3, 224, 224) for _ in range(3)]
    masks = [torch.tensor([index < num_real_views]) for index in range(3)]
    lang = torch.zeros(1, 200, dtype=torch.long)
    lang_mask = torch.zeros(1, 200, dtype=torch.bool)
    lang_mask[:, :live_text] = True

    embs, pad_masks, _ = tiny_model.embed_prefix(images, masks, lang, lang_mask)
    assert embs.shape[1] == 256 * 3 + 200
    assert int(pad_masks.sum()) == 256 * num_real_views + live_text


@pytest.mark.slow
def test_sample_actions_shape_and_determinism(tiny_model):
    """Flow matching is an ODE: fixed noise must give a bit-identical chunk."""
    model = tiny_model.eval()
    images = [torch.zeros(1, 3, 224, 224) for _ in range(3)]
    masks = [torch.tensor([True]), torch.tensor([False]), torch.tensor([False])]
    lang = torch.zeros(1, 200, dtype=torch.long)
    lang_mask = torch.ones(1, 200, dtype=torch.bool)
    noise = torch.randn(1, 4, 8, generator=torch.Generator().manual_seed(42))

    with torch.no_grad():
        a1 = model.sample_actions(
            images=images, image_masks=masks, lang_tokens=lang, lang_masks=lang_mask, noise=noise, num_steps=2
        )
        a2 = model.sample_actions(
            images=images, image_masks=masks, lang_tokens=lang, lang_masks=lang_mask, noise=noise, num_steps=2
        )
    assert a1.shape == (1, 4, 8)
    assert torch.isfinite(a1).all()
    assert torch.equal(a1, a2)


@pytest.mark.slow
def test_sample_actions_uses_request_generator(tiny_model):
    model = tiny_model.eval()
    images = [torch.zeros(1, 3, 224, 224) for _ in range(3)]
    masks = [torch.tensor([True]), torch.tensor([False]), torch.tensor([False])]
    lang = torch.zeros(1, 200, dtype=torch.long)
    lang_mask = torch.ones(1, 200, dtype=torch.bool)

    def run(seed):
        return model.sample_actions(
            images=images,
            image_masks=masks,
            lang_tokens=lang,
            lang_masks=lang_mask,
            num_steps=2,
            generator=torch.Generator().manual_seed(seed),
        )

    with torch.no_grad():
        first = run(42)
        repeated = run(42)
        different = run(43)

    assert torch.equal(first, repeated)
    assert not torch.equal(first, different)


# ----------------------------------------------------------------------------
# Serving dtype
# ----------------------------------------------------------------------------
# The check that decides what runs lives on the pipeline, not on Pi05Config:
# Pi05Config.dtype records what the checkpoint declares and never reaches a
# cast, while _initialize_model casts using the top-level OmniDiffusionConfig.
@pytest.mark.parametrize(
    "declared,expected",
    [
        ("float32", torch.float32),
        ("bfloat16", torch.bfloat16),
        (torch.float32, torch.float32),
        (torch.bfloat16, torch.bfloat16),
    ],
)
def test_resolve_dtype_accepts_supported_serving_dtypes(declared, expected):
    from vllm_omni.diffusion.data import OmniDiffusionConfig
    from vllm_omni.diffusion.models.pi05.pipeline_pi05 import Pi05Pipeline

    assert Pi05Pipeline._resolve_dtype(OmniDiffusionConfig(dtype=declared)) is expected


@pytest.mark.parametrize("declared", ["float16", "fp16", torch.float16])
def test_resolve_dtype_rejects_unvalidated_dtypes(declared):
    """float16 resolves to a real torch dtype, so a permissive lookup would run
    the model in a precision nothing here has been checked in.

    Only dtypes that survive ``OmniDiffusionConfig`` normalization are worth
    asserting on: it maps a string it does not recognize onto bfloat16 with a
    warning, so e.g. ``"float64"`` never reaches this guard as float64.
    """
    from vllm_omni.diffusion.data import OmniDiffusionConfig
    from vllm_omni.diffusion.models.pi05.pipeline_pi05 import Pi05Pipeline

    with pytest.raises(ValueError, match="Unsupported π0.5 dtype"):
        Pi05Pipeline._resolve_dtype(OmniDiffusionConfig(dtype=declared))


def test_bfloat16_layout_is_applied_before_checkpoint_load(monkeypatch, tmp_path):
    """Checkpoint tensors must be copied directly into their final dtypes.

    LeRobot casts only the inner PaliGemma/expert module; the outer action and
    time projections remain float32. The layout is established before loading
    so float32-retained weights never pass through bfloat16.
    """
    from vllm_omni.diffusion.models.pi05 import pipeline_pi05

    class _TinyPi05(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.paligemma_with_expert = torch.nn.Module()
            self.paligemma_with_expert.backbone = torch.nn.Linear(2, 2)
            self.paligemma_with_expert.vision_tower = torch.nn.Linear(2, 2)
            self.paligemma_with_expert.input_layernorm = torch.nn.LayerNorm(2)
            self.paligemma_with_expert.register_buffer("floating_buffer", torch.ones(1))
            self.action_in_proj = torch.nn.Linear(2, 2)
            self.action_out_proj = torch.nn.Linear(2, 2)
            self.time_mlp_in = torch.nn.Linear(2, 2)
            self.time_mlp_out = torch.nn.Linear(2, 2)

    # DiffusersLoader constructs custom pipelines under the requested default
    # dtype. Start fully in bfloat16 to verify that the helper restores every
    # float32-retained tensor rather than relying on the construction default.
    model = _TinyPi05().to(dtype=torch.bfloat16)
    observed_dtypes = {}

    def fake_load_checkpoint(_self, model_to_load):
        observed_dtypes.update({name: param.dtype for name, param in model_to_load.named_parameters()})

    (tmp_path / "model.safetensors").touch()
    pipeline = object.__new__(pipeline_pi05.Pi05Pipeline)
    pipeline.model_dir = str(tmp_path)
    pipeline.config = object()
    pipeline._device = torch.device("cpu")
    pipeline._torch_dtype = torch.bfloat16
    monkeypatch.setattr(pipeline_pi05, "Pi05ForActionPrediction", lambda _config: model)
    monkeypatch.setattr(pipeline_pi05.Pi05Pipeline, "_load_checkpoint", fake_load_checkpoint)

    initialized = pipeline._initialize_model()

    assert observed_dtypes["paligemma_with_expert.backbone.weight"] is torch.bfloat16
    assert model.paligemma_with_expert.floating_buffer.dtype is torch.bfloat16
    for name in (
        "paligemma_with_expert.vision_tower.weight",
        "paligemma_with_expert.input_layernorm.weight",
        "action_in_proj.weight",
        "action_out_proj.weight",
        "time_mlp_in.weight",
        "time_mlp_out.weight",
    ):
        assert observed_dtypes[name] is torch.float32
    assert initialized is model
    assert not initialized.training


@pytest.mark.slow
def test_bfloat16_runs_and_tracks_float32(tiny_model):
    """bfloat16 is a supported serving dtype, so it has to produce a usable
    chunk — and one that still resembles float32.

    The bound is deliberately loose. This is not a parity check (bfloat16
    deviates from float32 by ~1-2% of the action range on the real checkpoint,
    which is why recipes/lerobot/Pi05.md tells operators to validate it on their
    own task). What it catches is gross breakage: a cast that silently did not
    happen, activations overflowing to inf, or a dtype mismatch that would make
    the two outputs unrelated.
    """
    import copy

    images = [torch.zeros(1, 3, 224, 224) for _ in range(3)]
    masks = [torch.tensor([True]), torch.tensor([False]), torch.tensor([False])]
    lang = torch.zeros(1, 200, dtype=torch.long)
    lang_mask = torch.ones(1, 200, dtype=torch.bool)
    noise = torch.randn(1, 4, 8, generator=torch.Generator().manual_seed(7))

    def run(model, dtype):
        with torch.no_grad():
            return model.sample_actions(
                images=[img.to(dtype) for img in images],
                image_masks=masks,
                lang_tokens=lang,
                lang_masks=lang_mask,
                noise=noise.to(dtype),
                num_steps=2,
            ).float()

    reference = run(tiny_model.eval(), torch.float32)
    from vllm_omni.diffusion.models.pi05.pipeline_pi05 import _set_inference_dtype

    bf16_model = copy.deepcopy(tiny_model).eval()
    _set_inference_dtype(bf16_model, torch.bfloat16)

    # The cast actually happened — otherwise everything below passes trivially.
    assert any(param.dtype is torch.bfloat16 for param in bf16_model.parameters())
    assert bf16_model.action_in_proj.weight.dtype is torch.float32
    assert bf16_model.action_out_proj.weight.dtype is torch.float32
    assert bf16_model.time_mlp_in.weight.dtype is torch.float32
    assert bf16_model.time_mlp_out.weight.dtype is torch.float32

    actual = run(bf16_model, torch.bfloat16)
    assert actual.shape == reference.shape
    assert torch.isfinite(actual).all()

    scale = float(reference.abs().max())
    assert float((actual - reference).abs().max()) < max(0.5 * scale, 1e-2)


# ----------------------------------------------------------------------------
# State normalization runs π0's stats path
# ----------------------------------------------------------------------------
@pytest.mark.parametrize(
    "stats,state,expected",
    [
        # A denominator that is small but not zero. ``normalize`` divides by
        # ``std + eps``; replacing a small std with 1.0 would be a different map.
        ({"mode": "mean_std", "mean": [0.0], "std": [1e-7]}, [1e-7], 1e-7 / (1e-7 + 1e-8)),
        ({"mode": "quantile", "q01": [0.0], "q99": [1e-7]}, [1e-7], 1.0),
        ({"mode": "quantile", "q01": [0.0], "q99": [1e-9]}, [1e-9], 1.0),
        # An exactly zero range maps the bound to -1 instead of dividing by zero.
        ({"mode": "quantile", "q01": [2.0], "q99": [2.0]}, [2.0], -1.0),
    ],
)
def test_normalize_state_uses_the_shared_denominator_rule(stats, state, expected):
    """State normalization and action unnormalization are the same affine map,
    so a narrow range keeps producing LeRobot's answer."""
    assert _normalized(state, stats)[0] == pytest.approx(expected, rel=1e-5)


# ----------------------------------------------------------------------------
# State wire contract
# ----------------------------------------------------------------------------
@pytest.mark.parametrize("bad", [float("nan"), float("inf"), float("-inf")])
def test_state_rejects_non_finite_values(bad):
    """``np.digitize`` maps nan and inf to the top bin and -inf to -1, all legal
    state tokens, so a dead sensor would reach the model as an ordinary prompt."""
    with pytest.raises(ValueError, match="NaN or Inf"):
        as_state_vector([0.0, bad], 2)


def test_state_rejects_a_multi_row_matrix():
    """Flattening would read a (2, 16) request as a 32-value state."""
    with pytest.raises(ValueError, match=r"\(D,\) or \(1, D\)"):
        as_state_vector(np.zeros((2, 16), dtype=np.float32), 32)


# ----------------------------------------------------------------------------
# Per-request denoising steps reach the model
# ----------------------------------------------------------------------------
class _SpyModel:
    """Records the ``num_steps`` ``forward`` hands to ``sample_actions``."""

    def __init__(self):
        self.seen: list = []
        self.generators: list = []

    def sample_actions(self, *, images, image_masks, lang_tokens, lang_masks, num_steps, generator):
        del images, image_masks, lang_tokens, lang_masks
        self.seen.append(num_steps)
        self.generators.append(generator)
        return torch.zeros(1, 50, 32)


def _spy_pipeline(spy: _SpyModel):
    from vllm_omni.diffusion.models.pi05.pipeline_pi05 import Pi05Pipeline

    pipeline = object.__new__(Pi05Pipeline)
    pipeline.config = Pi05Config.from_model_config(_LEROBOT_CFG)
    pipeline.tokenizer = _FakeTokenizer()
    pipeline._device = torch.device("cpu")
    pipeline.model = spy
    pipeline.processor = Pi05Processor(pipeline.config, _FakeTokenizer(), torch.device("cpu"))
    return pipeline


def _spy_request(num_inference_steps, generator=None):
    params = type("_Params", (), {})()
    params.extra_args = {
        "robot_obs": {
            "state": np.zeros(32, dtype=np.float32),
            "images": {key: np.zeros((4, 4, 3), dtype=np.uint8) for key in _EXPECTED_CAMERA_ORDER},
        }
    }
    params.num_inference_steps = num_inference_steps
    params.generator = generator
    request = type("_Request", (), {})()
    request.sampling_params = params
    return request


@pytest.mark.parametrize("requested", [None, 2, 25])
def test_request_denoising_steps_reach_sample_actions(requested):
    """The OpenPI E2E cannot see this: a pipeline that ignored the override and
    ran the static default would still return a finite chunk of the right shape."""
    spy = _SpyModel()
    output = _spy_pipeline(spy).forward(_spy_request(requested))

    assert spy.seen == [requested]
    assert output.output["actions"].shape == (50, 32)


def test_request_generator_reaches_sample_actions():
    generator = torch.Generator().manual_seed(42)
    spy = _SpyModel()

    _spy_pipeline(spy).forward(_spy_request(2, generator))

    assert spy.generators == [generator]


@pytest.mark.parametrize("bad", [0, -1, 2.5, True, "4"])
def test_request_denoising_steps_must_be_a_positive_integer(bad):
    spy = _SpyModel()
    with pytest.raises(ValueError, match="num_inference_steps must be a positive integer"):
        _spy_pipeline(spy).forward(_spy_request(bad))
    assert spy.seen == []