lerobot/tests/policies/groot/test_groot_n1_7.py

#!/usr/bin/env python

# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import inspect
import json
import sys
from types import SimpleNamespace
from unittest.mock import patch

import numpy as np
import pytest
import torch
from safetensors.torch import load_file
from torch import nn

from lerobot.configs import FeatureType, PolicyFeature
from lerobot.policies.factory import make_policy_config, make_pre_post_processors
from lerobot.policies.groot.configuration_groot import (
    GROOT_ACTION_DECODE_TRANSFORM_LIBERO,
    GROOT_N1_7,
    GROOT_N1_7_BASE_MODEL,
    GrootConfig,
    infer_groot_n1_7_action_execution_horizon,
    infer_groot_n1_7_action_horizon,
    normalize_groot_model_version,
)
from lerobot.policies.groot.modeling_groot import GrootPolicy
from lerobot.policies.groot.processor_groot import (
    GrootActionUnpackUnnormalizeStep,
    GrootN17ActionDecodeStep,
    GrootN17PackInputsStep,
    GrootN17VLMEncodeStep,
    N1_7_NATIVE_ACTION_HORIZON,
    _make_relative_action_training_stats,
    _transform_n1_7_image_for_vlm_albumentations,
    _transform_n1_7_image_for_vlm_torch,
    make_groot_pre_post_processors,
)
from lerobot.processor import (
    AbsoluteActionsProcessorStep,
    PolicyProcessorPipeline,
    RelativeActionsProcessorStep,
)
from lerobot.types import TransitionKey
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE


def _groot_features(
    state_dim: int, action_dim: int
) -> tuple[dict[str, PolicyFeature], dict[str, PolicyFeature]]:
    return (
        {
            f"{OBS_IMAGES}.front": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 256, 256)),
            OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(state_dim,)),
        },
        {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(action_dim,))},
    )


def _groot_config() -> GrootConfig:
    input_features, output_features = _groot_features(state_dim=8, action_dim=7)
    kwargs = {"action_decode_transform": GROOT_ACTION_DECODE_TRANSFORM_LIBERO}
    return GrootConfig(
        input_features=input_features,
        output_features=output_features,
        device="cpu",
        use_bf16=False,
        **kwargs,
    )


def _native_action_chunk(rows: list[list[float]]) -> torch.Tensor:
    chunk = torch.tensor(rows, dtype=torch.float32)
    if chunk.shape[0] >= N1_7_NATIVE_ACTION_HORIZON:
        return chunk[:N1_7_NATIVE_ACTION_HORIZON]
    tail = chunk[-1:].repeat(N1_7_NATIVE_ACTION_HORIZON - chunk.shape[0], 1)
    return torch.cat([chunk, tail], dim=0)


def _raw_n1_7_libero_config(model_path) -> GrootConfig:
    input_features, output_features = _groot_features(state_dim=8, action_dim=7)
    return GrootConfig(
        base_model_path=str(model_path),
        embodiment_tag="libero_sim",
        input_features=input_features,
        output_features=output_features,
        device="cpu",
        use_bf16=False,
        action_decode_transform=GROOT_ACTION_DECODE_TRANSFORM_LIBERO,
    )


def test_n1_7_backbone_accepts_transformers_5_layout_and_forwards_mm_token_type_ids(monkeypatch):
    pytest.importorskip("transformers")

    from transformers.feature_extraction_utils import BatchFeature

    import lerobot.policies.groot.groot_n1_7 as groot_n1_7

    class FakeLanguageModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.layers = nn.ModuleList([nn.Linear(1, 1) for _ in range(2)])

    class FakeInnerModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.language_model = FakeLanguageModel()
            self.visual = nn.Linear(1, 1)

    class FakeQwen3VLForConditionalGeneration(nn.Module):
        config = SimpleNamespace(image_token_id=42, video_token_id=43)

        def __init__(self):
            super().__init__()
            self.model = FakeInnerModel()
            self.forward_kwargs = None

        @classmethod
        def from_pretrained(cls, *args, **kwargs):
            return cls()

        @classmethod
        def _from_config(cls, *args, **kwargs):
            return cls()

        def eval(self):
            super().eval()
            return self

        def forward(self, **kwargs):
            self.forward_kwargs = kwargs
            assert "mm_token_type_ids" in kwargs
            batch_size, sequence_length = kwargs["input_ids"].shape
            features = torch.arange(batch_size * sequence_length * 4, dtype=torch.float32).view(
                batch_size, sequence_length, 4
            )
            return SimpleNamespace(hidden_states=[features])

    monkeypatch.setattr(
        groot_n1_7,
        "metadata",
        SimpleNamespace(version=lambda package: "5.3.0" if package == "transformers" else "0"),
        raising=False,
    )
    monkeypatch.setattr(groot_n1_7, "Qwen3VLForConditionalGeneration", FakeQwen3VLForConditionalGeneration)

    backbone = groot_n1_7.Qwen3Backbone(
        model_name="nvidia/Cosmos-Reason2-2B",
        select_layer=1,
        use_flash_attention=False,
    )

    assert len(backbone.language_model.layers) == 1
    output = backbone.forward(
        BatchFeature(
            data={
                "input_ids": torch.tensor([[1, 42, 2]]),
                "attention_mask": torch.tensor([[1, 1, 0]]),
                "mm_token_type_ids": torch.tensor([[0, 1, 0]]),
                "pixel_values": torch.zeros(1, 3, 2, 2),
                "image_grid_thw": torch.ones(1, 3, dtype=torch.long),
            }
        )
    )

    assert backbone.model.forward_kwargs["mm_token_type_ids"].tolist() == [[0, 1, 0]]
    assert output["backbone_features"].shape == (1, 3, 4)

    output = backbone.forward(
        BatchFeature(
            data={
                "input_ids": torch.tensor([[1, 42, 43, 2]]),
                "attention_mask": torch.tensor([[1, 1, 1, 0]]),
                "pixel_values": torch.zeros(1, 3, 2, 2),
                "image_grid_thw": torch.ones(1, 3, dtype=torch.long),
                "pixel_values_videos": torch.zeros(1, 3, 2, 2),
                "video_grid_thw": torch.ones(1, 3, dtype=torch.long),
            }
        )
    )

    assert backbone.model.forward_kwargs["mm_token_type_ids"].tolist() == [[0, 1, 2, 0]]
    assert backbone.model.forward_kwargs["mm_token_type_ids"].dtype == torch.int32
    assert output["backbone_features"].shape == (1, 4, 4)


def test_n1_7_backbone_preserves_missing_qwen_optional_dependency_error(monkeypatch):
    pytest.importorskip("transformers")

    import lerobot.policies.groot.groot_n1_7 as groot_n1_7

    monkeypatch.setattr(
        groot_n1_7,
        "metadata",
        SimpleNamespace(version=lambda package: "5.3.0" if package == "transformers" else "0"),
        raising=False,
    )
    monkeypatch.setattr(groot_n1_7, "Qwen3VLForConditionalGeneration", None)

    with pytest.raises(ImportError, match="Qwen3VLForConditionalGeneration is required"):
        groot_n1_7.Qwen3Backbone(
            model_name="nvidia/Cosmos-Reason2-2B",
            select_layer=0,
            use_flash_attention=False,
        )


def _write_raw_n1_7_libero_checkpoint(path):
    path.mkdir()
    (path / "config.json").write_text(
        json.dumps(
            {
                "model_type": "Gr00tN1d7",
                "architectures": ["Gr00tN1d7"],
                "model_name": "nvidia/Cosmos-Reason2-2B",
                "action_horizon": 40,
                "max_state_dim": 132,
                "max_action_dim": 132,
                "image_target_size": [256, 256],
            }
        )
    )
    (path / "processor_config.json").write_text(
        json.dumps(
            {
                "processor_class": "Gr00tN1d7Processor",
                "processor_kwargs": {
                    "clip_outliers": True,
                    "formalize_language": True,
                    "image_crop_size": [230, 230],
                    "image_target_size": [256, 256],
                    "shortest_image_edge": 256,
                    "crop_fraction": 0.95,
                    "use_albumentations": True,
                    "letter_box_transform": False,
                    "max_action_horizon": 40,
                    "max_state_dim": 132,
                    "max_action_dim": 132,
                    "use_percentiles": True,
                    "use_relative_action": True,
                    "modality_configs": {
                        "libero_sim": {
                            "video": {
                                "delta_indices": [0],
                                "modality_keys": ["image", "wrist_image"],
                            },
                            "state": {
                                "delta_indices": [0],
                                "modality_keys": ["x", "y", "z", "roll", "pitch", "yaw", "gripper"],
                            },
                            "action": {
                                "delta_indices": list(range(16)),
                                "modality_keys": ["x", "y", "z", "roll", "pitch", "yaw", "gripper"],
                            },
                            "language": {
                                "delta_indices": [0],
                                "modality_keys": ["annotation.human.action.task_description"],
                            },
                        }
                    },
                },
            }
        )
    )
    (path / "embodiment_id.json").write_text(json.dumps({"libero_sim": 42}))
    (path / "statistics.json").write_text(
        json.dumps(
            {
                "libero_sim": {
                    "state": {
                        "x": _stats([0.0]),
                        "y": _stats([1.0]),
                        "z": _stats([2.0]),
                        "roll": _stats([3.0]),
                        "pitch": _stats([4.0]),
                        "yaw": _stats([5.0]),
                        "gripper": _stats([6.0, 7.0]),
                    },
                    "action": {
                        "x": _stats([10.0]),
                        "y": _stats([11.0]),
                        "z": _stats([12.0]),
                        "roll": _stats([13.0]),
                        "pitch": _stats([14.0]),
                        "yaw": _stats([15.0]),
                        "gripper": _stats([16.0]),
                    },
                    "relative_action": {},
                }
            }
        )
    )


def _stats(values):
    return {
        "min": values,
        "max": [value + 100.0 for value in values],
        "mean": [value + 50.0 for value in values],
        "std": [1.0 for _ in values],
        "q01": [value + 1.0 for value in values],
        "q99": [value + 99.0 for value in values],
    }


def _expected_albumentations_eval_image(image_np, cv2, *, target_size, shortest_edge, crop_fraction):
    del target_size

    def resize_shortest_edge(frame):
        height, width = frame.shape[:2]
        scale = shortest_edge / float(min(height, width))
        resized_height = max(1, int(round(height * scale)))
        resized_width = max(1, int(round(width * scale)))
        return cv2.resize(frame, (resized_width, resized_height), interpolation=cv2.INTER_AREA)

    image_np = resize_shortest_edge(image_np)
    height, width = image_np.shape[:2]
    crop_h = max(1, int(height * crop_fraction))
    crop_w = max(1, int(width * crop_fraction))
    top = (height - crop_h) // 2
    left = (width - crop_w) // 2
    return resize_shortest_edge(image_np[top : top + crop_h, left : left + crop_w])


class _DummyGrootModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = nn.Parameter(torch.zeros(()))
        self.config = SimpleNamespace(compute_dtype="float32")
        self.compute_dtype = "float32"
        self.forward_inputs = None
        self.get_action_options = None

    def forward(self, inputs):
        self.forward_inputs = dict(inputs)
        return {"loss": self.weight + 1.0}

    def get_action(self, inputs, options=None):
        self.forward_inputs = dict(inputs)
        self.get_action_options = options
        batch_size = inputs["state"].shape[0]
        return {"action_pred": torch.zeros(batch_size, 40, 132, device=self.weight.device)}


def test_groot_defaults_use_n1_7():
    config = GrootConfig(device="cpu")

    assert config.base_model_path == GROOT_N1_7_BASE_MODEL
    assert config.max_state_dim == 132
    assert config.max_action_dim == 132
    assert config.chunk_size == 40
    assert config.n_action_steps == 40
    assert len(config.action_delta_indices) == 40


@pytest.mark.parametrize("legacy_version", ["n1.5", "n1_5", "n15", "1.5"])
def test_groot_normalize_model_version_rejects_n1_5_aliases(legacy_version):
    # model_version is no longer a GrootConfig field, but normalize_groot_model_version is still
    # live (e.g. via infer_groot_model_version) and must keep rejecting N1.5 with removal guidance.
    with pytest.raises(ValueError, match="Unsupported GR00T model_version"):
        normalize_groot_model_version(legacy_version)


def test_groot_normalize_model_version_accepts_n1_7():
    assert normalize_groot_model_version(GROOT_N1_7) == GROOT_N1_7


def test_groot_n1_7_accepts_named_action_decode_transform():
    config = GrootConfig(
        action_decode_transform="libero",
        device="cpu",
    )

    assert config.action_decode_transform == GROOT_ACTION_DECODE_TRANSFORM_LIBERO


@pytest.mark.parametrize("legacy_transform", ["libero_gripper", "libero-gripper"])
def test_groot_n1_7_rejects_legacy_libero_gripper_action_decode_transform(legacy_transform):
    with pytest.raises(ValueError, match="Unsupported GR00T N1.7 action decode transform"):
        GrootConfig(
            action_decode_transform=legacy_transform,
            device="cpu",
        )


def test_groot_config_rejects_mismatched_n1_5_path_for_n1_7():
    with pytest.raises(ValueError, match="does not match base_model_path"):
        GrootConfig(
            base_model_path="nvidia/GR00T-N1.5-3B",
            device="cpu",
        )


def test_groot_n1_7_can_be_selected_from_policy_config_factory_without_external_gr00t():
    sys.modules.pop("gr00t", None)

    config = make_policy_config("groot", device="cpu")

    assert isinstance(config, GrootConfig)
    assert "gr00t" not in sys.modules


def test_groot_predict_action_chunk_accepts_rtc_kwargs():
    signature = inspect.signature(GrootPolicy.predict_action_chunk)

    assert any(parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in signature.parameters.values())
    signature.bind(object(), {}, inference_delay=2, prev_chunk_left_over=None)


def test_groot_predict_action_chunk_forwards_n1_7_rtc_prefix(monkeypatch):
    pytest.importorskip("transformers")

    from lerobot.policies.groot.groot_n1_7 import GR00TN17

    dummy_model = _DummyGrootModel()
    monkeypatch.setattr(GR00TN17, "from_pretrained", classmethod(lambda cls, **kwargs: dummy_model))
    config = _groot_config()
    policy = GrootPolicy(config)
    policy.config.rtc_config = SimpleNamespace(execution_horizon=6)

    prev_chunk = torch.arange(8 * 7, dtype=torch.float32).view(8, 7)

    actions = policy.predict_action_chunk(
        {"state": torch.zeros(1, 1, 132)},
        inference_delay=3,
        prev_chunk_left_over=prev_chunk,
    )

    assert actions.shape == (1, 40, 7)
    assert dummy_model.get_action_options == {
        "action_horizon": 8,
        "rtc_overlap_steps": 6,
        "rtc_frozen_steps": 3,
        "rtc_ramp_rate": 6.0,
    }
    assert dummy_model.forward_inputs["action"].shape == (1, 8, 132)
    torch.testing.assert_close(dummy_model.forward_inputs["action"][0, :, :7], prev_chunk)
    torch.testing.assert_close(dummy_model.forward_inputs["action"][0, :, 7:], torch.zeros(8, 125))


def test_groot_predict_action_chunk_strips_padded_n1_7_rtc_prefix(monkeypatch):
    pytest.importorskip("transformers")

    from lerobot.policies.groot.groot_n1_7 import GR00TN17

    dummy_model = _DummyGrootModel()
    monkeypatch.setattr(GR00TN17, "from_pretrained", classmethod(lambda cls, **kwargs: dummy_model))
    config = _groot_config()
    policy = GrootPolicy(config)
    policy.config.rtc_config = SimpleNamespace(execution_horizon=6)

    prev_chunk = torch.cat(
        (
            torch.arange(4 * 7, dtype=torch.float32).view(4, 7) + 1.0,
            torch.zeros(2, 7),
        )
    )

    policy.predict_action_chunk(
        {"state": torch.zeros(1, 1, 132)},
        inference_delay=5,
        prev_chunk_left_over=prev_chunk,
    )

    assert dummy_model.get_action_options == {
        "action_horizon": 4,
        "rtc_overlap_steps": 4,
        "rtc_frozen_steps": 4,
        "rtc_ramp_rate": 6.0,
    }
    assert dummy_model.forward_inputs["action"].shape == (1, 4, 132)
    torch.testing.assert_close(dummy_model.forward_inputs["action"][0, :, :7], prev_chunk[:4])
    torch.testing.assert_close(dummy_model.forward_inputs["action"][0, :, 7:], torch.zeros(4, 125))


def test_groot_n1_7_predict_action_chunk_truncates_to_checkpoint_valid_horizon(tmp_path, monkeypatch):
    pytest.importorskip("transformers")

    from lerobot.policies.groot.groot_n1_7 import GR00TN17

    model_path = tmp_path / "libero_spatial"
    _write_raw_n1_7_libero_checkpoint(model_path)

    class HorizonModel(_DummyGrootModel):
        def get_action(self, inputs, options=None):
            del options
            batch_size = inputs["state"].shape[0]
            steps = torch.arange(40, dtype=torch.float32).view(1, 40, 1).expand(batch_size, 40, 132)
            return {"action_pred": steps}

    monkeypatch.setattr(GR00TN17, "from_pretrained", classmethod(lambda cls, **kwargs: HorizonModel()))
    input_features, output_features = _groot_features(state_dim=8, action_dim=7)
    config = GrootConfig(
        base_model_path=str(model_path),
        embodiment_tag="libero_sim",
        input_features=input_features,
        output_features=output_features,
        device="cpu",
        use_bf16=False,
        chunk_size=40,
        n_action_steps=40,
    )
    policy = GrootPolicy(config)

    actions = policy.predict_action_chunk({"state": torch.zeros(1, 1, 132)})

    assert actions.shape == (1, 16, 7)
    torch.testing.assert_close(actions[0, :, 0], torch.arange(16, dtype=torch.float32))


def test_groot_from_pretrained_rejects_mismatched_caller_config(tmp_path):
    model_path = tmp_path / "GR00T-N1.7-local"
    model_path.mkdir()
    input_features, output_features = _groot_features(state_dim=8, action_dim=7)

    # An N1.7 config paired with a legacy N1.5 base path is a mismatch and must be
    # rejected. The mismatch is detected during config validation (__post_init__),
    # so construction itself raises before from_pretrained is reached.
    with pytest.raises(ValueError, match="does not match base_model_path"):
        config = GrootConfig(
            base_model_path="nvidia/GR00T-N1.5-3B",
            input_features=input_features,
            output_features=output_features,
            device="cpu",
            use_bf16=False,
            action_decode_transform=GROOT_ACTION_DECODE_TRANSFORM_LIBERO,
        )
        GrootPolicy.from_pretrained(model_path, config=config)


def test_groot_from_pretrained_keeps_matching_caller_config(tmp_path, monkeypatch):
    pytest.importorskip("transformers")

    from lerobot.policies.groot.groot_n1_7 import GR00TN17

    model_path = tmp_path / "GR00T-N1.7-local"
    model_path.mkdir()
    config = _groot_config()

    monkeypatch.setattr(GR00TN17, "from_pretrained", classmethod(lambda cls, **kwargs: _DummyGrootModel()))

    policy = GrootPolicy.from_pretrained(model_path, config=config)

    assert policy.config.base_model_path == str(model_path)


def test_groot_from_pretrained_infers_n1_7_from_ambiguous_local_config(tmp_path, monkeypatch):
    pytest.importorskip("transformers")

    from lerobot.policies.groot.groot_n1_7 import GR00TN17

    model_path = tmp_path / "local-checkpoint"
    model_path.mkdir()
    (model_path / "config.json").write_text('{"model_type": "Gr00tN1d7"}')

    monkeypatch.setattr(GR00TN17, "from_pretrained", classmethod(lambda cls, **kwargs: _DummyGrootModel()))

    policy = GrootPolicy.from_pretrained(model_path)

    assert policy.config.base_model_path == str(model_path)


def test_raw_n1_7_libero_checkpoint_processors_use_checkpoint_assets(tmp_path):
    model_path = tmp_path / "libero_spatial"
    _write_raw_n1_7_libero_checkpoint(model_path)
    config = _raw_n1_7_libero_config(model_path)

    preprocessor, postprocessor = make_pre_post_processors(config, pretrained_path=str(model_path))

    pack_inputs = next(step for step in preprocessor.steps if isinstance(step, GrootN17PackInputsStep))
    vlm_encode = next(step for step in preprocessor.steps if isinstance(step, GrootN17VLMEncodeStep))
    decode_actions = next(step for step in postprocessor.steps if isinstance(step, GrootN17ActionDecodeStep))

    assert pack_inputs.embodiment_tag == "libero_sim"
    assert pack_inputs.embodiment_mapping["libero_sim"] == 42
    assert pack_inputs.formalize_language is True
    assert pack_inputs.valid_action_horizon == 16
    assert pack_inputs.action_horizon == 40
    assert pack_inputs.max_state_dim == 132
    assert pack_inputs.max_action_dim == 132
    assert pack_inputs.clip_outliers is True
    assert pack_inputs.video_modality_keys == ["image", "wrist_image"]
    assert pack_inputs.stats[OBS_STATE]["min"] == [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]
    assert pack_inputs.stats[OBS_STATE]["max"] == [
        99.0,
        100.0,
        101.0,
        102.0,
        103.0,
        104.0,
        105.0,
        106.0,
    ]
    assert pack_inputs.stats[ACTION]["min"] == [11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0]
    assert vlm_encode.image_crop_size == [230, 230]
    assert vlm_encode.image_target_size == [256, 256]
    assert vlm_encode.shortest_image_edge == 256
    assert vlm_encode.crop_fraction == 0.95
    assert vlm_encode.use_albumentations is True
    assert vlm_encode.letter_box_transform is False
    assert decode_actions.raw_stats["action"]["gripper"]["q99"] == [115.0]
    assert decode_actions.env_action_dim == 7
    assert decode_actions.use_percentiles is True
    assert decode_actions.use_relative_action is True
    assert decode_actions.action_decode_transform == GROOT_ACTION_DECODE_TRANSFORM_LIBERO


def test_raw_n1_7_checkpoint_requires_percentile_stats_when_config_uses_percentiles(tmp_path):
    model_path = tmp_path / "libero_spatial"
    _write_raw_n1_7_libero_checkpoint(model_path)
    statistics = json.loads((model_path / "statistics.json").read_text())
    del statistics["libero_sim"]["state"]["x"]["q01"]
    (model_path / "statistics.json").write_text(json.dumps(statistics))
    config = _raw_n1_7_libero_config(model_path)

    with pytest.raises(KeyError, match="q01.*state.x"):
        make_pre_post_processors(config, pretrained_path=str(model_path))


def test_raw_n1_7_checkpoint_processors_prefer_checkpoint_stats_when_dataset_stats_supplied(tmp_path):
    model_path = tmp_path / "libero_spatial"
    _write_raw_n1_7_libero_checkpoint(model_path)
    config = _raw_n1_7_libero_config(model_path)
    dataset_stats = {
        OBS_STATE: {
            "min": torch.full((8,), -8.0),
            "max": torch.full((8,), 8.0),
        },
        ACTION: {
            "min": torch.full((7,), -7.0),
            "max": torch.full((7,), 7.0),
        },
    }

    preprocessor, postprocessor = make_pre_post_processors(
        config,
        pretrained_path=str(model_path),
        dataset_stats=dataset_stats,
    )

    pack_inputs = next(step for step in preprocessor.steps if isinstance(step, GrootN17PackInputsStep))
    decode_actions = next(step for step in postprocessor.steps if isinstance(step, GrootN17ActionDecodeStep))
    torch.testing.assert_close(
        torch.as_tensor(pack_inputs.stats[OBS_STATE]["min"]),
        torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]),
    )
    torch.testing.assert_close(
        torch.as_tensor(pack_inputs.stats[ACTION]["max"]),
        torch.tensor([109.0, 110.0, 111.0, 112.0, 113.0, 114.0, 115.0]),
    )
    assert decode_actions.raw_stats["action"]["gripper"]["q99"] == [115.0]
    assert decode_actions.action_decode_transform == GROOT_ACTION_DECODE_TRANSFORM_LIBERO


def test_groot_n1_7_saved_processors_round_trip_checkpoint_specific_fields(tmp_path):
    model_path = tmp_path / "libero_spatial"
    _write_raw_n1_7_libero_checkpoint(model_path)
    config = _raw_n1_7_libero_config(model_path)
    preprocessor, postprocessor = make_pre_post_processors(config, pretrained_path=str(model_path))
    save_dir = tmp_path / "saved_processors"

    preprocessor.save_pretrained(save_dir)
    postprocessor.save_pretrained(save_dir)

    loaded_preprocessor = PolicyProcessorPipeline.from_pretrained(
        save_dir,
        config_filename="policy_preprocessor.json",
    )
    loaded_postprocessor = PolicyProcessorPipeline.from_pretrained(
        save_dir,
        config_filename="policy_postprocessor.json",
    )
    pack_inputs = next(step for step in loaded_preprocessor.steps if isinstance(step, GrootN17PackInputsStep))
    vlm_encode = next(step for step in loaded_preprocessor.steps if isinstance(step, GrootN17VLMEncodeStep))
    decode_actions = next(
        step for step in loaded_postprocessor.steps if isinstance(step, GrootN17ActionDecodeStep)
    )

    assert pack_inputs.valid_action_horizon == 16
    assert pack_inputs.action_horizon == 40
    assert pack_inputs.video_modality_keys == ["image", "wrist_image"]
    assert pack_inputs.clip_outliers is True
    assert vlm_encode.letter_box_transform is False
    torch.testing.assert_close(
        pack_inputs.stats[OBS_STATE]["min"],
        torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]),
    )
    assert decode_actions.env_action_dim == 7
    assert decode_actions.action_decode_transform == GROOT_ACTION_DECODE_TRANSFORM_LIBERO
    assert decode_actions.raw_stats["action"]["gripper"]["q99"] == [115.0]


def test_converted_raw_n1_7_processors_load_without_legacy_action_unpack_override(tmp_path):
    model_path = tmp_path / "libero_spatial"
    _write_raw_n1_7_libero_checkpoint(model_path)
    config = _raw_n1_7_libero_config(model_path)
    preprocessor, postprocessor = make_pre_post_processors(config, pretrained_path=str(model_path))
    save_dir = tmp_path / "converted_pretrained_model"

    config.save_pretrained(save_dir)
    preprocessor.save_pretrained(save_dir)
    postprocessor.save_pretrained(save_dir)

    loaded_preprocessor, loaded_postprocessor = make_pre_post_processors(
        config,
        pretrained_path=str(save_dir),
        preprocessor_overrides={"rename_observations_processor": {"rename_map": {}}},
    )

    assert any(isinstance(step, GrootN17PackInputsStep) for step in loaded_preprocessor.steps)
    assert any(isinstance(step, GrootN17ActionDecodeStep) for step in loaded_postprocessor.steps)


def test_converted_raw_n1_7_absolute_action_processors_load_without_relative_steps(tmp_path):
    model_path = tmp_path / "libero_spatial"
    _write_raw_n1_7_libero_checkpoint(model_path)
    config = _raw_n1_7_libero_config(model_path)
    preprocessor, postprocessor = make_pre_post_processors(config, pretrained_path=str(model_path))
    save_dir = tmp_path / "absolute_pretrained_model"

    config.save_pretrained(save_dir)
    preprocessor.save_pretrained(save_dir)
    postprocessor.save_pretrained(save_dir)

    loaded_preprocessor, loaded_postprocessor = make_pre_post_processors(
        config,
        pretrained_path=str(save_dir),
        preprocessor_overrides={"rename_observations_processor": {"rename_map": {}}},
    )

    assert any(isinstance(step, GrootN17PackInputsStep) for step in loaded_preprocessor.steps)
    assert any(isinstance(step, GrootN17ActionDecodeStep) for step in loaded_postprocessor.steps)
    assert not any(isinstance(step, RelativeActionsProcessorStep) for step in loaded_preprocessor.steps)
    assert not any(isinstance(step, AbsoluteActionsProcessorStep) for step in loaded_postprocessor.steps)


def test_converted_raw_n1_7_relative_action_processors_reconnect_after_load(tmp_path):
    model_path = tmp_path / "libero_spatial"
    _write_raw_n1_7_libero_checkpoint(model_path)
    config = _raw_n1_7_libero_config(model_path)
    preprocessor, postprocessor = make_pre_post_processors(config, pretrained_path=str(model_path))
    save_dir = tmp_path / "relative_pretrained_model"
    action_names = [
        "shoulder_pan.pos",
        "shoulder_lift.pos",
        "elbow_flex.pos",
        "wrist_flex.pos",
        "wrist_roll.pos",
        "gripper.pos",
    ]

    config.save_pretrained(save_dir)
    preprocessor.save_pretrained(save_dir)
    postprocessor.save_pretrained(save_dir)

    preprocessor_config_path = save_dir / "policy_preprocessor.json"
    preprocessor_config = json.loads(preprocessor_config_path.read_text())
    preprocessor_config["steps"].insert(
        2,
        {
            "registry_name": "relative_actions_processor",
            "config": {
                "enabled": True,
                "exclude_joints": ["gripper"],
                "action_names": action_names,
            },
        },
    )
    preprocessor_config_path.write_text(json.dumps(preprocessor_config, indent=4) + "\n")

    postprocessor_config_path = save_dir / "policy_postprocessor.json"
    postprocessor_config = json.loads(postprocessor_config_path.read_text())
    postprocessor_config["steps"].insert(
        -1,
        {
            "registry_name": "absolute_actions_processor",
            "config": {"enabled": True},
        },
    )
    postprocessor_config_path.write_text(json.dumps(postprocessor_config, indent=4) + "\n")

    loaded_preprocessor, loaded_postprocessor = make_pre_post_processors(
        config,
        pretrained_path=str(save_dir),
        preprocessor_overrides={"rename_observations_processor": {"rename_map": {}}},
    )

    relative_step = next(
        step for step in loaded_preprocessor.steps if isinstance(step, RelativeActionsProcessorStep)
    )
    absolute_step = next(
        step for step in loaded_postprocessor.steps if isinstance(step, AbsoluteActionsProcessorStep)
    )

    assert relative_step.enabled is True
    assert relative_step.exclude_joints == ["gripper"]
    assert relative_step.action_names == action_names
    assert absolute_step.relative_step is relative_step


def test_groot_n1_7_pack_inputs_rejects_state_dim_above_core_max():
    step = GrootN17PackInputsStep(
        max_state_dim=2,
        max_action_dim=4,
        normalize_min_max=False,
    )
    transition = {
        TransitionKey.OBSERVATION: {
            OBS_STATE: torch.zeros(1, 3),
        },
        TransitionKey.COMPLEMENTARY_DATA: {"task": ["Move"]},
    }

    with pytest.raises(ValueError, match="State dimension 3 exceeds max_state_dim 2"):
        step(transition)


def test_groot_n1_7_pack_inputs_rejects_action_shape_above_core_limits():
    step = GrootN17PackInputsStep(
        action_horizon=2,
        max_state_dim=2,
        max_action_dim=2,
        normalize_min_max=False,
    )
    transition = {
        TransitionKey.OBSERVATION: {
            OBS_STATE: torch.zeros(1, 2),
        },
        TransitionKey.ACTION: torch.zeros(1, 2, 3),
        TransitionKey.COMPLEMENTARY_DATA: {"task": ["Move"]},
    }

    with pytest.raises(ValueError, match="Action dimension 3 exceeds max_action_dim 2"):
        step(transition)

    transition[TransitionKey.ACTION] = torch.zeros(1, 3, 2)
    with pytest.raises(ValueError, match="Action horizon 3 exceeds action_horizon 2"):
        step(transition)


def test_groot_n1_7_pack_inputs_clips_and_masks_only_valid_action_horizon():
    step = GrootN17PackInputsStep(
        action_horizon=40,
        valid_action_horizon=16,
        max_state_dim=4,
        max_action_dim=4,
        normalize_min_max=True,
        clip_outliers=True,
        stats={
            OBS_STATE: {"min": [0.0, 0.0], "max": [1.0, 1.0]},
            ACTION: {"min": [0.0, 0.0], "max": [1.0, 1.0]},
        },
    )
    transition = {
        TransitionKey.OBSERVATION: {
            OBS_STATE: torch.tensor([[2.0, -1.0]]),
        },
        TransitionKey.ACTION: torch.full((1, 16, 2), 1.0),
        TransitionKey.COMPLEMENTARY_DATA: {"task": ["Move"]},
    }

    output = step(transition)

    torch.testing.assert_close(
        output[TransitionKey.OBSERVATION]["state"][0, 0, :2],
        torch.tensor([1.0, -1.0]),
    )
    assert output[TransitionKey.ACTION].shape == (1, 40, 4)
    torch.testing.assert_close(output[TransitionKey.ACTION][0, 16:], torch.zeros(24, 4))
    action_mask = output[TransitionKey.COMPLEMENTARY_DATA]["action_mask"]
    assert action_mask.shape == (1, 40, 4)
    assert action_mask[0, :16, :2].sum().item() == 32
    assert action_mask[0, 16:].sum().item() == 0
    assert action_mask[0, :, 2:].sum().item() == 0


def test_groot_n1_7_pack_inputs_normalizes_state_with_q01_q99_clips_and_pads():
    step = GrootN17PackInputsStep(
        action_horizon=4,
        max_state_dim=6,
        max_action_dim=7,
        normalize_min_max=True,
        clip_outliers=True,
        stats={
            OBS_STATE: {
                "min": [0.0, 10.0, -2.0, 4.0],
                "max": [10.0, 10.0, 2.0, 8.0],
            }
        },
    )
    transition = {
        TransitionKey.OBSERVATION: {
            OBS_STATE: torch.tensor([[5.0, 42.0, -6.0, 10.0]]),
        },
        TransitionKey.COMPLEMENTARY_DATA: {"task": ["Move"]},
    }

    output = step(transition)

    expected = torch.tensor([[[0.0, 0.0, -1.0, 1.0, 0.0, 0.0]]])
    torch.testing.assert_close(output[TransitionKey.OBSERVATION]["state"], expected)


def test_groot_n1_7_libero_open_gripper_state_normalizes_near_core_oracle():
    step = GrootN17PackInputsStep(
        action_horizon=40,
        max_state_dim=132,
        max_action_dim=7,
        normalize_min_max=True,
        clip_outliers=True,
        stats={
            OBS_STATE: {
                "min": [
                    -0.27276572585105896,
                    -0.237214133143425,
                    0.916006326675415,
                    2.779496669769287,
                    -1.3187512159347534,
                    -0.4198998212814331,
                    0.001503719249740243,
                    -0.03989770635962486,
                ],
                "max": [
                    0.1352936029434204,
                    0.362916499376297,
                    1.286232590675354,
                    3.2829697132110596,
                    0.9332759976387024,
                    0.6325722336769104,
                    0.03993396461009979,
                    -0.0016719202976673841,
                ],
            }
        },
    )
    transition = {
        TransitionKey.OBSERVATION: {
            OBS_STATE: torch.tensor(
                [
                    [
                        -0.20846466720104218,
                        0.0,
                        1.1732795238494873,
                        3.1403393745422363,
                        0.0007735038525424898,
                        -0.0892220064997673,
                        0.020833000540733337,
                        -0.020833000540733337,
                    ]
                ]
            ),
        },
        TransitionKey.COMPLEMENTARY_DATA: {"task": ["Move"]},
    }

    output = step(transition)

    normalized = output[TransitionKey.OBSERVATION]["state"][0, 0, :8]
    expected = torch.tensor(
        [
            -0.6848445534706116,
            -0.2094583511352539,
            0.3898160457611084,
            0.4334142208099365,
            0.17185509204864502,
            -0.3716168999671936,
            0.005941033363342285,
            -0.002521216869354248,
        ]
    )
    torch.testing.assert_close(normalized, expected, atol=1e-6, rtol=1e-6)
    assert normalized[6:].abs().max().item() < 0.01


def test_groot_n1_7_pack_inputs_normalizes_action_chunk_per_dimension_before_padding():
    step = GrootN17PackInputsStep(
        action_horizon=5,
        valid_action_horizon=3,
        max_state_dim=4,
        max_action_dim=5,
        normalize_min_max=True,
        clip_outliers=True,
        stats={
            OBS_STATE: {"min": [0.0, 0.0], "max": [1.0, 1.0]},
            ACTION: {
                "min": [-2.0, 10.0, 100.0],
                "max": [2.0, 30.0, 101.0],
            },
        },
    )
    transition = {
        TransitionKey.OBSERVATION: {
            OBS_STATE: torch.tensor([[0.5, 0.5]]),
        },
        TransitionKey.ACTION: torch.tensor(
            [
                [
                    [-2.0, 30.0, 100.25],
                    [0.0, 20.0, 101.0],
                    [2.0, 10.0, 100.0],
                ]
            ]
        ),
        TransitionKey.COMPLEMENTARY_DATA: {"task": ["Move"]},
    }

    output = step(transition)

    expected_actions = torch.tensor(
        [
            [
                [-1.0, 1.0, -0.5, 0.0, 0.0],
                [0.0, 0.0, 1.0, 0.0, 0.0],
                [1.0, -1.0, -1.0, 0.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0],
            ]
        ]
    )
    torch.testing.assert_close(output[TransitionKey.ACTION], expected_actions)
    action_mask = output[TransitionKey.COMPLEMENTARY_DATA]["action_mask"]
    assert action_mask.shape == (1, 5, 5)
    assert action_mask[0, :3, :3].sum().item() == 9
    assert action_mask[0, 3:].sum().item() == 0
    assert action_mask[0, :, 3:].sum().item() == 0


def test_groot_n1_7_pack_inputs_raises_when_relative_groups_cannot_normalize():
    # Relative groups carry per-chunk-timestep stats; if the action horizon exceeds the available
    # stat rows, grouped normalization cannot apply and the flat fallback would silently wrongly scale.
    step = GrootN17PackInputsStep(
        action_horizon=3,
        valid_action_horizon=3,
        max_state_dim=2,
        max_action_dim=2,
        normalize_min_max=True,
        raw_stats={
            "state": {"single_arm": {"min": [0.0, 0.0], "max": [1.0, 1.0]}},
            "action": {"single_arm": {"min": [0.0, 0.0], "max": [1.0, 1.0]}},
            # only one horizon row, but the action chunk has horizon 3
            "relative_action": {"single_arm": {"min": [[-1.0, -1.0]], "max": [[1.0, 1.0]]}},
        },
        modality_config={
            "state": {"modality_keys": ["single_arm"]},
            "action": {
                "modality_keys": ["single_arm"],
                "action_configs": [
                    {"rep": "RELATIVE", "type": "NON_EEF", "format": "DEFAULT", "state_key": None}
                ],
                "delta_indices": [0, 1, 2],
            },
        },
    )
    transition = {
        TransitionKey.OBSERVATION: {OBS_STATE: torch.zeros(1, 2)},
        TransitionKey.ACTION: torch.zeros(1, 3, 2),
        TransitionKey.COMPLEMENTARY_DATA: {"task": ["Move"]},
    }

    with pytest.raises(ValueError, match="could not apply native grouped normalization"):
        step(transition)


def test_groot_n1_7_pack_inputs_trains_native_relative_groups_with_absolute_gripper():
    step = GrootN17PackInputsStep(
        action_horizon=2,
        valid_action_horizon=2,
        max_state_dim=6,
        max_action_dim=6,
        normalize_min_max=True,
        clip_outliers=False,
        stats={
            OBS_STATE: {
                "min": [-100.0, -100.0, -100.0, -100.0, -100.0, 0.0],
                "max": [100.0, 100.0, 100.0, 100.0, 100.0, 100.0],
            },
            ACTION: {
                "min": [-10.0, -10.0, -10.0, -10.0, -10.0, 0.0],
                "max": [10.0, 10.0, 10.0, 10.0, 10.0, 100.0],
            },
        },
        raw_stats={
            "state": {
                "single_arm": {"min": [-100.0] * 5, "max": [100.0] * 5},
                "gripper": {"min": [0.0], "max": [100.0]},
            },
            "action": {
                "single_arm": {"min": [-100.0] * 5, "max": [100.0] * 5},
                "gripper": {"min": [0.0], "max": [100.0]},
            },
            "relative_action": {
                "single_arm": {"min": [-10.0] * 5, "max": [10.0] * 5},
            },
        },
        modality_config={
            "state": {"modality_keys": ["single_arm", "gripper"]},
            "action": {
                "modality_keys": ["single_arm", "gripper"],
                "action_configs": [
                    {"rep": "RELATIVE", "type": "NON_EEF", "format": "DEFAULT", "state_key": None},
                    {"rep": "ABSOLUTE", "type": "NON_EEF", "format": "DEFAULT", "state_key": None},
                ],
                "delta_indices": [0, 1],
            },
        },
    )
    transition = {
        TransitionKey.OBSERVATION: {
            OBS_STATE: torch.tensor([[10.0, 20.0, 30.0, 40.0, 50.0, 25.0]]),
        },
        TransitionKey.ACTION: torch.tensor(
            [
                [
                    [12.0, 18.0, 35.0, 30.0, 55.0, 0.0],
                    [9.0, 21.0, 27.0, 43.0, 50.0, 100.0],
                ]
            ]
        ),
        TransitionKey.COMPLEMENTARY_DATA: {"task": ["Move"]},
    }

    output = step(transition)

    expected_actions = torch.tensor(
        [
            [
                [0.2, -0.2, 0.5, -1.0, 0.5, -1.0],
                [-0.1, 0.1, -0.3, 0.3, 0.0, 1.0],
            ]
        ]
    )
    torch.testing.assert_close(output[TransitionKey.ACTION], expected_actions)


def test_groot_policy_ignores_rtc_leftovers_for_relative_actions():
    policy = object.__new__(GrootPolicy)
    policy.config = SimpleNamespace(use_relative_actions=True)
    policy._warned_native_relative_rtc_prefix_disabled = False
    inputs = {"state": torch.zeros(1, 1, 132)}

    output_inputs, options = policy._prepare_n1_7_rtc_inputs(
        inputs,
        inference_delay=1,
        prev_chunk_left_over=torch.ones(8, 6),
    )

    assert output_inputs is inputs
    assert options is None


def test_groot_n1_7_pack_inputs_adds_inference_action_horizon_mask():
    step = GrootN17PackInputsStep(
        action_horizon=40,
        valid_action_horizon=16,
        max_state_dim=8,
        max_action_dim=7,
        normalize_min_max=False,
    )
    transition = {
        TransitionKey.OBSERVATION: {
            OBS_STATE: torch.zeros(2, 8),
        },
        TransitionKey.COMPLEMENTARY_DATA: {"task": ["Move", "Place"]},
    }

    output = step(transition)

    action_mask = output[TransitionKey.COMPLEMENTARY_DATA]["action_mask"]
    assert action_mask.shape == (2, 40)
    assert action_mask[:, :16].sum().item() == 32
    assert action_mask[:, 16:].sum().item() == 0
    assert output[TransitionKey.COMPLEMENTARY_DATA]["embodiment_id"].dtype == torch.int32


def test_groot_n1_7_pack_inputs_masks_padded_action_horizons():
    step = GrootN17PackInputsStep(
        action_horizon=4,
        valid_action_horizon=4,
        max_state_dim=3,
        max_action_dim=5,
        normalize_min_max=False,
    )
    action = torch.arange(2 * 4 * 3, dtype=torch.float32).view(2, 4, 3)
    action_is_pad = torch.tensor(
        [
            [False, True, False, True],
            [True, False, False, False],
        ]
    )
    transition = {
        TransitionKey.OBSERVATION: {
            OBS_STATE: torch.zeros(2, 3),
        },
        TransitionKey.ACTION: action.clone(),
        TransitionKey.COMPLEMENTARY_DATA: {
            "task": ["Move", "Place"],
            "action_is_pad": action_is_pad,
        },
    }

    output = step(transition)

    expected_valid = (~action_is_pad).float()
    action_mask = output[TransitionKey.COMPLEMENTARY_DATA]["action_mask"]
    assert action_mask.shape == (2, 4, 5)
    torch.testing.assert_close(action_mask[..., :3], expected_valid.unsqueeze(-1).expand(-1, -1, 3))
    assert action_mask[..., 3:].sum().item() == 0

    packed_action = output[TransitionKey.ACTION]
    assert packed_action.shape == (2, 4, 5)
    torch.testing.assert_close(packed_action[0, 0, :3], action[0, 0])
    torch.testing.assert_close(packed_action[0, 2, :3], action[0, 2])
    assert packed_action[0, 1].abs().sum().item() == 0
    assert packed_action[0, 3].abs().sum().item() == 0
    assert packed_action[1, 0].abs().sum().item() == 0


def test_groot_n1_7_pack_inputs_orders_video_by_checkpoint_modality_keys():
    step = GrootN17PackInputsStep(
        normalize_min_max=False,
        video_modality_keys=["image", "wrist_image"],
    )
    transition = {
        TransitionKey.OBSERVATION: {
            f"{OBS_IMAGES}.zz_extra": torch.full((1, 3, 2, 2), 33, dtype=torch.uint8),
            f"{OBS_IMAGES}.image2": torch.full((1, 3, 2, 2), 22, dtype=torch.uint8),
            f"{OBS_IMAGES}.image": torch.full((1, 3, 2, 2), 11, dtype=torch.uint8),
            OBS_STATE: torch.zeros(1, 8),
        },
        TransitionKey.COMPLEMENTARY_DATA: {"task": ["Move"]},
    }

    output = step(transition)

    video = output[TransitionKey.OBSERVATION]["video"]
    assert video.shape == (1, 1, 2, 2, 2, 3)
    assert np.unique(video[0, 0, 0]).tolist() == [11]
    assert np.unique(video[0, 0, 1]).tolist() == [22]
    assert f"{OBS_IMAGES}.zz_extra" not in output[TransitionKey.OBSERVATION]
    assert f"{OBS_IMAGES}.image" not in output[TransitionKey.OBSERVATION]
    assert f"{OBS_IMAGES}.image2" not in output[TransitionKey.OBSERVATION]


def test_groot_n1_7_postprocessor_clips_normalized_action_before_unnormalizing():
    step = GrootActionUnpackUnnormalizeStep(
        env_action_dim=3,
        normalize_min_max=True,
        clip_normalized_action=True,
        stats={
            ACTION: {
                "min": [0.0, 0.0, 0.0],
                "max": [10.0, 10.0, 10.0],
            }
        },
    )
    transition = {
        TransitionKey.ACTION: torch.tensor([[-2.0, 0.0, 2.0]]),
    }

    output = step(transition)

    torch.testing.assert_close(output[TransitionKey.ACTION], torch.tensor([[0.0, 5.0, 10.0]]))


def test_groot_n1_7_action_decode_applies_named_libero_transform_from_modality_key():
    unit_stats = {
        "min": [0.0],
        "max": [1.0],
        "mean": [0.5],
        "std": [1.0],
        "q01": [0.0],
        "q99": [1.0],
    }
    step = GrootN17ActionDecodeStep(
        env_action_dim=3,
        raw_stats={
            "action": {
                "x": unit_stats,
                "gripper": unit_stats,
                "y": unit_stats,
            }
        },
        modality_config={
            "action": {
                "modality_keys": ["x", "gripper", "y"],
                "action_configs": [{}, {}, {}],
            }
        },
        action_decode_transform=GROOT_ACTION_DECODE_TRANSFORM_LIBERO,
    )
    action = torch.tensor(
        [
            [
                [-1.0, -1.0, 1.0],
                [1.0, 0.0, -1.0],
                [0.0, 1.0, 0.0],
            ]
        ]
    )

    output = step({TransitionKey.ACTION: action})

    expected = torch.tensor(
        [
            [
                [0.0, 1.0, 1.0],
                [1.0, -0.0, 0.0],
                [0.5, -1.0, 0.5],
            ]
        ]
    )
    torch.testing.assert_close(output[TransitionKey.ACTION], expected)


def test_groot_n1_7_action_decode_truncates_to_valid_horizon_for_relative_stats():
    arm_min = [[-1.0] * 5 for _ in range(16)]
    arm_max = [[1.0] * 5 for _ in range(16)]
    raw_stats = {
        "state": {
            "single_arm": _stats([0.0] * 5),
            "gripper": _stats([0.0]),
        },
        "action": {
            "single_arm": _stats([0.0] * 5),
            "gripper": {
                "min": [0.0],
                "max": [10.0],
                "mean": [5.0],
                "std": [1.0],
                "q01": [0.0],
                "q99": [10.0],
            },
        },
        "relative_action": {
            "single_arm": {
                "min": arm_min,
                "max": arm_max,
                "mean": [[0.0] * 5 for _ in range(16)],
                "std": [[1.0] * 5 for _ in range(16)],
                "q01": arm_min,
                "q99": arm_max,
            },
        },
    }
    modality_config = {
        "state": {
            "modality_keys": ["single_arm", "gripper"],
        },
        "action": {
            "delta_indices": list(range(16)),
            "modality_keys": ["single_arm", "gripper"],
            "action_configs": [
                {"rep": "RELATIVE", "type": "NON_EEF", "format": "DEFAULT", "state_key": None},
                {"rep": "ABSOLUTE", "type": "NON_EEF", "format": "DEFAULT", "state_key": None},
            ],
        },
    }
    pack_step = GrootN17PackInputsStep(
        raw_stats=raw_stats,
        modality_config=modality_config,
        normalize_min_max=False,
    )
    pack_step(
        {
            TransitionKey.OBSERVATION: {OBS_STATE: torch.zeros(1, 6)},
            TransitionKey.COMPLEMENTARY_DATA: {},
        }
    )
    decode_step = GrootN17ActionDecodeStep(
        env_action_dim=6,
        raw_stats=raw_stats,
        modality_config=modality_config,
        use_relative_action=True,
        pack_step=pack_step,
    )

    output = decode_step({TransitionKey.ACTION: torch.zeros(1, 40, 6)})

    decoded = output[TransitionKey.ACTION]
    assert decoded.shape == (1, 16, 6)
    torch.testing.assert_close(decoded[..., :5], torch.zeros(1, 16, 5))
    torch.testing.assert_close(decoded[..., 5], torch.full((1, 16), 5.0))


def test_groot_n1_7_action_decode_rejects_stepwise_native_relative_actions():
    raw_stats = {
        "state": {
            "single_arm": _stats([0.0] * 5),
            "gripper": _stats([0.0]),
        },
        "action": {
            "single_arm": _stats([0.0] * 5),
            "gripper": _stats([0.0]),
        },
        "relative_action": {
            "single_arm": _stats([0.0] * 5),
        },
    }
    modality_config = {
        "state": {
            "modality_keys": ["single_arm", "gripper"],
        },
        "action": {
            "modality_keys": ["single_arm", "gripper"],
            "action_configs": [
                {"rep": "RELATIVE", "type": "NON_EEF", "format": "DEFAULT", "state_key": None},
                {"rep": "ABSOLUTE", "type": "NON_EEF", "format": "DEFAULT", "state_key": None},
            ],
        },
    }
    pack_step = GrootN17PackInputsStep(
        raw_stats=raw_stats,
        modality_config=modality_config,
        normalize_min_max=False,
    )
    pack_step(
        {
            TransitionKey.OBSERVATION: {OBS_STATE: torch.zeros(1, 6)},
            TransitionKey.COMPLEMENTARY_DATA: {},
        }
    )
    decode_step = GrootN17ActionDecodeStep(
        env_action_dim=6,
        raw_stats=raw_stats,
        modality_config=modality_config,
        use_relative_action=True,
        pack_step=pack_step,
    )

    with pytest.raises(NotImplementedError, match="cannot decode native relative actions one step at a time"):
        decode_step({TransitionKey.ACTION: torch.zeros(1, 6)})


def test_groot_n1_7_action_decode_requires_gripper_key_for_libero_transform():
    step = GrootN17ActionDecodeStep(
        env_action_dim=1,
        raw_stats={
            "action": {
                "x": {
                    "min": [0.0],
                    "max": [1.0],
                },
            }
        },
        modality_config={
            "action": {
                "modality_keys": ["x"],
                "action_configs": [{}],
            }
        },
        action_decode_transform=GROOT_ACTION_DECODE_TRANSFORM_LIBERO,
    )

    with pytest.raises(KeyError, match="gripper"):
        step({TransitionKey.ACTION: torch.zeros(1, 1, 1)})


def test_groot_n1_7_fallback_processors_wire_libero_transform_to_postprocessor():
    config = _groot_config()
    dataset_stats = {
        OBS_STATE: {
            "min": torch.zeros(8),
            "max": torch.ones(8),
        },
        ACTION: {
            "min": torch.zeros(7),
            "max": torch.ones(7),
        },
    }

    _, postprocessor = make_groot_pre_post_processors(config, dataset_stats=dataset_stats)

    action_decode_step = next(
        step for step in postprocessor.steps if isinstance(step, GrootActionUnpackUnnormalizeStep)
    )
    assert action_decode_step.libero_gripper_action is True


def test_groot_n1_7_loaded_fallback_postprocessor_honors_config_action_decode_transform(tmp_path):
    input_features, output_features = _groot_features(state_dim=8, action_dim=7)
    dataset_stats = {
        OBS_STATE: {
            "min": torch.zeros(8),
            "max": torch.ones(8),
        },
        ACTION: {
            "min": torch.zeros(7),
            "max": torch.ones(7),
        },
    }
    disabled_config = GrootConfig(
        input_features=input_features,
        output_features=output_features,
        device="cpu",
        use_bf16=False,
        action_decode_transform=None,
    )
    preprocessor, postprocessor = make_groot_pre_post_processors(
        disabled_config,
        dataset_stats=dataset_stats,
    )
    save_dir = tmp_path / "saved_fallback_processors"
    disabled_config.save_pretrained(save_dir)
    preprocessor.save_pretrained(save_dir)
    postprocessor.save_pretrained(save_dir)

    saved_postprocessor = json.loads((save_dir / "policy_postprocessor.json").read_text())
    saved_decode_config = next(
        step["config"]
        for step in saved_postprocessor["steps"]
        if step["registry_name"] == "groot_action_unpack_unnormalize_v2"
    )
    assert saved_decode_config["libero_gripper_action"] is False

    enabled_config = GrootConfig(
        input_features=input_features,
        output_features=output_features,
        device="cpu",
        use_bf16=False,
        action_decode_transform=GROOT_ACTION_DECODE_TRANSFORM_LIBERO,
    )
    _, loaded_postprocessor = make_pre_post_processors(enabled_config, pretrained_path=str(save_dir))
    action_decode_step = next(
        step for step in loaded_postprocessor.steps if isinstance(step, GrootActionUnpackUnnormalizeStep)
    )

    assert action_decode_step.libero_gripper_action is True
    output = action_decode_step({TransitionKey.ACTION: torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -1.0]])})
    torch.testing.assert_close(output[TransitionKey.ACTION][0, -1], torch.tensor(1.0))


def test_groot_n1_7_postprocessor_converts_libero_gripper_convention():
    step = GrootActionUnpackUnnormalizeStep(
        env_action_dim=7,
        normalize_min_max=True,
        stats={
            ACTION: {
                "min": [0.0] * 7,
                "max": [1.0] * 7,
            }
        },
        libero_gripper_action=True,
    )
    transition = {
        TransitionKey.ACTION: torch.tensor(
            [
                [-1.0, -1.0, -1.0, -1.0, -1.0, -1.0, -1.0],
                [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
            ]
        )
    }

    output = step(transition)

    torch.testing.assert_close(output[TransitionKey.ACTION][:, -1], torch.tensor([1.0, -1.0]))


def test_groot_n1_7_postprocessor_decodes_selected_action_and_gripper_thresholds():
    step = GrootActionUnpackUnnormalizeStep(
        env_action_dim=7,
        normalize_min_max=True,
        clip_normalized_action=True,
        stats={
            ACTION: {
                "min": [0.0, 10.0, 20.0, 30.0, 40.0, 50.0, 0.0],
                "max": [2.0, 14.0, 26.0, 38.0, 50.0, 62.0, 1.0],
            }
        },
        libero_gripper_action=True,
    )
    selected_actions = torch.tensor(
        [
            [-1.0, -0.5, 0.0, 0.5, 1.0, 2.0, -0.5],
            [-1.0, -0.5, 0.0, 0.5, 1.0, 2.0, 0.0],
            [-1.0, -0.5, 0.0, 0.5, 1.0, 2.0, 0.5],
        ]
    )

    output = step({TransitionKey.ACTION: selected_actions})

    expected_prefix = torch.tensor([0.0, 11.0, 23.0, 36.0, 50.0, 62.0])
    torch.testing.assert_close(output[TransitionKey.ACTION][:, :6], expected_prefix.expand(3, 6))
    torch.testing.assert_close(output[TransitionKey.ACTION][:, -1], torch.tensor([1.0, -0.0, -1.0]))


def test_groot_n1_7_postprocessor_decodes_action_chunks_without_dropping_timesteps():
    step = GrootActionUnpackUnnormalizeStep(
        env_action_dim=7,
        normalize_min_max=True,
        clip_normalized_action=True,
        stats={
            ACTION: {
                "min": [0.0, 10.0, 20.0, 30.0, 40.0, 50.0, 0.0],
                "max": [2.0, 14.0, 26.0, 38.0, 50.0, 62.0, 1.0],
            }
        },
        libero_gripper_action=True,
    )
    action_chunk = torch.tensor(
        [
            [
                [-1.0, 0.0, 1.0, -0.5, 0.5, 2.0, -1.0, 99.0],
                [0.25, -0.25, 0.75, -0.75, 1.0, -1.0, 0.0, 99.0],
                [1.0, -1.0, 0.0, 0.5, -0.5, 0.0, 0.5, 99.0],
            ]
        ]
    )

    output = step({TransitionKey.ACTION: action_chunk})

    expected_prefix = torch.tensor(
        [
            [
                [0.0, 12.0, 26.0, 32.0, 47.5, 62.0],
                [1.25, 11.5, 25.25, 31.0, 50.0, 50.0],
                [2.0, 10.0, 23.0, 36.0, 42.5, 56.0],
            ]
        ]
    )
    assert output[TransitionKey.ACTION].shape == (1, 3, 7)
    torch.testing.assert_close(output[TransitionKey.ACTION][..., :6], expected_prefix)
    torch.testing.assert_close(output[TransitionKey.ACTION][..., -1], torch.tensor([[1.0, -0.0, -1.0]]))


def test_groot_from_pretrained_rejects_caller_config_mismatch_from_local_config(tmp_path):
    model_path = tmp_path / "local-checkpoint"
    model_path.mkdir()
    (model_path / "config.json").write_text('{"model_type": "Gr00tN1d7"}')
    input_features, output_features = _groot_features(state_dim=8, action_dim=7)

    # An N1.7 config paired with a legacy N1.5 base path is a mismatch and must be
    # rejected. The mismatch is detected during config validation (__post_init__),
    # so construction itself raises before from_pretrained is reached.
    with pytest.raises(ValueError, match="does not match base_model_path"):
        config = GrootConfig(
            base_model_path="nvidia/GR00T-N1.5-3B",
            input_features=input_features,
            output_features=output_features,
            device="cpu",
            use_bf16=False,
            action_decode_transform=GROOT_ACTION_DECODE_TRANSFORM_LIBERO,
        )
        GrootPolicy.from_pretrained(model_path, config=config)


def test_groot_n1_7_processors_are_registered_lazily_without_external_gr00t():
    sys.modules.pop("gr00t", None)
    config = _groot_config()

    preprocessor, _ = make_groot_pre_post_processors(config)
    step_types = {type(step) for step in preprocessor.steps}

    assert GrootN17PackInputsStep in step_types
    assert GrootN17VLMEncodeStep in step_types
    assert "gr00t" not in sys.modules


def test_groot_n1_7_pack_inputs_preserves_per_sample_language():
    step = GrootN17PackInputsStep(
        action_horizon=2,
        max_state_dim=4,
        max_action_dim=3,
        formalize_language=True,
        normalize_min_max=False,
    )
    transition = {
        TransitionKey.OBSERVATION: {
            OBS_STATE: torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
        },
        TransitionKey.ACTION: torch.tensor([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]),
        TransitionKey.COMPLEMENTARY_DATA: {
            "task": ["Pick Red Block!", "Place Blue Cube."],
        },
    }

    output = step(transition)

    assert output[TransitionKey.COMPLEMENTARY_DATA]["language"] == [
        "pick red block",
        "place blue cube",
    ]
    torch.testing.assert_close(
        output[TransitionKey.OBSERVATION]["state"][:, 0, :2],
        torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
    )


def test_groot_n1_7_language_formalization_preserves_core_task_identifier_and_batch():
    step = GrootN17PackInputsStep(
        action_horizon=2,
        max_state_dim=8,
        max_action_dim=7,
        formalize_language=True,
        normalize_min_max=False,
    )
    transition = {
        TransitionKey.OBSERVATION: {
            OBS_STATE: torch.zeros(2, 8),
        },
        TransitionKey.COMPLEMENTARY_DATA: {
            "task": [
                "Pick_Up_The_Black_Bowl_Next_To_The_Ramekin_And_Place_It_On_The_Plate!!!",
                "MOVE, the YELLOW mug -- to Zone_2.",
            ],
        },
    }

    output = step(transition)

    assert output[TransitionKey.COMPLEMENTARY_DATA]["language"] == [
        "pick_up_the_black_bowl_next_to_the_ramekin_and_place_it_on_the_plate",
        "move the yellow mug  to zone_2",
    ]


def test_groot_n1_7_vlm_encode_uses_per_sample_language():
    class FakeProcessor:
        def __init__(self):
            self.rendered_texts = []
            self.encoded_texts = None

        def apply_chat_template(self, conversation, tokenize, add_generation_prompt):
            content = conversation[0]["content"]
            assert [item["type"] for item in content] == ["image", "text"]
            text = content[-1]["text"]
            self.rendered_texts.append(text)
            return f"rendered:{text}"

        def __call__(self, text, images, return_tensors, padding):
            self.encoded_texts = text
            return {
                "input_ids": torch.arange(len(text)).view(len(text), 1),
                "attention_mask": torch.ones(len(text), 1, dtype=torch.long),
            }

    fake_proc = FakeProcessor()
    step = GrootN17VLMEncodeStep()
    step._proc = fake_proc
    transition = {
        TransitionKey.OBSERVATION: {
            "video": np.zeros((2, 1, 1, 2, 2, 3), dtype=np.uint8),
        },
        TransitionKey.COMPLEMENTARY_DATA: {
            "language": ["first task", "second task"],
        },
    }

    output = step(transition)

    assert fake_proc.rendered_texts == ["first task", "second task"]
    assert fake_proc.encoded_texts == ["rendered:first task", "rendered:second task"]
    assert "video" not in output[TransitionKey.OBSERVATION]
    torch.testing.assert_close(
        output[TransitionKey.COMPLEMENTARY_DATA]["input_ids"],
        torch.tensor([[0], [1]]),
    )


def test_groot_n1_7_vlm_encode_packs_images_time_major_then_camera_order():
    class FakeProcessor:
        def __init__(self):
            self.add_generation_prompts = []
            self.conversation_content_types = []
            self.conversation_image_values = []
            self.conversation_texts = []
            self.encoded_texts = None
            self.encoded_image_values = None

        def apply_chat_template(self, conversation, tokenize, add_generation_prompt):
            assert tokenize is False
            self.add_generation_prompts.append(add_generation_prompt)
            content = conversation[0]["content"]
            self.conversation_content_types.append([item["type"] for item in content])
            self.conversation_image_values.append(
                [int(np.asarray(item["image"])[0, 0, 0]) for item in content if item["type"] == "image"]
            )
            text = content[-1]["text"]
            self.conversation_texts.append(text)
            return f"rendered:{text}"

        def __call__(self, text, images, return_tensors, padding):
            assert return_tensors == "pt"
            assert padding is True
            self.encoded_texts = text
            self.encoded_image_values = [int(np.asarray(image)[0, 0, 0]) for image in images]
            return {
                "input_ids": torch.arange(len(text)).view(len(text), 1),
                "attention_mask": torch.ones(len(text), 1, dtype=torch.long),
                "pixel_values": torch.arange(len(images)).view(len(images), 1),
                "image_grid_thw": torch.ones(len(images), 3, dtype=torch.long),
            }

    fake_proc = FakeProcessor()
    step = GrootN17VLMEncodeStep()
    step._proc = fake_proc
    video = np.zeros((2, 2, 2, 2, 2, 3), dtype=np.uint8)
    image_id = 1
    for batch_idx in range(2):
        for timestep in range(2):
            for view_idx in range(2):
                video[batch_idx, timestep, view_idx, :, :, :] = image_id
                image_id += 1
    transition = {
        TransitionKey.OBSERVATION: {"video": video},
        TransitionKey.COMPLEMENTARY_DATA: {"language": ["task a", "task b"]},
    }

    output = step(transition)

    assert fake_proc.conversation_image_values == [[1, 2, 3, 4], [5, 6, 7, 8]]
    assert fake_proc.conversation_content_types == [
        ["image", "image", "image", "image", "text"],
        ["image", "image", "image", "image", "text"],
    ]
    assert fake_proc.encoded_image_values == [1, 2, 3, 4, 5, 6, 7, 8]
    assert fake_proc.conversation_texts == ["task a", "task b"]
    assert fake_proc.encoded_texts == ["rendered:task a", "rendered:task b"]
    assert fake_proc.add_generation_prompts == [False, False]
    assert "video" not in output[TransitionKey.OBSERVATION]
    assert set(output[TransitionKey.COMPLEMENTARY_DATA]) >= {
        "input_ids",
        "attention_mask",
        "pixel_values",
        "image_grid_thw",
    }


def test_groot_n1_7_vlm_image_transform_matches_albumentations_eval_path():
    cv2 = pytest.importorskip("cv2", exc_type=ImportError)
    from PIL import Image

    image_np = (np.arange(360 * 360 * 3, dtype=np.uint32) % 251).astype(np.uint8).reshape(360, 360, 3)

    transformed = _transform_n1_7_image_for_vlm_albumentations(
        Image.fromarray(image_np),
        image_crop_size=[230, 230],
        image_target_size=[256, 256],
        shortest_image_edge=256,
        crop_fraction=0.95,
    )

    expected = cv2.resize(image_np, (256, 256), interpolation=cv2.INTER_AREA)
    crop_edge = int(256 * 0.95)
    crop_start = (256 - crop_edge) // 2
    expected = expected[crop_start : crop_start + crop_edge, crop_start : crop_start + crop_edge]
    expected = cv2.resize(expected, (256, 256), interpolation=cv2.INTER_AREA)

    assert transformed.shape == (256, 256, 3)
    np.testing.assert_array_equal(np.asarray(transformed), expected)


def test_groot_n1_7_albumentations_letterbox_is_opt_in():
    pytest.importorskip("cv2", exc_type=ImportError)

    image = np.full((3, 5, 3), 255, dtype=np.uint8)

    default = _transform_n1_7_image_for_vlm_albumentations(
        image,
        image_crop_size=None,
        image_target_size=[10, 10],
        shortest_image_edge=10,
        crop_fraction=None,
    )
    letterboxed = _transform_n1_7_image_for_vlm_albumentations(
        image,
        image_crop_size=None,
        image_target_size=[10, 10],
        shortest_image_edge=10,
        crop_fraction=None,
        letter_box_transform=True,
    )

    assert default.shape == (10, 17, 3)
    assert default.min() == 255
    assert letterboxed.shape == (10, 10, 3)
    assert letterboxed.min() < 255


def test_groot_n1_7_torch_letterbox_is_opt_in():
    image = torch.full((3, 3, 5), 255, dtype=torch.uint8)

    default = _transform_n1_7_image_for_vlm_torch(
        image,
        image_crop_size=None,
        image_target_size=[10, 10],
        shortest_image_edge=10,
        crop_fraction=None,
    )
    letterboxed = _transform_n1_7_image_for_vlm_torch(
        image,
        image_crop_size=None,
        image_target_size=[10, 10],
        shortest_image_edge=10,
        crop_fraction=None,
        letter_box_transform=True,
    )

    assert tuple(default.shape) == (3, 10, 10)
    assert int(default.min()) == 255
    assert tuple(letterboxed.shape) == (3, 10, 10)
    assert int(letterboxed.min()) < 255


def test_groot_n1_7_vlm_encode_transforms_non_square_two_camera_sample_like_core_albumentations():
    cv2 = pytest.importorskip("cv2", exc_type=ImportError)

    class FakeProcessor:
        def __init__(self):
            self.images = None

        def apply_chat_template(self, conversation, tokenize, add_generation_prompt):
            content = conversation[0]["content"]
            assert [item["type"] for item in content] == ["image", "image", "text"]
            return content[-1]["text"]

        def __call__(self, text, images, return_tensors, padding):
            self.images = images
            return {
                "input_ids": torch.ones(len(text), 1, dtype=torch.long),
                "attention_mask": torch.ones(len(text), 1, dtype=torch.long),
            }

    camera_a = np.arange(3 * 5 * 3, dtype=np.uint8).reshape(3, 5, 3)
    camera_b = (np.arange(3 * 5 * 3, dtype=np.uint16).reshape(3, 5, 3) * 3 % 251).astype(np.uint8)
    video = np.stack([camera_a, camera_b], axis=0).reshape(1, 1, 2, 3, 5, 3)
    fake_proc = FakeProcessor()
    step = GrootN17VLMEncodeStep(
        image_target_size=[8, 8],
        shortest_image_edge=10,
        crop_fraction=0.6,
        use_albumentations=True,
    )
    step._proc = fake_proc

    step(
        {
            TransitionKey.OBSERVATION: {"video": video},
            TransitionKey.COMPLEMENTARY_DATA: {"language": ["move"]},
        }
    )

    assert fake_proc.images is not None
    assert len(fake_proc.images) == 2
    np.testing.assert_array_equal(
        np.asarray(fake_proc.images[0]),
        _expected_albumentations_eval_image(
            camera_a,
            cv2,
            target_size=[8, 8],
            shortest_edge=10,
            crop_fraction=0.6,
        ),
    )
    np.testing.assert_array_equal(
        np.asarray(fake_proc.images[1]),
        _expected_albumentations_eval_image(
            camera_b,
            cv2,
            target_size=[8, 8],
            shortest_edge=10,
            crop_fraction=0.6,
        ),
    )


def test_groot_n1_7_vlm_encode_config_round_trips_model_name():
    step = GrootN17VLMEncodeStep(
        model_name="local-cosmos",
        image_crop_size=[230, 230],
        image_target_size=[256, 256],
        shortest_image_edge=256,
        crop_fraction=0.95,
        use_albumentations=True,
        letter_box_transform=True,
    )

    restored = GrootN17VLMEncodeStep(**step.get_config())

    assert restored.model_name == "local-cosmos"
    assert restored.image_crop_size == [230, 230]
    assert restored.image_target_size == [256, 256]
    assert restored.shortest_image_edge == 256
    assert restored.crop_fraction == 0.95
    assert restored.use_albumentations is True
    assert restored.letter_box_transform is True


def test_groot_n1_7_processor_uses_qwen_component_assets(monkeypatch):
    pytest.importorskip("transformers")

    from lerobot.policies.groot import processor_groot

    calls = []

    class FakeTokenizer:
        chat_template = "fake-chat-template"
        padding_side = "right"

        @classmethod
        def from_pretrained(cls, model_name, **kwargs):
            calls.append(("tokenizer", model_name, kwargs))
            return cls()

    class FakeImageProcessor:
        @classmethod
        def from_pretrained(cls, model_name, **kwargs):
            calls.append(("image_processor", model_name, kwargs))
            return cls()

    class FakeVideoProcessor:
        @classmethod
        def from_pretrained(cls, model_name, **kwargs):
            calls.append(("video_processor", model_name, kwargs))
            return cls()

    class FakeProcessor:
        from_pretrained_called = False

        def __init__(self, *, image_processor, tokenizer, video_processor, chat_template):
            self.image_processor = image_processor
            self.tokenizer = tokenizer
            self.video_processor = video_processor
            self.chat_template = chat_template

        @classmethod
        def from_pretrained(cls, *args, **kwargs):
            cls.from_pretrained_called = True
            raise AssertionError("Cosmos does not publish processor_config.json")

    monkeypatch.setattr(processor_groot, "AutoTokenizer", FakeTokenizer)
    monkeypatch.setattr(processor_groot, "Qwen2VLImageProcessor", FakeImageProcessor)
    monkeypatch.setattr(processor_groot, "Qwen3VLVideoProcessor", FakeVideoProcessor)
    monkeypatch.setattr(processor_groot, "Qwen3VLProcessor", FakeProcessor)

    processor = processor_groot._build_n1_7_processor("nvidia/Cosmos-Reason2-2B")

    assert [call[:2] for call in calls] == [
        ("tokenizer", "nvidia/Cosmos-Reason2-2B"),
        ("image_processor", "nvidia/Cosmos-Reason2-2B"),
        ("video_processor", "nvidia/Cosmos-Reason2-2B"),
    ]
    assert all(call[2] == {"trust_remote_code": True} for call in calls)
    assert processor.tokenizer.padding_side == "left"
    assert processor.chat_template == "fake-chat-template"
    assert not FakeProcessor.from_pretrained_called


def test_groot_n1_7_saved_processors_reload_through_factory(tmp_path):
    config = _groot_config()
    dataset_stats = {
        OBS_STATE: {
            "min": torch.zeros(8),
            "max": torch.ones(8),
        },
        ACTION: {
            "min": torch.zeros(7),
            "max": torch.ones(7),
        },
    }
    preprocessor, postprocessor = make_groot_pre_post_processors(config, dataset_stats=dataset_stats)
    preprocessor.save_pretrained(tmp_path)
    postprocessor.save_pretrained(tmp_path)

    loaded_preprocessor, loaded_postprocessor = make_pre_post_processors(
        config,
        pretrained_path=str(tmp_path),
        dataset_stats=dataset_stats,
    )

    pack_step = next(step for step in loaded_preprocessor.steps if isinstance(step, GrootN17PackInputsStep))
    unpack_step = loaded_postprocessor.steps[0]
    assert pack_step.normalize_min_max
    torch.testing.assert_close(pack_step.stats[OBS_STATE]["min"], dataset_stats[OBS_STATE]["min"])
    torch.testing.assert_close(pack_step.stats[ACTION]["max"], dataset_stats[ACTION]["max"])
    torch.testing.assert_close(unpack_step.stats[OBS_STATE]["min"], dataset_stats[OBS_STATE]["min"])
    torch.testing.assert_close(unpack_step.stats[ACTION]["max"], dataset_stats[ACTION]["max"])
    assert unpack_step.env_action_dim == 7


def test_groot_n1_7_saved_processors_reload_through_factory_preserves_saved_stats(tmp_path):
    config = _groot_config()
    saved_stats = {
        OBS_STATE: {
            "min": torch.full((8,), -2.0),
            "max": torch.full((8,), 2.0),
        },
        ACTION: {
            "min": torch.full((7,), -3.0),
            "max": torch.full((7,), 3.0),
        },
    }
    preprocessor, postprocessor = make_groot_pre_post_processors(config, dataset_stats=saved_stats)
    preprocessor.save_pretrained(tmp_path)
    postprocessor.save_pretrained(tmp_path)

    loaded_preprocessor, loaded_postprocessor = make_pre_post_processors(
        config,
        pretrained_path=str(tmp_path),
    )

    pack_step = next(step for step in loaded_preprocessor.steps if isinstance(step, GrootN17PackInputsStep))
    unpack_step = loaded_postprocessor.steps[0]
    assert pack_step.normalize_min_max
    torch.testing.assert_close(pack_step.stats[OBS_STATE]["min"], saved_stats[OBS_STATE]["min"])
    torch.testing.assert_close(pack_step.stats[ACTION]["max"], saved_stats[ACTION]["max"])
    torch.testing.assert_close(unpack_step.stats[OBS_STATE]["min"], saved_stats[OBS_STATE]["min"])
    torch.testing.assert_close(unpack_step.stats[ACTION]["max"], saved_stats[ACTION]["max"])
    assert unpack_step.env_action_dim == 7


def test_groot_n1_7_relative_action_training_processors_save_native_grouped_stats(tmp_path):
    input_features, output_features = _groot_features(state_dim=6, action_dim=6)
    action_names = [
        "shoulder_pan.pos",
        "shoulder_lift.pos",
        "elbow_flex.pos",
        "wrist_flex.pos",
        "wrist_roll.pos",
        "gripper.pos",
    ]
    config = GrootConfig(
        input_features=input_features,
        output_features=output_features,
        device="cpu",
        use_bf16=False,
        action_decode_transform=None,
        use_relative_actions=True,
        relative_exclude_joints=["gripper"],
    )
    absolute_dataset_stats = {
        OBS_STATE: {
            "min": torch.tensor([-50.0, -60.0, -70.0, -80.0, -90.0, 0.0]),
            "max": torch.tensor([50.0, 60.0, 70.0, 80.0, 90.0, 100.0]),
        },
        ACTION: {
            "min": torch.tensor([-100.0, -110.0, -120.0, -130.0, -140.0, 0.0]),
            "max": torch.tensor([100.0, 110.0, 120.0, 130.0, 140.0, 100.0]),
        },
    }
    samples = [
        {
            OBS_STATE: torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0, 0.0]),
            ACTION: _native_action_chunk(
                [
                    [8.0, 17.0, 26.0, 35.0, 44.0, 0.0],
                    [12.0, 23.0, 34.0, 45.0, 56.0, 100.0],
                ]
            ),
        },
        {
            OBS_STATE: torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0, 50.0]),
            ACTION: _native_action_chunk(
                [
                    [-1.0, -2.0, -3.0, -4.0, -5.0, 25.0],
                    [1.0, 2.0, 3.0, 4.0, 5.0, 75.0],
                ]
            ),
        },
    ]

    class _RelativeStatsDataset:
        meta = SimpleNamespace(
            stats=absolute_dataset_stats,
            features={ACTION: {"names": action_names}},
        )

        def __len__(self):
            return len(samples)

        def __getitem__(self, idx):
            return samples[idx]

    relative_dataset_stats = _make_relative_action_training_stats(
        _RelativeStatsDataset(),
        exclude_joints=["gripper"],
        action_names=action_names,
        preserve_action_horizon=True,
    )
    expected_relative_action_min_prefix = torch.tensor(
        [-2.0, -3.0, -4.0, -5.0, -6.0, 1.0, 2.0, 3.0, 4.0, 5.0]
    )
    expected_relative_action_max_prefix = torch.tensor(
        [-1.0, -2.0, -3.0, -4.0, -5.0, 2.0, 3.0, 4.0, 5.0, 6.0]
    )

    preprocessor, postprocessor = make_groot_pre_post_processors(
        config, dataset_stats=relative_dataset_stats, dataset_meta=_RelativeStatsDataset.meta
    )
    preprocessor.save_pretrained(tmp_path)
    postprocessor.save_pretrained(tmp_path)

    preprocessor_config = json.loads((tmp_path / "policy_preprocessor.json").read_text())
    assert not any(
        step.get("registry_name") == "relative_actions_processor" for step in preprocessor_config["steps"]
    )
    pack_entry = next(
        step
        for step in preprocessor_config["steps"]
        if step.get("registry_name") == "groot_n1_7_pack_inputs_v1"
    )
    pack_config = pack_entry["config"]
    assert pack_config["modality_config"]["action"]["modality_keys"] == ["single_arm", "gripper"]
    assert pack_config["modality_config"]["action"]["action_configs"] == [
        {"rep": "RELATIVE", "type": "NON_EEF", "format": "DEFAULT", "state_key": None},
        {"rep": "ABSOLUTE", "type": "NON_EEF", "format": "DEFAULT", "state_key": None},
    ]
    pack_relative_min = pack_config["raw_stats"]["relative_action"]["single_arm"]["min"]
    assert pack_relative_min[:2] == [
        [-2.0, -3.0, -4.0, -5.0, -6.0],
        [1.0, 2.0, 3.0, 4.0, 5.0],
    ]
    assert len(pack_relative_min) == N1_7_NATIVE_ACTION_HORIZON
    assert (
        pack_config["raw_stats"]["relative_action"]["single_arm"]["count"] == [2] * N1_7_NATIVE_ACTION_HORIZON
    )
    assert pack_config["raw_stats"]["action"]["gripper"]["min"] == [0.0]
    assert pack_config["raw_stats"]["action"]["gripper"]["max"] == [100.0]

    pack_state = load_file(tmp_path / pack_entry["state_file"])
    expected_flat_dim = N1_7_NATIVE_ACTION_HORIZON * 5 + 1
    assert pack_state[f"{ACTION}.min"].shape == (expected_flat_dim,)
    assert pack_state[f"{ACTION}.max"].shape == (expected_flat_dim,)
    torch.testing.assert_close(pack_state[f"{ACTION}.min"][:10], expected_relative_action_min_prefix)
    torch.testing.assert_close(pack_state[f"{ACTION}.max"][:10], expected_relative_action_max_prefix)
    assert pack_state[f"{ACTION}.min"][-1].item() == 0.0
    assert pack_state[f"{ACTION}.max"][-1].item() == 100.0

    postprocessor_config = json.loads((tmp_path / "policy_postprocessor.json").read_text())
    assert not any(
        step.get("registry_name") == "absolute_actions_processor" for step in postprocessor_config["steps"]
    )
    decode_entry = next(
        step
        for step in postprocessor_config["steps"]
        if step.get("registry_name") == "groot_n1_7_action_decode_v1"
    )
    decode_config = decode_entry["config"]
    assert decode_config["use_relative_action"] is True
    decode_relative_max = decode_config["raw_stats"]["relative_action"]["single_arm"]["max"]
    assert decode_relative_max[:2] == [
        [-1.0, -2.0, -3.0, -4.0, -5.0],
        [2.0, 3.0, 4.0, 5.0, 6.0],
    ]
    assert len(decode_relative_max) == N1_7_NATIVE_ACTION_HORIZON
    assert (
        decode_config["raw_stats"]["relative_action"]["single_arm"]["count"]
        == [2] * N1_7_NATIVE_ACTION_HORIZON
    )
    assert decode_config["raw_stats"]["action"]["gripper"]["max"] == [100.0]


def test_groot_n1_7_relative_action_processors_compute_stats_from_runtime_dataset_meta(monkeypatch, tmp_path):
    pytest.importorskip("datasets")

    input_features, output_features = _groot_features(state_dim=6, action_dim=6)
    action_names = [
        "shoulder_pan.pos",
        "shoulder_lift.pos",
        "elbow_flex.pos",
        "wrist_flex.pos",
        "wrist_roll.pos",
        "gripper.pos",
    ]
    config = GrootConfig(
        input_features=input_features,
        output_features=output_features,
        device="cpu",
        use_bf16=False,
        action_decode_transform=None,
        chunk_size=2,
        n_action_steps=2,
        use_relative_actions=True,
        relative_exclude_joints=["gripper"],
    )
    absolute_dataset_stats = {
        OBS_STATE: {
            "min": torch.tensor([-50.0, -60.0, -70.0, -80.0, -90.0, 0.0]),
            "max": torch.tensor([50.0, 60.0, 70.0, 80.0, 90.0, 100.0]),
        },
        ACTION: {
            "min": torch.tensor([-100.0, -110.0, -120.0, -130.0, -140.0, 0.0]),
            "max": torch.tensor([100.0, 110.0, 120.0, 130.0, 140.0, 100.0]),
        },
    }
    samples = [
        {
            OBS_STATE: torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0, 0.0]),
            ACTION: _native_action_chunk(
                [
                    [8.0, 17.0, 26.0, 35.0, 44.0, 0.0],
                    [12.0, 23.0, 34.0, 45.0, 56.0, 100.0],
                ]
            ),
        },
        {
            OBS_STATE: torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0, 50.0]),
            ACTION: _native_action_chunk(
                [
                    [-1.0, -2.0, -3.0, -4.0, -5.0, 25.0],
                    [1.0, 2.0, 3.0, 4.0, 5.0, 75.0],
                ]
            ),
        },
    ]
    runtime_meta = SimpleNamespace(
        repo_id="local/relative",
        root=tmp_path,
        revision="main",
        fps=30,
        stats=absolute_dataset_stats,
        features={ACTION: {"names": action_names}},
    )

    class _RelativeStatsDataset:
        meta = runtime_meta

        def __len__(self):
            return len(samples)

        def __getitem__(self, idx):
            return samples[idx]

    def _fake_lerobot_dataset(repo_id, **kwargs):
        assert repo_id == runtime_meta.repo_id
        assert kwargs["root"] == runtime_meta.root
        assert kwargs["revision"] == runtime_meta.revision
        assert kwargs["download_videos"] is False
        assert kwargs["delta_timestamps"][ACTION] == [
            index / runtime_meta.fps for index in range(N1_7_NATIVE_ACTION_HORIZON)
        ]
        return _RelativeStatsDataset()

    monkeypatch.setattr("lerobot.policies.groot.processor_groot.LeRobotDataset", _fake_lerobot_dataset)
    config._runtime_dataset_meta = runtime_meta

    preprocessor, postprocessor = make_groot_pre_post_processors(config, dataset_stats=absolute_dataset_stats)

    assert not any(isinstance(step, RelativeActionsProcessorStep) for step in preprocessor.steps)
    assert isinstance(postprocessor.steps[0], GrootN17ActionDecodeStep)
    pack_step = next(step for step in preprocessor.steps if isinstance(step, GrootN17PackInputsStep))
    assert pack_step.action_horizon == N1_7_NATIVE_ACTION_HORIZON
    assert pack_step.valid_action_horizon == 2
    pack_relative_min = pack_step.raw_stats["relative_action"]["single_arm"]["min"]
    assert pack_relative_min[:2] == [
        [-2.0, -3.0, -4.0, -5.0, -6.0],
        [1.0, 2.0, 3.0, 4.0, 5.0],
    ]
    assert len(pack_relative_min) == N1_7_NATIVE_ACTION_HORIZON
    assert pack_step.raw_stats["relative_action"]["single_arm"]["count"] == [2] * N1_7_NATIVE_ACTION_HORIZON
    assert pack_step.raw_stats["action"]["gripper"]["max"] == [100.0]


def test_groot_n1_7_generated_relative_stats_match_oss_gr00t_reference_numbers():
    input_features, output_features = _groot_features(state_dim=6, action_dim=6)
    action_names = [
        "shoulder_pan.pos",
        "shoulder_lift.pos",
        "elbow_flex.pos",
        "wrist_flex.pos",
        "wrist_roll.pos",
        "gripper.pos",
    ]
    config = GrootConfig(
        input_features=input_features,
        output_features=output_features,
        device="cpu",
        use_bf16=False,
        action_decode_transform=None,
        chunk_size=3,
        n_action_steps=3,
        use_relative_actions=True,
        relative_exclude_joints=["gripper"],
    )
    absolute_dataset_stats = {
        OBS_STATE: {
            "min": torch.tensor([-20.0, -30.0, -40.0, -50.0, -60.0, 0.0]),
            "max": torch.tensor([80.0, 70.0, 60.0, 50.0, 40.0, 100.0]),
            "mean": torch.tensor([30.0, 20.0, 10.0, 0.0, -10.0, 50.0]),
            "std": torch.tensor([10.0, 10.0, 10.0, 10.0, 10.0, 10.0]),
            "q01": torch.tensor([-10.0, -20.0, -30.0, -40.0, -50.0, 10.0]),
            "q99": torch.tensor([70.0, 60.0, 50.0, 40.0, 30.0, 90.0]),
        },
        ACTION: {
            "min": torch.tensor([-5.0, -20.0, 0.0, -25.0, 10.0, 20.0]),
            "max": torch.tensor([20.0, 30.0, 45.0, 60.0, 70.0, 90.0]),
            "mean": torch.tensor([5.0, 5.0, 20.0, 20.0, 40.0, 55.0]),
            "std": torch.tensor([5.0, 10.0, 10.0, 20.0, 20.0, 25.0]),
            "q01": torch.tensor([-4.0, -19.0, 1.0, -24.0, 11.0, 20.0]),
            "q99": torch.tensor([19.0, 29.0, 44.0, 59.0, 69.0, 90.0]),
        },
    }
    state_a = torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0, 25.0])
    state_b = torch.tensor([0.0, -10.0, 10.0, -20.0, 20.0, 75.0])
    action_a = _native_action_chunk(
        [
            [11.0, 22.0, 33.0, 44.0, 55.0, 20.0],
            [12.0, 24.0, 36.0, 48.0, 60.0, 80.0],
            [13.0, 26.0, 39.0, 52.0, 65.0, 90.0],
        ]
    )
    action_b = _native_action_chunk(
        [
            [-1.0, -8.0, 13.0, -16.0, 25.0, 30.0],
            [-2.0, -6.0, 16.0, -12.0, 30.0, 40.0],
            [-3.0, -4.0, 19.0, -8.0, 35.0, 50.0],
        ]
    )
    samples = [
        {OBS_STATE: state_a, ACTION: action_a},
        {OBS_STATE: state_b, ACTION: action_b},
    ]

    class _Dataset:
        meta = SimpleNamespace(
            stats=absolute_dataset_stats,
            features={ACTION: {"names": action_names}},
        )

        def __len__(self):
            return len(samples)

        def __getitem__(self, idx):
            return samples[idx]

    relative_dataset_stats = _make_relative_action_training_stats(
        _Dataset(),
        exclude_joints=["gripper"],
        action_names=action_names,
        preserve_action_horizon=True,
    )

    # Static reference values from OSS GR00T's JointActionChunk.relative_chunking +
    # calculate_stats_for_key path: stats are computed per chunk timestep, not
    # flattened over all timesteps.
    oss_arm_min = torch.tensor(
        [
            [-1.0, 2.0, 3.0, 4.0, 5.0],
            [-2.0, 4.0, 6.0, 8.0, 10.0],
            [-3.0, 6.0, 9.0, 12.0, 15.0],
        ]
    )
    oss_arm_max = torch.tensor(
        [
            [1.0, 2.0, 3.0, 4.0, 5.0],
            [2.0, 4.0, 6.0, 8.0, 10.0],
            [3.0, 6.0, 9.0, 12.0, 15.0],
        ]
    )
    oss_arm_mean = torch.tensor(
        [
            [0.0, 2.0, 3.0, 4.0, 5.0],
            [0.0, 4.0, 6.0, 8.0, 10.0],
            [0.0, 6.0, 9.0, 12.0, 15.0],
        ]
    )
    oss_arm_std = torch.tensor(
        [
            [1.0, 0.0, 0.0, 0.0, 0.0],
            [2.0, 0.0, 0.0, 0.0, 0.0],
            [3.0, 0.0, 0.0, 0.0, 0.0],
        ]
    )
    oss_arm_q01 = torch.tensor(
        [
            [-0.98, 2.0, 3.0, 4.0, 5.0],
            [-1.96, 4.0, 6.0, 8.0, 10.0],
            [-2.94, 6.0, 9.0, 12.0, 15.0],
        ]
    )
    oss_arm_q99 = torch.tensor(
        [
            [0.98, 2.0, 3.0, 4.0, 5.0],
            [1.96, 4.0, 6.0, 8.0, 10.0],
            [2.94, 6.0, 9.0, 12.0, 15.0],
        ]
    )

    torch.testing.assert_close(torch.as_tensor(relative_dataset_stats[ACTION]["min"][:3, :5]), oss_arm_min)
    torch.testing.assert_close(torch.as_tensor(relative_dataset_stats[ACTION]["max"][:3, :5]), oss_arm_max)
    torch.testing.assert_close(torch.as_tensor(relative_dataset_stats[ACTION]["mean"][:3, :5]), oss_arm_mean)
    torch.testing.assert_close(torch.as_tensor(relative_dataset_stats[ACTION]["std"][:3, :5]), oss_arm_std)
    torch.testing.assert_close(torch.as_tensor(relative_dataset_stats[ACTION]["q01"][:3, :5]), oss_arm_q01)
    torch.testing.assert_close(torch.as_tensor(relative_dataset_stats[ACTION]["q99"][:3, :5]), oss_arm_q99)
    assert torch.as_tensor(relative_dataset_stats[ACTION]["min"]).shape[0] == N1_7_NATIVE_ACTION_HORIZON

    preprocessor, postprocessor = make_groot_pre_post_processors(
        config,
        dataset_stats=relative_dataset_stats,
        dataset_meta=_Dataset.meta,
    )
    pack_step = next(step for step in preprocessor.steps if isinstance(step, GrootN17PackInputsStep))
    decode_step = next(step for step in postprocessor.steps if isinstance(step, GrootN17ActionDecodeStep))

    assert pack_step.use_percentiles is True
    pack_relative_min = torch.as_tensor(pack_step.raw_stats["relative_action"]["single_arm"]["min"])
    pack_relative_q99 = torch.as_tensor(pack_step.raw_stats["relative_action"]["single_arm"]["q99"])
    assert pack_relative_min.shape == (N1_7_NATIVE_ACTION_HORIZON, 5)
    assert pack_relative_q99.shape == (N1_7_NATIVE_ACTION_HORIZON, 5)
    torch.testing.assert_close(pack_relative_min[:3], oss_arm_min)
    torch.testing.assert_close(pack_relative_q99[:3], oss_arm_q99)
    assert pack_step.stats[ACTION]["min"][:15] == pytest.approx(oss_arm_min.flatten().tolist())
    assert pack_step.stats[ACTION]["max"][:15] == pytest.approx(oss_arm_max.flatten().tolist())
    assert pack_step.stats[ACTION]["min"][-1] == pytest.approx(20.0)
    assert pack_step.stats[ACTION]["max"][-1] == pytest.approx(90.0)

    packed = pack_step(
        {
            TransitionKey.OBSERVATION: {OBS_STATE: state_a.unsqueeze(0)},
            TransitionKey.ACTION: action_a.unsqueeze(0),
            TransitionKey.COMPLEMENTARY_DATA: {"task": ["Move the vial"]},
        }
    )
    expected_normalized = torch.tensor(
        [
            [1.0, 0.0, 0.0, 0.0, 0.0, -1.0],
            [1.0, 0.0, 0.0, 0.0, 0.0, 5.0 / 7.0],
            [1.0, 0.0, 0.0, 0.0, 0.0, 1.0],
        ]
    )
    torch.testing.assert_close(packed[TransitionKey.ACTION][0, :3, :6], expected_normalized)

    decoded = decode_step({TransitionKey.ACTION: packed[TransitionKey.ACTION]})
    assert decoded[TransitionKey.ACTION].shape == (1, N1_7_NATIVE_ACTION_HORIZON, 6)
    torch.testing.assert_close(
        decoded[TransitionKey.ACTION][:, :3],
        action_a.unsqueeze(0)[:, :3],
        atol=1e-5,
        rtol=1e-5,
    )


def test_groot_n1_7_relative_action_stats_skip_padded_tail_chunks():
    samples = [
        {
            OBS_STATE: torch.tensor([10.0, 100.0]),
            ACTION: torch.tensor([[11.0, 101.0], [12.0, 102.0], [13.0, 103.0]]),
            f"{ACTION}_is_pad": torch.tensor([False, False, False]),
        },
        {
            OBS_STATE: torch.tensor([20.0, 200.0]),
            ACTION: torch.tensor([[18.0, 198.0], [16.0, 196.0], [14.0, 194.0]]),
            f"{ACTION}_is_pad": torch.tensor([False, False, False]),
        },
        {
            OBS_STATE: torch.tensor([0.0, 0.0]),
            ACTION: torch.tensor([[999.0, 999.0], [888.0, 888.0], [777.0, 777.0]]),
            f"{ACTION}_is_pad": torch.tensor([False, False, True]),
        },
    ]

    class _Dataset:
        meta = SimpleNamespace(stats={})

        def __len__(self):
            return len(samples)

        def __getitem__(self, idx):
            return samples[idx]

    relative_dataset_stats = _make_relative_action_training_stats(
        _Dataset(),
        exclude_joints=[],
        action_names=None,
        preserve_action_horizon=True,
    )

    torch.testing.assert_close(
        torch.as_tensor(relative_dataset_stats[ACTION]["count"]),
        torch.tensor([2, 2, 2]),
    )
    torch.testing.assert_close(
        torch.as_tensor(relative_dataset_stats[ACTION]["min"]),
        torch.tensor([[-2.0, -2.0], [-4.0, -4.0], [-6.0, -6.0]]),
    )
    torch.testing.assert_close(
        torch.as_tensor(relative_dataset_stats[ACTION]["max"]),
        torch.tensor([[1.0, 1.0], [2.0, 2.0], [3.0, 3.0]]),
    )


def test_groot_policy_selects_n1_7_model_class(monkeypatch):
    pytest.importorskip("transformers")

    from lerobot.policies.groot.groot_n1_7 import GR00TN17

    called = {}

    def fake_from_pretrained(cls, **kwargs):
        called.update(kwargs)
        return _DummyGrootModel()

    monkeypatch.setattr(GR00TN17, "from_pretrained", classmethod(fake_from_pretrained))

    policy = GrootPolicy(_groot_config())

    assert called["pretrained_model_name_or_path"] == GROOT_N1_7_BASE_MODEL
    assert isinstance(policy._groot_model, _DummyGrootModel)


def test_groot_policy_forwards_n1_7_qwen_inputs(monkeypatch):
    pytest.importorskip("transformers")

    from lerobot.policies.groot.groot_n1_7 import GR00TN17

    dummy_model = _DummyGrootModel()
    monkeypatch.setattr(GR00TN17, "from_pretrained", classmethod(lambda cls, **kwargs: dummy_model))
    policy = GrootPolicy(_groot_config())

    batch = {
        "state": torch.zeros(2, 1, 132),
        "action": torch.zeros(2, 40, 132),
        "action_mask": torch.ones(2, 40, 132),
        "embodiment_id": torch.zeros(2, dtype=torch.long),
        "input_ids": torch.ones(2, 8, dtype=torch.long),
        "attention_mask": torch.ones(2, 8, dtype=torch.long),
        "pixel_values": torch.zeros(4, 3, 16, 16),
        "image_grid_thw": torch.ones(4, 3, dtype=torch.long),
        "mm_token_type_ids": torch.zeros(2, 8, dtype=torch.int32),
        "pixel_values_videos": torch.zeros(1, 3, 16, 16),
        "video_grid_thw": torch.ones(1, 3, dtype=torch.long),
        "next.state": torch.ones(2, 1, 132),
        "info": {"ignored": True},
    }

    loss, metrics = policy.forward(batch)

    assert loss.item() == pytest.approx(1.0)
    assert metrics == {"loss": pytest.approx(1.0)}
    assert set(dummy_model.forward_inputs) == {
        "state",
        "action",
        "action_mask",
        "embodiment_id",
        "input_ids",
        "attention_mask",
        "pixel_values",
        "image_grid_thw",
        "mm_token_type_ids",
        "pixel_values_videos",
        "video_grid_thw",
    }


def test_groot_n1_7_libero_execution_horizon_uses_core_eight_action_cadence(tmp_path):
    model_path = tmp_path / "libero_spatial"
    _write_raw_n1_7_libero_checkpoint(model_path)

    assert infer_groot_n1_7_action_horizon(model_path, "libero_sim") == 16
    assert infer_groot_n1_7_action_execution_horizon(model_path, "libero_sim") == 8


def test_groot_select_action_rejects_relative_action_policies():
    policy = object.__new__(GrootPolicy)
    object.__setattr__(policy, "config", SimpleNamespace(use_relative_actions=True))

    with pytest.raises(NotImplementedError, match="select_action does not support relative-action policies"):
        policy.select_action({})


def test_groot_n1_7_select_action_uses_checkpoint_valid_horizon(tmp_path, monkeypatch):
    pytest.importorskip("transformers")

    from lerobot.policies.groot.groot_n1_7 import GR00TN17

    model_path = tmp_path / "libero_spatial"
    _write_raw_n1_7_libero_checkpoint(model_path)

    class HorizonModel(_DummyGrootModel):
        def get_action(self, inputs):
            assert inputs["action_mask"].shape == (1, 40)
            assert inputs["action_mask"][0, :16].sum().item() == 16
            assert inputs["action_mask"][0, 16:].sum().item() == 0
            batch_size = inputs["state"].shape[0]
            steps = torch.arange(40, dtype=torch.float32).view(1, 40, 1).expand(batch_size, 40, 132)
            return {"action_pred": steps}

    monkeypatch.setattr(GR00TN17, "from_pretrained", classmethod(lambda cls, **kwargs: HorizonModel()))
    input_features, output_features = _groot_features(state_dim=8, action_dim=7)
    config = GrootConfig(
        base_model_path=str(model_path),
        embodiment_tag="libero_sim",
        input_features=input_features,
        output_features=output_features,
        device="cpu",
        use_bf16=False,
        n_action_steps=40,
    )
    policy = GrootPolicy(config)
    batch = {
        "state": torch.zeros(1, 1, 132),
        "embodiment_id": torch.zeros(1, dtype=torch.long),
        "input_ids": torch.ones(1, 2, dtype=torch.long),
        "attention_mask": torch.ones(1, 2, dtype=torch.long),
        "pixel_values": torch.zeros(1, 3, 2, 2),
        "image_grid_thw": torch.ones(1, 3, dtype=torch.long),
        "action_mask": torch.cat((torch.ones(1, 16), torch.zeros(1, 24)), dim=1),
    }

    first_action = policy.select_action(batch)

    assert policy._action_queue_steps == 8
    assert len(policy._action_queue) == 7
    torch.testing.assert_close(first_action[0, 0], torch.tensor(0.0))

    for expected_step in range(1, 8):
        action = policy.select_action(batch)
        torch.testing.assert_close(action[0, 0], torch.tensor(float(expected_step)))

    refreshed_action = policy.select_action(batch)
    torch.testing.assert_close(refreshed_action[0, 0], torch.tensor(0.0))


def test_qwen3_backbone_uses_nested_transformers_model_contract(monkeypatch):
    pytest.importorskip("transformers")
    from transformers.feature_extraction_utils import BatchFeature

    import lerobot.policies.groot.groot_n1_7 as groot_n1_7

    class FakeLanguageModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.layers = nn.ModuleList([nn.Linear(1, 1) for _ in range(3)])

    class FakeVisual(nn.Module):
        def __init__(self):
            super().__init__()
            self.proj = nn.Linear(1, 1)

    class FakeInnerModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.language_model = FakeLanguageModel()
            self.visual = FakeVisual()

    class FakeQwenForConditionalGeneration(nn.Module):
        config = SimpleNamespace(image_token_id=42)

        def __init__(self):
            super().__init__()
            self.model = FakeInnerModel()

        @classmethod
        def from_pretrained(cls, *args, **kwargs):
            return cls()

        def eval(self):
            super().eval()
            return self

        def forward(self, **kwargs):
            batch_size, sequence_length = kwargs["input_ids"].shape
            features = torch.arange(batch_size * sequence_length * 4, dtype=torch.float32).view(
                batch_size, sequence_length, 4
            )
            return SimpleNamespace(hidden_states=[features, features + 1])

    monkeypatch.setattr(
        groot_n1_7,
        "Qwen3VLForConditionalGeneration",
        FakeQwenForConditionalGeneration,
    )
    backbone = groot_n1_7.Qwen3Backbone(
        model_name="fake-qwen",
        select_layer=2,
        tune_llm=False,
        tune_visual=False,
        use_flash_attention=False,
    )

    assert not hasattr(backbone.model, "language_model")
    assert len(backbone.language_model.layers) == 2
    assert not any(parameter.requires_grad for parameter in backbone.language_model.parameters())
    assert not any(parameter.requires_grad for parameter in backbone.visual.parameters())

    output = backbone.forward(
        BatchFeature(
            data={
                "input_ids": torch.tensor([[1, 42, 2], [42, 3, 4]]),
                "attention_mask": torch.tensor([[1, 1, 0], [1, 1, 1]]),
                "pixel_values": torch.zeros(2, 3, 2, 2),
                "image_grid_thw": torch.ones(2, 3, dtype=torch.long),
            }
        )
    )

    assert output["backbone_features"].shape == (2, 3, 4)
    torch.testing.assert_close(
        output["image_mask"],
        torch.tensor([[False, True, False], [True, False, False]]),
    )
    torch.testing.assert_close(
        output["backbone_attention_mask"],
        torch.tensor([[True, True, False], [True, True, True]]),
    )


def test_qwen3_backbone_can_initialize_from_config_without_downloading_weights(monkeypatch):
    pytest.importorskip("transformers")

    import lerobot.policies.groot.groot_n1_7 as groot_n1_7

    class FakeLanguageModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.layers = nn.ModuleList([nn.Linear(1, 1) for _ in range(3)])

    class FakeVisual(nn.Module):
        def __init__(self):
            super().__init__()
            self.proj = nn.Linear(1, 1)

    class FakeInnerModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.language_model = FakeLanguageModel()
            self.visual = FakeVisual()

    class FakeQwenForConditionalGeneration(nn.Module):
        config = SimpleNamespace(image_token_id=42)
        from_pretrained_called = False
        from_config_called = False

        def __init__(self):
            super().__init__()
            self.model = FakeInnerModel()

        @classmethod
        def from_pretrained(cls, *args, **kwargs):
            cls.from_pretrained_called = True
            raise AssertionError("Qwen backbone weights should not be loaded separately")

        @classmethod
        def _from_config(cls, config, **kwargs):
            cls.from_config_called = True
            return cls()

        def eval(self):
            super().eval()
            return self

    monkeypatch.setattr(groot_n1_7, "Qwen3VLForConditionalGeneration", FakeQwenForConditionalGeneration)

    backbone = groot_n1_7.Qwen3Backbone(
        model_name="nvidia/Cosmos-Reason2-2B",
        select_layer=2,
        load_pretrained_weights=False,
    )

    assert isinstance(backbone.model, FakeQwenForConditionalGeneration)
    assert FakeQwenForConditionalGeneration.from_config_called
    assert not FakeQwenForConditionalGeneration.from_pretrained_called


def test_gr00t_n1_7_from_pretrained_defers_backbone_weight_loading(monkeypatch, tmp_path):
    pytest.importorskip("transformers")

    from huggingface_hub.errors import HFValidationError

    import lerobot.policies.groot.groot_n1_7 as groot_n1_7

    called = {}

    class FakeLoadedModel:
        def __init__(self):
            self.config = SimpleNamespace(tune_top_llm_layers=0)
            self.backbone = SimpleNamespace(set_trainable_parameters=lambda **kwargs: None)
            self.action_head = SimpleNamespace(set_trainable_parameters=lambda **kwargs: None)

    def fake_snapshot_download(*args, **kwargs):
        raise HFValidationError("local path")

    def fake_super_from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        called["pretrained_model_name_or_path"] = pretrained_model_name_or_path
        called.update(kwargs)
        return FakeLoadedModel()

    monkeypatch.setattr(groot_n1_7, "snapshot_download", fake_snapshot_download)
    monkeypatch.setattr(
        groot_n1_7.PreTrainedModel,
        "from_pretrained",
        classmethod(fake_super_from_pretrained),
    )

    loaded = groot_n1_7.GR00TN17.from_pretrained(str(tmp_path))

    assert isinstance(loaded, FakeLoadedModel)
    assert called["pretrained_model_name_or_path"] == str(tmp_path)
    assert called["load_backbone_weights"] is False


def test_gr00t_n1_7_action_head_meta_init_defers_beta_distribution():
    pytest.importorskip("diffusers")

    from lerobot.policies.groot.groot_n1_7 import GR00TN17ActionHead, GR00TN17Config

    config = GR00TN17Config(
        backbone_embedding_dim=32,
        hidden_size=32,
        input_embedding_dim=32,
        max_state_dim=7,
        max_action_dim=5,
        action_horizon=4,
        state_history_length=1,
        max_num_embodiments=4,
        use_alternate_vl_dit=False,
        use_vlln=False,
        add_pos_embed=False,
        vl_self_attention_cfg={"num_layers": 0},
        diffusion_model_cfg={
            "positional_embeddings": None,
            "num_layers": 1,
            "num_attention_heads": 2,
            "attention_head_dim": 16,
            "norm_type": "ada_norm",
            "dropout": 0.0,
            "final_dropout": False,
            "output_dim": 32,
            "interleave_self_attention": False,
        },
    )

    with torch.device("meta"):
        meta_action_head = GR00TN17ActionHead(config)

    assert meta_action_head._beta_dist is None
    assert any(parameter.is_meta for parameter in meta_action_head.parameters())

    action_head = GR00TN17ActionHead(config)
    sample = action_head.sample_time(batch_size=3, device=torch.device("cpu"), dtype=torch.float32)

    assert action_head._beta_dist is not None
    assert sample.shape == (3,)
    assert torch.isfinite(sample).all()


def test_gr00t_n1_7_model_forward_with_mocked_backbone():
    pytest.importorskip("diffusers")
    pytest.importorskip("transformers")

    from transformers.feature_extraction_utils import BatchFeature

    from lerobot.policies.groot.groot_n1_7 import GR00TN17, GR00TN17Config

    config = GR00TN17Config(
        backbone_embedding_dim=32,
        hidden_size=32,
        input_embedding_dim=32,
        max_state_dim=7,
        max_action_dim=5,
        action_horizon=4,
        state_history_length=1,
        num_inference_timesteps=2,
        max_num_embodiments=4,
        use_alternate_vl_dit=False,
        use_vlln=True,
        vl_self_attention_cfg={"num_layers": 0},
        state_dropout_prob=0.0,
        diffusion_model_cfg={
            "positional_embeddings": None,
            "num_layers": 1,
            "num_attention_heads": 2,
            "attention_head_dim": 16,
            "norm_type": "ada_norm",
            "dropout": 0.0,
            "final_dropout": False,
            "output_dim": 32,
            "interleave_self_attention": False,
        },
    )

    class MockBackbone(nn.Module):
        def __init__(self):
            super().__init__()
            self.weight = nn.Parameter(torch.zeros(()))

        def prepare_input(self, inputs):
            return BatchFeature(data=inputs)

        def forward(self, inputs):
            batch_size = inputs["state"].shape[0]
            return BatchFeature(
                data={
                    "backbone_features": torch.randn(batch_size, 3, config.backbone_embedding_dim),
                    "backbone_attention_mask": torch.ones(batch_size, 3, dtype=torch.bool),
                    "image_mask": torch.zeros(batch_size, 3, dtype=torch.bool),
                }
            )

        def set_trainable_parameters(self, *args, **kwargs):
            return None

    with patch(
        "lerobot.policies.groot.groot_n1_7.get_backbone_cls",
        return_value=lambda **kwargs: MockBackbone(),
    ):
        model = GR00TN17(config)

    inputs = {
        "state": torch.randn(2, config.state_history_length, config.max_state_dim),
        "action": torch.randn(2, config.action_horizon, config.max_action_dim),
        "action_mask": torch.ones(2, config.action_horizon, config.max_action_dim),
        "embodiment_id": torch.zeros(2, dtype=torch.long),
    }

    output = model.forward(inputs)
    assert output["loss"].dim() == 0
    assert torch.isfinite(output["loss"])

    inference_inputs = {key: value for key, value in inputs.items() if key != "action"}
    action_output = model.get_action(inference_inputs)
    assert action_output["action_pred"].shape == (2, config.action_horizon, config.max_action_dim)