Merge remote-tracking branch 'origin/main' into feat/smolvla-on-steerable

Co-authored-by: Cursor <cursoragent@cursor.com>

# Conflicts:
#	src/lerobot/configs/train.py
#	src/lerobot/datasets/__init__.py
#	src/lerobot/policies/factory.py
#	src/lerobot/policies/groot/groot_n1.py
#	src/lerobot/scripts/lerobot_eval.py
#	src/lerobot/scripts/lerobot_train.py
#	uv.lock

tests/policies/evo1/test_evo1.py

@@ -0,0 +1,840 @@
+#!/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.
+
+from __future__ import annotations
+
+import pytest
+import torch
+from torch import nn
+
+import lerobot.policies.evo1.evo1_model as evo1_model
+import lerobot.policies.evo1.modeling_evo1 as modeling_evo1
+from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.policies.evo1.configuration_evo1 import Evo1Config
+from lerobot.policies.evo1.flow_matching import FlowmatchingActionHead
+from lerobot.policies.evo1.internvl3_embedder import (
+    IMAGENET_MEAN,
+    IMAGENET_STD,
+    _batched_pixel_values,
+)
+from lerobot.policies.evo1.processor_evo1 import (
+    Evo1ActionProcessorStep,
+    Evo1PadActionProcessorStep,
+    Evo1PadStateProcessorStep,
+    evo1_batch_to_transition,
+    make_evo1_pre_post_processors,
+    reconcile_evo1_processors,
+)
+from lerobot.policies.factory import get_policy_class, make_policy_config
+from lerobot.policies.rtc.configuration_rtc import RTCConfig
+from lerobot.policies.rtc.modeling_rtc import RTCProcessor
+from lerobot.processor import (
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    PolicyProcessorPipeline,
+    UnnormalizerProcessorStep,
+)
+from lerobot.processor.converters import (
+    batch_to_transition,
+    policy_action_to_transition,
+    transition_to_batch,
+    transition_to_policy_action,
+)
+from lerobot.utils.constants import (
+    ACTION,
+    OBS_IMAGES,
+    OBS_STATE,
+    POLICY_POSTPROCESSOR_DEFAULT_NAME,
+    POLICY_PREPROCESSOR_DEFAULT_NAME,
+)
+
+STATE_DIM = 4
+ACTION_DIM = 3
+MAX_STATE_DIM = 6
+MAX_ACTION_DIM = 5
+CHUNK_SIZE = 2
+EMBED_DIM = 8
+
+
+class DummyEvo1Model(nn.Module):
+    def __init__(self, config, vlm_hub_kwargs=None):
+        super().__init__()
+        self.config = config
+        self.embedder = nn.Dropout(p=0.0)
+        self.action_head = nn.Linear(1, 1)
+        self.get_vl_embeddings_calls = 0
+        self.grad_enabled_calls = []
+        self.embedder_training_calls = []
+
+    def set_finetune_flags(self):
+        return None
+
+    def get_vl_embeddings(self, images, image_mask, prompt=None, return_cls_only=False):
+        self.get_vl_embeddings_calls += 1
+        self.grad_enabled_calls.append(torch.is_grad_enabled())
+        self.embedder_training_calls.append(self.embedder.training)
+        # images is a list of per-camera (B, C, H, W) tensors, so the batch dim is images[0].shape[0].
+        batch_size = images[0].shape[0]
+        tokens = torch.ones(batch_size, 4, EMBED_DIM, requires_grad=torch.is_grad_enabled())
+        valid_mask = torch.ones(batch_size, 4, dtype=torch.bool)
+        return tokens, valid_mask
+
+    def forward(
+        self,
+        fused_tokens,
+        state=None,
+        actions_gt=None,
+        action_mask=None,
+        embodiment_ids=None,
+        context_mask=None,
+        **kwargs,
+    ):
+        batch_size = fused_tokens.shape[0]
+        if actions_gt is None:
+            return torch.ones(batch_size, CHUNK_SIZE * MAX_ACTION_DIM)
+        pred_velocity = torch.zeros(batch_size, CHUNK_SIZE * MAX_ACTION_DIM)
+        noise = torch.zeros_like(actions_gt)
+        return pred_velocity, noise
+
+
+class ChunkCountingDummyModel(DummyEvo1Model):
+    """Emits per-step distinguishable actions so queue ordering and re-prediction are observable."""
+
+    def __init__(self, config, vlm_hub_kwargs=None):
+        super().__init__(config, vlm_hub_kwargs)
+        self.chunks_predicted = 0
+
+    def forward(
+        self,
+        fused_tokens,
+        state=None,
+        actions_gt=None,
+        action_mask=None,
+        embodiment_ids=None,
+        context_mask=None,
+        **kwargs,
+    ):
+        if actions_gt is not None:
+            return super().forward(fused_tokens, state, actions_gt, action_mask, embodiment_ids, context_mask)
+        self.chunks_predicted += 1
+        batch_size = fused_tokens.shape[0]
+        step_values = torch.arange(CHUNK_SIZE, dtype=torch.float32) + 10.0 * self.chunks_predicted
+        chunk = step_values.repeat_interleave(MAX_ACTION_DIM).unsqueeze(0).repeat(batch_size, 1)
+        return chunk
+
+
+def make_config(training_stage="stage1", **kwargs):
+    config_kwargs = {
+        "device": "cpu",
+        "vlm_model_name": "dummy-internvl3",
+        "training_stage": training_stage,
+        "chunk_size": CHUNK_SIZE,
+        "n_action_steps": 1,
+        "max_state_dim": MAX_STATE_DIM,
+        "max_action_dim": MAX_ACTION_DIM,
+        "max_views": 2,
+        "embed_dim": EMBED_DIM,
+        "hidden_dim": 16,
+        "state_hidden_dim": 16,
+        "num_heads": 2,
+        "num_layers": 1,
+        "num_inference_timesteps": 2,
+        "input_features": {
+            OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(STATE_DIM,)),
+            f"{OBS_IMAGES}.front": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 16, 16)),
+        },
+        "output_features": {
+            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(ACTION_DIM,)),
+        },
+    }
+    config_kwargs.update(kwargs)
+    return Evo1Config(**config_kwargs)
+
+
+def make_batch(include_action=True):
+    batch = {
+        "task": ["pick the block", "place the block"],
+        OBS_STATE: torch.randn(2, STATE_DIM),
+        f"{OBS_IMAGES}.front": torch.rand(2, 3, 16, 16),
+    }
+    if include_action:
+        batch[ACTION] = torch.randn(2, CHUNK_SIZE, ACTION_DIM)
+    return batch
+
+
+def make_stats(state_dim=STATE_DIM, action_dim=ACTION_DIM):
+    return {
+        OBS_STATE: {
+            "min": torch.full((state_dim,), -2.0),
+            "max": torch.full((state_dim,), 2.0),
+        },
+        ACTION: {
+            "min": torch.full((action_dim,), -1.0),
+            "max": torch.full((action_dim,), 1.0),
+        },
+    }
+
+
+def make_flowmatching_head(**overrides):
+    kwargs = {
+        "embed_dim": EMBED_DIM,
+        "hidden_dim": 16,
+        "action_dim": CHUNK_SIZE * ACTION_DIM,
+        "horizon": CHUNK_SIZE,
+        "per_action_dim": ACTION_DIM,
+        "num_heads": 2,
+        "num_layers": 1,
+        "num_inference_timesteps": 2,
+        "state_dim": STATE_DIM,
+        "state_hidden_dim": 16,
+        "num_categories": 1,
+    }
+    kwargs.update(overrides)
+    return FlowmatchingActionHead(**kwargs)
+
+
+def test_evo1_factory_registration():
+    cfg = make_policy_config(
+        "evo1",
+        device="cpu",
+        vlm_model_name="dummy-internvl3",
+        input_features={
+            OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(STATE_DIM,)),
+            f"{OBS_IMAGES}.front": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 16, 16)),
+        },
+        output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(ACTION_DIM,))},
+    )
+
+    assert isinstance(cfg, Evo1Config)
+    assert get_policy_class("evo1") is modeling_evo1.Evo1Policy
+
+
+def test_evo1_stage_defaults_and_consistency():
+    stage1 = make_config(training_stage="stage1")
+    assert (stage1.finetune_vlm, stage1.finetune_language_model, stage1.finetune_vision_model) == (
+        False,
+        False,
+        False,
+    )
+    assert stage1.finetune_action_head is True
+
+    stage2 = make_config(training_stage="stage2")
+    assert (stage2.finetune_vlm, stage2.finetune_language_model, stage2.finetune_vision_model) == (
+        True,
+        True,
+        True,
+    )
+    assert stage2.finetune_action_head is True
+
+    stage2_from_stage1_checkpoint_flags = make_config(
+        training_stage="stage2",
+        finetune_vlm=False,
+        finetune_language_model=False,
+        finetune_vision_model=False,
+        finetune_action_head=False,
+    )
+    assert (
+        stage2_from_stage1_checkpoint_flags.finetune_vlm,
+        stage2_from_stage1_checkpoint_flags.finetune_language_model,
+        stage2_from_stage1_checkpoint_flags.finetune_vision_model,
+    ) == (
+        True,
+        True,
+        True,
+    )
+    assert stage2_from_stage1_checkpoint_flags.finetune_action_head is True
+
+    explicit_off = make_config(
+        training_stage="stage2",
+        apply_training_stage_defaults=False,
+        finetune_vlm=False,
+        finetune_language_model=False,
+        finetune_vision_model=False,
+        finetune_action_head=False,
+    )
+    assert (
+        explicit_off.finetune_vlm,
+        explicit_off.finetune_language_model,
+        explicit_off.finetune_vision_model,
+    ) == (
+        False,
+        False,
+        False,
+    )
+    assert explicit_off.finetune_action_head is False
+
+    # An explicit finetune_vlm=False without branch-level flags freezes both branches instead of
+    # raising an inconsistency error.
+    frozen_vlm = make_config(
+        training_stage="stage2",
+        apply_training_stage_defaults=False,
+        finetune_vlm=False,
+    )
+    assert (
+        frozen_vlm.finetune_vlm,
+        frozen_vlm.finetune_language_model,
+        frozen_vlm.finetune_vision_model,
+    ) == (False, False, False)
+
+    try:
+        make_config(
+            training_stage="stage2",
+            apply_training_stage_defaults=False,
+            finetune_vlm=True,
+            finetune_language_model=False,
+        )
+    except ValueError as exc:
+        assert "Inconsistent EVO1 finetune config" in str(exc)
+    else:
+        raise AssertionError("Expected inconsistent finetune config to raise ValueError")
+
+
+def test_evo1_rejects_non_square_image_resolution():
+    with pytest.raises(ValueError, match="square image_resolution"):
+        make_config(image_resolution=(448, 320))
+
+
+def test_evo1_rejects_out_of_range_default_embodiment_id():
+    with pytest.raises(ValueError, match="default_embodiment_id"):
+        make_config(default_embodiment_id=3, num_categories=2)
+
+
+def test_evo1_model_uses_image_resolution_and_trainable_checkpointing(monkeypatch):
+    captured: dict = {}
+
+    class SpyEmbedder(nn.Module):
+        def __init__(self, **kwargs):
+            super().__init__()
+            captured.clear()
+            captured.update(kwargs)
+
+    monkeypatch.setattr(evo1_model, "InternVL3Embedder", SpyEmbedder)
+
+    stage1 = make_config(training_stage="stage1", image_resolution=(224, 224))
+    evo1_model.Evo1Model(stage1)
+    assert captured["image_size"] == 224
+    # VLM is frozen in stage1, so gradient checkpointing is gated off.
+    assert captured["enable_gradient_checkpointing"] is False
+
+    stage2 = make_config(training_stage="stage2", image_resolution=(224, 224))
+    evo1_model.Evo1Model(stage2)
+    assert captured["enable_gradient_checkpointing"] is True
+
+
+class FakeInternVLModel(nn.Module):
+    """Minimal stand-in with the native HF InternVL submodule layout."""
+
+    def __init__(self):
+        super().__init__()
+        self.language_model = nn.Linear(2, 2)
+        self.vision_tower = nn.Linear(2, 2)
+        self.multi_modal_projector = nn.Linear(2, 2)
+
+
+class FakeEmbedder(nn.Module):
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.model = FakeInternVLModel()
+
+
+def test_set_finetune_flags_targets_native_hf_internvl_submodules(monkeypatch):
+    monkeypatch.setattr(evo1_model, "InternVL3Embedder", FakeEmbedder)
+
+    stage2_model = evo1_model.Evo1Model(make_config(training_stage="stage2"))
+    stage2_model.set_finetune_flags()
+    vlm = stage2_model.embedder.model
+    assert all(p.requires_grad for p in vlm.language_model.parameters())
+    assert all(p.requires_grad for p in vlm.vision_tower.parameters())
+    assert all(p.requires_grad for p in vlm.multi_modal_projector.parameters())
+    assert all(p.requires_grad for p in stage2_model.action_head.parameters())
+
+    stage1_model = evo1_model.Evo1Model(make_config(training_stage="stage1"))
+    stage1_model.set_finetune_flags()
+    vlm = stage1_model.embedder.model
+    assert not any(p.requires_grad for p in vlm.parameters())
+    assert all(p.requires_grad for p in stage1_model.action_head.parameters())
+
+
+def test_set_finetune_flags_fails_loudly_on_unknown_vlm_layout(monkeypatch):
+    class LegacyLayoutModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.language_model = nn.Linear(2, 2)
+            self.vision_model = nn.Linear(2, 2)  # trust_remote_code-era attribute name
+            self.mlp1 = nn.Linear(2, 2)
+
+    class FakeEmbedder(nn.Module):
+        def __init__(self, **kwargs):
+            super().__init__()
+            self.model = LegacyLayoutModel()
+
+    monkeypatch.setattr(evo1_model, "InternVL3Embedder", FakeEmbedder)
+    model = evo1_model.Evo1Model(make_config(training_stage="stage2"))
+    with pytest.raises(AttributeError, match="vision_tower"):
+        model.set_finetune_flags()
+
+
+def test_evo1_policy_processors_pad_state_crop_action_and_binarize_gripper():
+    libero_action_dim = 7
+    config = make_config(
+        max_state_dim=MAX_STATE_DIM,
+        max_action_dim=8,
+        postprocess_action_dim=libero_action_dim,
+        binarize_gripper=True,
+        output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(libero_action_dim,))},
+    )
+    stats = make_stats(action_dim=libero_action_dim)
+
+    preprocessor, postprocessor = make_evo1_pre_post_processors(config, dataset_stats=stats)
+
+    assert isinstance(preprocessor.steps[2], Evo1PadStateProcessorStep)
+    assert isinstance(preprocessor.steps[3], Evo1PadActionProcessorStep)
+    assert isinstance(preprocessor.steps[4], NormalizerProcessorStep)
+    assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
+    assert isinstance(postprocessor.steps[1], Evo1ActionProcessorStep)
+
+    normalizer = preprocessor.steps[4]
+    assert normalizer.features[OBS_STATE].shape == (MAX_STATE_DIM,)
+    assert normalizer.features[ACTION].shape == (8,)
+    assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (MAX_STATE_DIM,)
+    assert normalizer._tensor_stats[ACTION]["min"].shape == (8,)
+
+    processed_batch = preprocessor(
+        {
+            "task": "pick the block",
+            OBS_STATE: torch.zeros(STATE_DIM),
+            ACTION: torch.zeros(libero_action_dim),
+            f"{OBS_IMAGES}.front": torch.rand(3, 16, 16),
+        }
+    )
+    processed_state = processed_batch[OBS_STATE]
+    assert processed_state.shape == (1, MAX_STATE_DIM)
+    assert torch.allclose(processed_state, torch.zeros_like(processed_state))
+    assert processed_batch[ACTION].shape == (1, 8)
+    assert torch.allclose(processed_batch[ACTION], torch.zeros_like(processed_batch[ACTION]))
+    assert processed_batch["action_mask"].shape == (1, 8)
+    assert processed_batch["action_mask"][:, :libero_action_dim].all()
+    assert not processed_batch["action_mask"][:, libero_action_dim:].any()
+
+    action = torch.tensor(
+        [
+            [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.5, 0.7],
+            [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7],
+        ],
+        dtype=torch.float32,
+    )
+    processed = postprocessor(action)
+
+    assert processed.shape == (2, 7)
+    assert processed.dtype == torch.float32
+    assert torch.allclose(processed[:, :6], action[:, :6])
+    assert torch.equal(processed[:, 6], torch.tensor([1.0, -1.0]))
+
+
+def test_evo1_postprocessor_returns_float32_for_bf16_actions():
+    config = make_config()
+    _preprocessor, postprocessor = make_evo1_pre_post_processors(config, dataset_stats=make_stats())
+
+    processed = postprocessor(torch.zeros(2, MAX_ACTION_DIM, dtype=torch.bfloat16))
+    assert processed.dtype == torch.float32
+
+
+def test_evo1_processor_save_load_round_trip_applies_config_overrides(tmp_path):
+    train_config = make_config()
+    preprocessor, postprocessor = make_evo1_pre_post_processors(train_config, dataset_stats=make_stats())
+    preprocessor.save_pretrained(tmp_path)
+    postprocessor.save_pretrained(tmp_path)
+
+    loaded_pre = PolicyProcessorPipeline.from_pretrained(
+        tmp_path,
+        config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json",
+        to_transition=batch_to_transition,
+        to_output=transition_to_batch,
+    )
+    loaded_post = PolicyProcessorPipeline.from_pretrained(
+        tmp_path,
+        config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json",
+        to_transition=policy_action_to_transition,
+        to_output=transition_to_policy_action,
+    )
+
+    # Simulate eval-time CLI overrides applied on top of the loaded pipelines.
+    eval_config = make_config(binarize_gripper=True, postprocess_action_dim=ACTION_DIM)
+    loaded_pre, loaded_post = reconcile_evo1_processors(eval_config, loaded_pre, loaded_post)
+
+    assert loaded_pre.to_transition is evo1_batch_to_transition
+    assert sum(isinstance(step, Evo1ActionProcessorStep) for step in loaded_post.steps) == 1
+    action_step = next(step for step in loaded_post.steps if isinstance(step, Evo1ActionProcessorStep))
+    assert action_step.binarize_gripper is True
+    assert action_step.action_dim == ACTION_DIM
+    # The float32 output dtype is part of the serialized pipeline itself.
+    device_step = next(step for step in loaded_post.steps if isinstance(step, DeviceProcessorStep))
+    assert device_step.float_dtype == "float32"
+
+    # Non-observation extras (embodiment_id, ...) must survive the reloaded preprocessor.
+    processed = loaded_pre(
+        {
+            "task": "pick the block",
+            OBS_STATE: torch.zeros(STATE_DIM),
+            f"{OBS_IMAGES}.front": torch.rand(3, 16, 16),
+            "embodiment_id": torch.tensor([0]),
+        }
+    )
+    assert "embodiment_id" in processed
+
+
+def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch):
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
+    policy = modeling_evo1.Evo1Policy(make_config())
+    preprocessor, _postprocessor = make_evo1_pre_post_processors(policy.config, dataset_stats=make_stats())
+    training_batch = preprocessor(make_batch(include_action=True))
+
+    assert training_batch[ACTION].shape == (2, CHUNK_SIZE, MAX_ACTION_DIM)
+    assert training_batch["action_mask"].shape == (2, CHUNK_SIZE, MAX_ACTION_DIM)
+    assert training_batch["action_mask"][:, :, :ACTION_DIM].all()
+    assert not training_batch["action_mask"][:, :, ACTION_DIM:].any()
+
+    loss, metrics = policy.forward(training_batch)
+    assert loss.ndim == 0
+    assert torch.isfinite(loss)
+    assert metrics["active_action_dims"] == ACTION_DIM * CHUNK_SIZE
+    assert policy.model.get_vl_embeddings_calls == 1
+
+    action_chunk = policy.predict_action_chunk(make_batch(include_action=False))
+    assert action_chunk.shape == (2, CHUNK_SIZE, MAX_ACTION_DIM)
+    assert action_chunk.dtype == torch.float32
+
+    policy.reset()
+    selected = policy.select_action(make_batch(include_action=False))
+    assert selected.shape == (2, MAX_ACTION_DIM)
+
+
+def test_evo1_forward_masks_padded_action_timesteps(monkeypatch):
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
+    policy = modeling_evo1.Evo1Policy(make_config())
+
+    batch = make_batch(include_action=True)
+    batch[ACTION] = torch.ones(2, CHUNK_SIZE, ACTION_DIM)
+    # Give the padded (past-episode-end) timestep a huge value: if it leaked into the loss, the
+    # loss would blow up far beyond 1.0.
+    batch[ACTION][:, -1, :] = 100.0
+    batch["action_is_pad"] = torch.zeros(2, CHUNK_SIZE, dtype=torch.bool)
+    batch["action_is_pad"][:, -1] = True
+
+    loss, metrics = policy.forward(batch)
+
+    # DummyEvo1Model predicts zero velocity and zero noise, so each active element contributes
+    # (0 - action)^2 = 1.0 for the in-episode ones-valued actions.
+    assert metrics["active_action_dims"] == ACTION_DIM * (CHUNK_SIZE - 1)
+    assert torch.isclose(loss, torch.tensor(1.0))
+
+
+def test_evo1_select_action_queue_orders_steps_and_repredicts(monkeypatch):
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", ChunkCountingDummyModel)
+    policy = modeling_evo1.Evo1Policy(make_config(n_action_steps=CHUNK_SIZE))
+
+    batch = make_batch(include_action=False)
+    first = policy.select_action(batch)
+    second = policy.select_action(batch)
+    third = policy.select_action(batch)
+
+    # First chunk provides steps 10, 11 in order; the third call triggers a fresh prediction (20).
+    assert torch.all(first == 10.0)
+    assert torch.all(second == 11.0)
+    assert torch.all(third == 20.0)
+    assert policy.model.chunks_predicted == 2
+
+
+def test_evo1_predict_action_chunk_rejects_rtc_kwargs_without_rtc_config(monkeypatch):
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
+    policy = modeling_evo1.Evo1Policy(make_config())
+    with pytest.raises(RuntimeError, match="RTC"):
+        policy.predict_action_chunk(make_batch(include_action=False), inference_delay=2)
+
+
+def test_evo1_rtc_processor_wiring(monkeypatch):
+    monkeypatch.setattr(evo1_model, "InternVL3Embedder", FakeEmbedder)
+    policy = modeling_evo1.Evo1Policy(make_config())
+    assert policy.rtc_processor is None
+    assert policy.model.rtc_processor is None
+
+    # The RTC rollout backend assigns rtc_config after loading and re-inits the processor.
+    policy.config.rtc_config = RTCConfig(execution_horizon=CHUNK_SIZE)
+    policy.init_rtc_processor()
+    assert isinstance(policy.rtc_processor, RTCProcessor)
+    assert policy.model.rtc_processor is policy.rtc_processor
+
+    # RTC drives predict_action_chunk directly; the select_action queue path is unsupported.
+    with pytest.raises(AssertionError, match="select_action"):
+        policy.select_action(make_batch(include_action=False))
+
+
+def test_flowmatching_rtc_guidance_pulls_prefix_toward_previous_chunk():
+    head = make_flowmatching_head(num_inference_timesteps=16)
+    processor = RTCProcessor(RTCConfig(execution_horizon=CHUNK_SIZE))
+    fused = torch.randn(2, 4, EMBED_DIM)
+    state = torch.randn(2, STATE_DIM)
+    action_mask = torch.ones(2, ACTION_DIM, dtype=torch.bool)
+    prev_chunk = torch.tensor([0.7, -0.4, 0.2]).expand(2, CHUNK_SIZE, ACTION_DIM).contiguous()
+
+    torch.manual_seed(0)
+    unguided = head.get_action(fused, state=state, action_mask=action_mask)
+    unguided = unguided.view(2, CHUNK_SIZE, ACTION_DIM)
+    torch.manual_seed(0)
+    guided = head.get_action(
+        fused,
+        state=state,
+        action_mask=action_mask,
+        inference_delay=1,
+        prev_chunk_left_over=prev_chunk,
+        rtc_processor=processor,
+    )
+    guided = guided.view(2, CHUNK_SIZE, ACTION_DIM)
+
+    # The frozen prefix (first inference_delay steps) must land far closer to the previous chunk
+    # than the unguided sample from the same noise does.
+    guided_dist = (guided[:, 0] - prev_chunk[:, 0]).abs().mean()
+    unguided_dist = (unguided[:, 0] - prev_chunk[:, 0]).abs().mean()
+    assert guided_dist < 0.5 * unguided_dist
+    assert torch.isfinite(guided).all()
+
+
+def test_flowmatching_rtc_first_chunk_without_leftover_matches_unguided():
+    head = make_flowmatching_head(num_inference_timesteps=4)
+    processor = RTCProcessor(RTCConfig(execution_horizon=CHUNK_SIZE))
+    fused = torch.randn(2, 4, EMBED_DIM)
+    state = torch.randn(2, STATE_DIM)
+    action_mask = torch.ones(2, ACTION_DIM, dtype=torch.bool)
+
+    torch.manual_seed(0)
+    unguided = head.get_action(fused, state=state, action_mask=action_mask)
+    torch.manual_seed(0)
+    first_chunk = head.get_action(
+        fused,
+        state=state,
+        action_mask=action_mask,
+        inference_delay=2,
+        prev_chunk_left_over=None,
+        rtc_processor=processor,
+    )
+
+    assert torch.allclose(unguided, first_chunk)
+
+
+def test_evo1_missing_configured_camera_needs_empty_cameras_budget(monkeypatch):
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
+    batch = make_batch(include_action=False)  # only provides the front camera
+
+    two_camera_features = {
+        OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(STATE_DIM,)),
+        f"{OBS_IMAGES}.front": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 16, 16)),
+        f"{OBS_IMAGES}.wrist": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 16, 16)),
+    }
+    strict_policy = modeling_evo1.Evo1Policy(make_config(input_features=dict(two_camera_features)))
+    with pytest.raises(ValueError, match="empty_cameras"):
+        strict_policy._collect_image_batches(batch)
+
+    # empty_cameras adds placeholder camera features that are never present in the batch; they
+    # become masked-out views instead of crashing with a KeyError.
+    padded_policy = modeling_evo1.Evo1Policy(make_config(empty_cameras=1))
+    assert len(padded_policy.config.image_features) == 2
+    camera_images, image_masks = padded_policy._collect_image_batches(batch)
+    assert len(camera_images) == 1
+    assert image_masks.tolist() == [[True, False], [True, False]]
+
+
+def test_stage1_frozen_vlm_embeddings_do_not_track_gradients(monkeypatch):
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
+    policy = modeling_evo1.Evo1Policy(make_config(training_stage="stage1"))
+    policy.train()
+
+    image_batches, image_masks = policy._collect_image_batches(make_batch(include_action=False))
+    fused_tokens, context_mask = policy._compute_fused_tokens(["pick", "place"], image_batches, image_masks)
+
+    assert policy.model.grad_enabled_calls == [False]
+    assert policy.model.embedder_training_calls == [False]
+    assert not fused_tokens.requires_grad
+    assert context_mask is not None
+    assert policy.model.embedder.training is False
+
+
+def test_stage2_vlm_embeddings_track_gradients(monkeypatch):
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
+    policy = modeling_evo1.Evo1Policy(make_config(training_stage="stage2"))
+    policy.train()
+
+    image_batches, image_masks = policy._collect_image_batches(make_batch(include_action=False))
+    fused_tokens, _context_mask = policy._compute_fused_tokens(["pick", "place"], image_batches, image_masks)
+
+    assert policy.model.grad_enabled_calls == [True]
+    assert policy.model.embedder_training_calls == [True]
+    assert fused_tokens.requires_grad
+
+
+def test_collect_image_batches_handles_unbatched_chw(monkeypatch):
+    # Regression for an issue where batch_size was read from shape[0] before normalizing
+    # per-camera tensor dims, so an unbatched (C, H, W) input was treated as batch_size=C.
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
+    policy = modeling_evo1.Evo1Policy(make_config())
+    batch = {
+        OBS_STATE: torch.randn(1, STATE_DIM),
+        f"{OBS_IMAGES}.front": torch.rand(3, 16, 16),
+    }
+
+    camera_images, image_masks = policy._collect_image_batches(batch)
+
+    # One present camera, returned as a batched (B, C, H, W) tensor with the unbatched CHW frame
+    # promoted to batch_size=1 (not read as batch_size=C).
+    assert len(camera_images) == 1
+    assert camera_images[0].shape == (1, 3, 16, 16)
+    assert image_masks.tolist() == [[True, False]]
+
+
+def test_evo1_state_mask_zeroes_masked_dims(monkeypatch):
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
+    policy = modeling_evo1.Evo1Policy(make_config())
+    batch = {
+        OBS_STATE: torch.ones(2, STATE_DIM),
+        "state_mask": torch.tensor([[True, True, False, False]] * 2),
+    }
+
+    states, mask = policy._prepare_state(batch)
+
+    assert torch.all(states[:, :2] == 1.0)
+    assert torch.all(states[:, 2:] == 0.0)
+    assert mask[:, :2].all()
+    assert not mask[:, 2:].any()
+
+
+def test_evo1_action_mask_accepts_chunk_size_one(monkeypatch):
+    monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
+    config = make_config(chunk_size=1, n_action_steps=1)
+    policy = modeling_evo1.Evo1Policy(config)
+    batch = make_batch(include_action=True)
+    batch[ACTION] = torch.randn(2, ACTION_DIM)
+    batch["action_mask"] = torch.ones(2, ACTION_DIM, dtype=torch.bool)
+
+    actions, action_mask = policy._prepare_actions(batch)
+
+    assert actions.shape == (2, 1, MAX_ACTION_DIM)
+    assert action_mask.shape == (2, 1, MAX_ACTION_DIM)
+    assert action_mask[:, :, :ACTION_DIM].all()
+    assert not action_mask[:, :, ACTION_DIM:].any()
+
+
+def test_flowmatching_state_encoder_for_horizon_one():
+    head = make_flowmatching_head(action_dim=ACTION_DIM, horizon=1)
+
+    assert head.state_encoder is not None
+    pred_velocity, noise = head(
+        torch.randn(2, 4, EMBED_DIM),
+        state=torch.randn(2, STATE_DIM),
+        actions_gt=torch.randn(2, 1, ACTION_DIM),
+        action_mask=torch.ones(2, 1, ACTION_DIM, dtype=torch.bool),
+    )
+
+    assert pred_velocity.shape == (2, ACTION_DIM)
+    assert noise.shape == (2, 1, ACTION_DIM)
+
+
+def test_flowmatching_get_action_real_path_respects_action_mask():
+    torch.manual_seed(0)
+    head = make_flowmatching_head()
+
+    action_mask = torch.zeros(2, ACTION_DIM, dtype=torch.bool)
+    action_mask[:, :2] = True
+    actions = head.get_action(
+        torch.randn(2, 4, EMBED_DIM),
+        state=torch.randn(2, STATE_DIM),
+        action_mask=action_mask,
+    )
+
+    assert actions.shape == (2, CHUNK_SIZE * ACTION_DIM)
+    assert torch.isfinite(actions).all()
+    action_seq = actions.view(2, CHUNK_SIZE, ACTION_DIM)
+    assert torch.all(action_seq[..., 2] == 0.0)
+
+
+def test_flowmatching_context_mask_blocks_masked_context_tokens():
+    head = make_flowmatching_head()
+    state = torch.randn(2, STATE_DIM)
+    action_mask = torch.ones(2, ACTION_DIM, dtype=torch.bool)
+    fused = torch.randn(2, 4, EMBED_DIM)
+    context_mask = torch.ones(2, 4, dtype=torch.bool)
+    context_mask[:, -1] = False
+    corrupted = fused.clone()
+    corrupted[:, -1] = 1e4
+
+    torch.manual_seed(0)
+    reference = head.get_action(fused, state=state, action_mask=action_mask, context_mask=context_mask)
+    torch.manual_seed(0)
+    with_garbage = head.get_action(corrupted, state=state, action_mask=action_mask, context_mask=context_mask)
+
+    assert torch.allclose(reference, with_garbage)
+
+
+def test_flowmatching_head_accepts_pooled_2d_context():
+    head = make_flowmatching_head()
+    pred_velocity, noise = head(
+        torch.randn(2, EMBED_DIM),  # pooled (B, E) context from return_cls_only
+        state=torch.randn(2, STATE_DIM),
+        actions_gt=torch.randn(2, CHUNK_SIZE, ACTION_DIM),
+        action_mask=torch.ones(2, CHUNK_SIZE, ACTION_DIM, dtype=torch.bool),
+    )
+    assert pred_velocity.shape == (2, CHUNK_SIZE * ACTION_DIM)
+
+    actions = head.get_action(
+        torch.randn(2, EMBED_DIM),
+        state=torch.randn(2, STATE_DIM),
+        action_mask=torch.ones(2, ACTION_DIM, dtype=torch.bool),
+    )
+    assert actions.shape == (2, CHUNK_SIZE * ACTION_DIM)
+
+
+def test_flowmatching_rejects_out_of_range_embodiment_ids():
+    head = make_flowmatching_head(num_categories=2)
+    with pytest.raises(ValueError, match="num_categories"):
+        head.get_action(
+            torch.randn(2, 4, EMBED_DIM),
+            state=torch.randn(2, STATE_DIM),
+            action_mask=torch.ones(2, ACTION_DIM, dtype=torch.bool),
+            embodiment_id=torch.tensor([0, 5]),
+        )
+
+
+def test_evo1_batched_pixel_values_shape_and_zero_padding():
+    torch.manual_seed(0)
+    batch_size, image_size, max_views = 2, 448, 3
+    camera_images = [torch.rand(batch_size, 3, 40, 50)]  # a single present camera
+    mean = torch.tensor(IMAGENET_MEAN)
+    std = torch.tensor(IMAGENET_STD)
+
+    pixel_values = _batched_pixel_values(
+        camera_images, max_views, image_size, mean, std, torch.float32, torch.device("cpu")
+    )
+
+    assert pixel_values.shape == (batch_size * max_views, 3, image_size, image_size)
+    grouped = pixel_values.reshape(batch_size, max_views, 3, image_size, image_size)
+    # Absent views (indices 1, 2) are zero images, normalized to the constant -mean/std.
+    expected_pad = (-mean / std).view(1, 3, 1, 1)
+    for view in (1, 2):
+        assert torch.allclose(
+            grouped[:, view], expected_pad.expand(batch_size, 3, image_size, image_size), atol=1e-5
+        )
+    # The present view is genuinely different from the constant pad value.
+    assert not torch.allclose(
+        grouped[:, 0], expected_pad.expand(batch_size, 3, image_size, image_size), atol=1e-3
+    )

tests/policies/fastwam/test_fastwam_policy.py

@@ -0,0 +1,391 @@
+#!/usr/bin/env python
+
+# Copyright 2024 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 json
+
+import pytest
+import torch
+from safetensors import safe_open
+from torch import nn
+
+pytest.importorskip("transformers", reason="fastwam requires the `fastwam` extra (transformers)")
+pytest.importorskip("diffusers", reason="fastwam requires the `fastwam` extra (diffusers)")
+
+from lerobot.configs import FeatureType, PolicyFeature, PreTrainedConfig
+from lerobot.policies import FastWAMConfig, get_policy_class, make_policy_config, make_pre_post_processors
+from lerobot.policies.fastwam.modeling_fastwam import FastWAMPolicy
+from lerobot.policies.fastwam.processor_fastwam import FastWAMActionToggleProcessorStep
+from lerobot.utils.constants import ACTION, OBS_STATE
+
+
+class FakeFastWAMCore(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.dit = nn.Linear(2, 2)
+
+    def training_loss(self, sample):
+        assert sample["video"].ndim == 5
+        assert sample["context"].ndim == 3
+        return sample[ACTION].sum() * 0.0 + torch.tensor(1.0), {"loss_action": 1.0}
+
+    def infer_action(self, **kwargs):
+        return {"action": torch.ones(1, kwargs["action_horizon"], 3)}
+
+
+def test_fastwam_is_registered_and_publicly_exported():
+    cfg = make_policy_config(
+        "fastwam",
+        action_dim=3,
+        proprio_dim=2,
+        action_horizon=4,
+        n_action_steps=2,
+        num_video_frames=5,
+        action_video_freq_ratio=1,
+        base_model_id=None,
+    )
+
+    assert isinstance(cfg, FastWAMConfig)
+    assert cfg.type == "fastwam"
+    assert get_policy_class("fastwam") is FastWAMPolicy
+
+
+def test_config_validates_features_model_ids_and_saved_auto_route(tmp_path):
+    cfg = FastWAMConfig()
+    cfg.save_pretrained(tmp_path)
+    saved = json.loads((tmp_path / "config.json").read_text())
+
+    assert saved["pretrained_path"] is None
+    assert cfg.image_features["observation.images.image"].type == FeatureType.VISUAL
+    assert cfg.action_feature.shape == (7,)
+    assert cfg.robot_state_feature.shape == (8,)
+    with pytest.raises(ValueError, match="image feature"):
+        FastWAMConfig(input_features={OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(8,))})
+    assert FastWAMConfig(tokenizer_model_id="somebody/other-tokenizer").tokenizer_model_id == (
+        "somebody/other-tokenizer"
+    )
+
+
+def test_preprocessor_passes_images_through_and_postprocessor_toggles_actions(tmp_path):
+    cfg = FastWAMConfig(
+        action_dim=3,
+        proprio_dim=2,
+        action_horizon=4,
+        n_action_steps=2,
+        num_video_frames=5,
+        action_video_freq_ratio=1,
+        image_size=(2, 2),
+        device="cpu",
+        toggle_action_dimensions=[-1],
+        input_features={
+            "observation.images.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 2, 2)),
+            OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(2,)),
+        },
+        output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(3,))},
+        base_model_id=None,
+    )
+    dataset_stats = {
+        "observation.images.image": {
+            "mean": torch.full((3, 1, 1), 0.2),
+            "std": torch.full((3, 1, 1), 0.1),
+        },
+        OBS_STATE: {
+            "mean": torch.tensor([1.0, 3.0]),
+            "std": torch.tensor([2.0, 4.0]),
+        },
+        ACTION: {
+            "mean": torch.zeros(3),
+            "std": torch.ones(3),
+        },
+    }
+
+    preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_stats)
+    processed = preprocessor(
+        {
+            "observation.images.image": torch.tensor(
+                [
+                    [[0.0, 0.5], [1.0, 0.5]],
+                    [[0.0, 0.5], [1.0, 0.5]],
+                    [[0.0, 0.5], [1.0, 0.5]],
+                ]
+            ),
+            OBS_STATE: torch.tensor([3.0, 7.0]),
+        }
+    )
+    preprocessor.save_pretrained(tmp_path, config_filename="policy_preprocessor.json")
+    postprocessor.save_pretrained(tmp_path, config_filename="policy_postprocessor.json")
+    _, loaded_postprocessor = make_pre_post_processors(cfg, pretrained_path=str(tmp_path))
+
+    # VISUAL normalization is IDENTITY
+    expected_image = torch.tensor(
+        [[[[0.0, 0.5], [1.0, 0.5]], [[0.0, 0.5], [1.0, 0.5]], [[0.0, 0.5], [1.0, 0.5]]]]
+    )
+    assert preprocessor.name == "policy_preprocessor"
+    assert postprocessor.name == "policy_postprocessor"
+    assert torch.allclose(processed["observation.images.image"], expected_image)
+    assert torch.allclose(processed[OBS_STATE], torch.tensor([[1.0, 1.0]]))
+    assert torch.equal(dataset_stats["observation.images.image"]["mean"], torch.full((3, 1, 1), 0.2))
+    assert any(isinstance(step, FastWAMActionToggleProcessorStep) for step in loaded_postprocessor.steps)
+    assert torch.equal(
+        loaded_postprocessor(torch.tensor([[0.25, 0.5, 1.0]])), torch.tensor([[0.25, 0.5, -1.0]])
+    )
+
+
+def test_policy_forward_and_predict_action_adapt_lerobot_batches(monkeypatch):
+    captured = []
+
+    class CapturingCore(FakeFastWAMCore):
+        def infer_action(self, **kwargs):
+            captured.append(
+                {
+                    "image_shape": tuple(kwargs["input_image"].shape),
+                    "proprio_shape": tuple(kwargs["proprio"].shape),
+                    "prompt": kwargs["prompt"],
+                }
+            )
+            return {"action": torch.full((1, kwargs["action_horizon"], 3), float(len(captured)))}
+
+    monkeypatch.setattr(FastWAMPolicy, "_build_core_model", lambda self, config: CapturingCore())
+    cfg = FastWAMConfig(
+        action_dim=3,
+        proprio_dim=2,
+        action_horizon=4,
+        n_action_steps=2,
+        num_video_frames=5,
+        action_video_freq_ratio=1,
+        image_size=(16, 16),
+        input_features={
+            "observation.images.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 16, 16)),
+            OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(2,)),
+        },
+        output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(3,))},
+        base_model_id=None,
+    )
+    policy = FastWAMPolicy(cfg)
+
+    loss, metrics = policy.forward(
+        {
+            "observation.images.image": torch.zeros(1, 3, 16, 16),
+            OBS_STATE: torch.zeros(1, 2),
+            ACTION: torch.zeros(1, 4, 3),
+            "context": torch.zeros(1, 5, 4096),
+            "context_mask": torch.ones(1, 5, dtype=torch.bool),
+        }
+    )
+    action = policy.predict_action_chunk(
+        {
+            "observation.images.image": torch.stack(
+                [
+                    torch.zeros(3, 16, 16),
+                    torch.ones(3, 16, 16),
+                ]
+            ),
+            OBS_STATE: torch.tensor([[0.0, 1.0], [2.0, 3.0]]),
+            "task": ["task 0", "task 1"],
+        }
+    )
+
+    assert loss.item() == 1.0
+    assert metrics["loss_action"] == 1.0
+    assert action.shape == (2, 4, 3)
+    assert action[:, 0, 0].tolist() == [1.0, 2.0]
+    assert [item["image_shape"] for item in captured] == [(1, 3, 16, 16), (1, 3, 16, 16)]
+    assert [item["proprio_shape"] for item in captured] == [(1, 2), (1, 2)]
+    assert [item["prompt"] for item in captured] == [
+        cfg.prompt_template.format(task="task 0"),
+        cfg.prompt_template.format(task="task 1"),
+    ]
+
+
+class CoreWithFrozenComponents(FakeFastWAMCore):
+    """Fake core mirroring the real one: frozen VAE / text encoder held as
+    *unregistered* attributes (via `object.__setattr__`) so they are excluded from
+    `state_dict()` and the saved checkpoint, but still moved by the `_apply` override."""
+
+    def __init__(self):
+        super().__init__()
+        object.__setattr__(self, "vae", nn.Linear(2, 2))
+        object.__setattr__(self, "text_encoder", nn.Linear(2, 2))
+        self.vae.requires_grad_(False)
+        self.text_encoder.requires_grad_(False)
+
+    def _apply(self, fn, *args, **kwargs):
+        super()._apply(fn, *args, **kwargs)
+        self.vae._apply(fn)
+        self.text_encoder._apply(fn)
+        return self
+
+
+def test_from_pretrained_uses_base_loader_and_skips_wan_backbone(monkeypatch, tmp_path):
+    cfg = FastWAMConfig(
+        action_dim=3,
+        proprio_dim=2,
+        action_horizon=4,
+        n_action_steps=2,
+        num_video_frames=5,
+        action_video_freq_ratio=1,
+        base_model_id=None,
+    )
+
+    def build_core(self, config):
+        core = CoreWithFrozenComponents()
+        with torch.no_grad():
+            core.dit.weight.fill_(0.5)
+        return core
+
+    monkeypatch.setattr(FastWAMPolicy, "_build_core_model", build_core)
+
+    reference = FastWAMPolicy(cfg)
+    with torch.no_grad():
+        reference.model.dit.weight.fill_(1.25)  # a distinctive, trained-looking weight
+    reference.save_pretrained(tmp_path)
+
+    # Building from Wan2.2 must never happen on a checkpoint load.
+    def fail_if_wan_pretrained_is_loaded(*args, **kwargs):
+        raise AssertionError("from_pretrained must not initialize or download the Wan2.2 backbone")
+
+    monkeypatch.setattr(
+        "lerobot.policies.fastwam.wan.modular.FastWAM.from_wan22_pretrained",
+        fail_if_wan_pretrained_is_loaded,
+    )
+
+    policy = FastWAMPolicy.from_pretrained(tmp_path)
+
+    assert isinstance(policy.model, CoreWithFrozenComponents)
+    # The bundled checkpoint weights overwrote the freshly built (0.5) DiT weights.
+    assert torch.allclose(policy.model.dit.weight, torch.full_like(policy.model.dit.weight, 1.25))
+
+
+def test_save_pretrained_excludes_frozen_components(monkeypatch, tmp_path):
+    cfg = FastWAMConfig(
+        action_dim=3,
+        proprio_dim=2,
+        action_horizon=4,
+        n_action_steps=2,
+        num_video_frames=5,
+        action_video_freq_ratio=1,
+        base_model_id=None,
+    )
+    monkeypatch.setattr(FastWAMPolicy, "_build_core_model", lambda self, config: CoreWithFrozenComponents())
+    policy = FastWAMPolicy(cfg)
+
+    save_dir = tmp_path / "saved"
+    policy.save_pretrained(save_dir)
+
+    assert (save_dir / "model.safetensors").is_file()
+    # No Wan sidecar files either: the frozen backbone comes from the diffusers repo.
+    assert not (save_dir / "Wan2.2_VAE.safetensors").exists()
+    assert not (save_dir / "google").exists()
+
+    with safe_open(save_dir / "model.safetensors", framework="pt") as f:
+        keys = set(f.keys())
+    # Lean checkpoint: only the trainable DiT is saved; the frozen VAE / UMT5 text
+    # encoder are excluded (loaded from the diffusers/transformers repos at init).
+    assert any(key.startswith("model.dit.") for key in keys)
+    assert not any(key.startswith("model.vae.") for key in keys)
+    assert not any(key.startswith("model.text_encoder.") for key in keys)
+
+
+def test_frozen_components_excluded_from_params_but_follow_device_moves(monkeypatch):
+    cfg = FastWAMConfig(
+        action_dim=3,
+        proprio_dim=2,
+        action_horizon=4,
+        n_action_steps=2,
+        num_video_frames=5,
+        action_video_freq_ratio=1,
+        base_model_id=None,
+    )
+    monkeypatch.setattr(FastWAMPolicy, "_build_core_model", lambda self, config: CoreWithFrozenComponents())
+    policy = FastWAMPolicy(cfg)
+
+    # Unregistered: excluded from state_dict and from the optimizer's parameter set.
+    sd = policy.state_dict()
+    assert not any(k.startswith("model.vae.") or k.startswith("model.text_encoder.") for k in sd)
+    param_names = [n for n, _ in policy.named_parameters()]
+    assert not any("vae" in n or "text_encoder" in n for n in param_names)
+
+    # ...but the `_apply` override still carries them through `.to()` (dtype stands in
+    # for device on a CPU box), so they never strand off the rest of the model.
+    policy.to(torch.float64)
+    assert policy.model.dit.weight.dtype == torch.float64  # registered
+    assert policy.model.vae.weight.dtype == torch.float64  # unregistered, moved via _apply
+    assert policy.model.text_encoder.weight.dtype == torch.float64
+
+
+def test_pretrained_config_round_trips_fastwam_features(tmp_path):
+    cfg = FastWAMConfig(action_dim=7, proprio_dim=8, image_size=(224, 448), base_model_id=None)
+    cfg.save_pretrained(tmp_path)
+
+    loaded = PreTrainedConfig.from_pretrained(tmp_path)
+
+    assert loaded.type == "fastwam"
+    assert loaded.image_features["observation.images.image"].type == FeatureType.VISUAL
+    assert loaded.action_feature.shape == (7,)
+    assert loaded.robot_state_feature.shape == (8,)
+
+
+def test_vae_adapter_empty_build_encode_decode_shapes():
+    """Offline glue check of the diffusers-backed VAE adapter (random weights).
+
+    Validates the encode/decode contract — 48 latent channels, 16x spatial / 4x
+    temporal compression, list-or-batch input, scaling round-trip — without any
+    weight download. (Numerical fidelity vs the original Wan VAE is a separate,
+    GPU + real-weights verification step.)
+    """
+    pytest.importorskip("diffusers")
+    from diffusers import AutoencoderKLWan
+
+    from lerobot.policies.fastwam.wan import WanVideoVAE38
+
+    # Production always loads a real pretrained VAE from the diffusers repo; here we
+    # build the same architecture with random weights and dummy standardization stats
+    # to exercise the adapter's shape/scaling contract offline (fidelity is checked
+    # separately, with real weights, on GPU).
+    arch = {
+        "base_dim": 160,
+        "decoder_base_dim": 256,
+        "z_dim": 48,
+        "dim_mult": [1, 2, 4, 4],
+        "num_res_blocks": 2,
+        "attn_scales": [],
+        "temporal_downsample": [False, True, True],
+        "dropout": 0.0,
+        "is_residual": True,
+        "in_channels": 12,
+        "out_channels": 12,
+        "patch_size": 2,
+        "scale_factor_spatial": 16,
+        "scale_factor_temporal": 4,
+        "clip_output": False,
+        "latents_mean": [0.0] * 48,
+        "latents_std": [1.0] * 48,
+    }
+    raw = AutoencoderKLWan.from_config(arch)
+    vae = WanVideoVAE38(dtype=torch.float32, device="cpu", pretrained=raw)
+    assert vae.z_dim == 48
+    assert vae.upsampling_factor == 16
+    assert vae.temporal_downsample_factor == 4
+
+    video = torch.rand(1, 3, 5, 32, 32) * 2 - 1  # [B,C,T,H,W] in [-1,1]
+    latents = vae.encode(video)
+    assert latents.shape == (1, 48, 2, 2, 2)  # T'=(5-1)//4+1, H'=W'=32//16
+
+    decoded = vae.decode(latents)
+    assert decoded.shape[0] == 1 and decoded.shape[1] == 3 and decoded.shape[-2:] == (32, 32)
+    assert decoded.min() >= -1.0 and decoded.max() <= 1.0
+
+    # list input is accepted and equals the batched path
+    assert torch.equal(vae.encode([video[0]]), latents)

tests/policies/groot/test_groot_lerobot.py

@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Test script for LeRobot's Groot policy forward and inference passes."""
+"""Test script for LeRobot's GR00T N1.7 policy forward and inference passes."""
 
 import gc
 import os
@@ -25,6 +25,8 @@ import numpy as np
 import pytest
 import torch
 
+pytest.importorskip("transformers", reason="groot requires the `groot` extra (transformers)")
+
 from lerobot.policies.groot.configuration_groot import GrootConfig
 from lerobot.policies.groot.modeling_groot import GrootPolicy
 from lerobot.policies.groot.processor_groot import make_groot_pre_post_processors
@@ -33,21 +35,26 @@ from lerobot.types import PolicyAction
 from lerobot.utils.device_utils import auto_select_torch_device
 from tests.utils import require_cuda
 
-pytest.importorskip("transformers")
-
 pytestmark = pytest.mark.skipif(
     os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
     reason="This test requires local Groot installation and is not meant for CI",
 )
 
 
-# Define constants for dummy data
+# Define constants for dummy data (GR00T N1.7 native conventions).
+# N1.7 internally uses a 40-step action chunk, 132-dim state/action, and 256px images
+# (see GrootConfig.__post_init__). Use a chunk-sized action horizon so the dummy batch
+# matches the model's native action space.
 DUMMY_STATE_DIM = 44
 DUMMY_ACTION_DIM = 44
-DUMMY_ACTION_HORIZON = 16
+DUMMY_ACTION_HORIZON = 40
 IMAGE_SIZE = 256
 DEVICE = auto_select_torch_device()
-MODEL_PATH = "aractingi/bimanual-handover-groot-10k"
+# GR00T N1.7 checkpoint (N1.5 is no longer supported). The N1.7-3B base model loads
+# via GrootPolicy.from_pretrained with root-level sharded safetensors.
+MODEL_PATH = "nvidia/GR00T-N1.7-3B"
+# Valid N1.7 embodiment tag carried by the checkpoint metadata.
+EMBODIMENT_TAG = "gr1_unified"
 
 
 def cleanup_memory():
@@ -88,13 +95,13 @@ def instantiate_lerobot_groot(
     PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
     PolicyProcessorPipeline[PolicyAction, PolicyAction],
 ]:
-    """Instantiate LeRobot Groot policy with preprocessor and postprocessor."""
+    """Instantiate LeRobot GR00T N1.7 policy with preprocessor and postprocessor."""
     if from_pretrained:
         policy = GrootPolicy.from_pretrained(
             pretrained_name_or_path=model_path,
             strict=False,
         )
-        policy.config.embodiment_tag = "gr1"
+        policy.config.embodiment_tag = EMBODIMENT_TAG
     else:
         config = GrootConfig(
             base_model_path=model_path,
@@ -102,7 +109,7 @@ def instantiate_lerobot_groot(
             chunk_size=DUMMY_ACTION_HORIZON,
             image_size=[IMAGE_SIZE, IMAGE_SIZE],
             device=DEVICE,
-            embodiment_tag="gr1",
+            embodiment_tag=EMBODIMENT_TAG,
         )
         policy = GrootPolicy(config)
 
@@ -148,8 +155,8 @@ def create_dummy_data(device=DEVICE):
 
 @require_cuda
 def test_lerobot_groot_inference():
-    """Test the inference pass (select_action) of LeRobot's Groot policy."""
-    print("Test: LeRobot Groot Inference Pass")
+    """Test the inference pass (select_action) of LeRobot's GR00T N1.7 policy."""
+    print("Test: LeRobot GR00T N1.7 Inference Pass")
 
     set_seed_all(42)
 
@@ -181,9 +188,9 @@ def test_lerobot_groot_inference():
 
 @require_cuda
 def test_lerobot_groot_forward_pass():
-    """Test the forward pass of LeRobot's Groot policy."""
+    """Test the forward pass of LeRobot's GR00T N1.7 policy."""
     print("\n" + "=" * 50)
-    print("Test: LeRobot Groot Forward Pass (Training Mode)")
+    print("Test: LeRobot GR00T N1.7 Forward Pass (Training Mode)")
 
     set_seed_all(42)
 

tests/policies/groot/test_groot_n1_7_oss_parity.py

@@ -0,0 +1,259 @@
+#!/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 hashlib
+import os
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+from lerobot.policies.groot.action_head.cross_attention_dit import AlternateVLDiT
+from lerobot.policies.groot.groot_n1_7 import GR00TN17
+from lerobot.policies.groot.processor_groot import (
+    GrootN17ActionDecodeStep,
+    GrootN17PackInputsStep,
+    GrootN17VLMEncodeStep,
+    _transform_n1_7_image_for_vlm_albumentations,
+)
+from lerobot.types import TransitionKey
+from lerobot.utils.constants import OBS_STATE
+
+OSS_REFERENCE_COMMIT = "ab88b50c718f6528e1df9dcbaf75865d1b604760"
+
+
+def _fixture_path(filename: str) -> Path:
+    fixture_dir = os.environ.get("GROOT_N17_OSS_PARITY_FIXTURE_DIR")
+    if fixture_dir is None:
+        pytest.skip("Set GROOT_N17_OSS_PARITY_FIXTURE_DIR to run external OSS parity fixtures.")
+    path = Path(fixture_dir) / filename
+    if not path.is_file():
+        pytest.skip(f"External OSS parity fixture not found: {path}")
+    return path
+
+
+def test_groot_n1_7_eval_image_transform_matches_oss_reference():
+    """Match the native N1.7 eval transform for a non-square SO-101 frame."""
+
+    y, x = np.indices((480, 640), dtype=np.uint16)
+    image = np.stack(
+        ((x + 3 * y) % 256, (2 * x + y) % 256, (x + 5 * y) % 256),
+        axis=-1,
+    ).astype(np.uint8)
+    actual = _transform_n1_7_image_for_vlm_albumentations(
+        image,
+        image_crop_size=[230, 230],
+        image_target_size=[256, 256],
+        shortest_image_edge=256,
+        crop_fraction=0.95,
+    )
+
+    assert actual.shape == (256, 340, 3)
+    assert hashlib.sha256(actual.tobytes()).hexdigest() == (
+        "c17e47af68a812aa79db3bb7b64b549ddf10148ac1b204a9686095018561ae9e"
+    )
+
+
+def test_groot_n1_7_vlm_chat_content_order_matches_oss_reference():
+    """Native OSS places all image items before the language item."""
+
+    class RecordingProcessor:
+        def __init__(self):
+            self.content_types = None
+
+        def apply_chat_template(self, conversation, tokenize, add_generation_prompt):
+            assert tokenize is False
+            assert add_generation_prompt is False
+            self.content_types = [item["type"] for item in conversation[0]["content"]]
+            return "rendered"
+
+        def __call__(self, **kwargs):
+            return {}
+
+    processor = RecordingProcessor()
+    step = GrootN17VLMEncodeStep(
+        image_crop_size=[230, 230],
+        image_target_size=[256, 256],
+        shortest_image_edge=256,
+        crop_fraction=0.95,
+        use_albumentations=True,
+        device="cpu",
+    )
+    step._proc = processor
+    transition = {
+        TransitionKey.OBSERVATION: {
+            "video": np.zeros((1, 1, 2, 480, 640, 3), dtype=np.uint8),
+        },
+        TransitionKey.COMPLEMENTARY_DATA: {"language": ["pick up the vial"]},
+    }
+
+    step(transition)
+
+    assert processor.content_types == ["image", "image", "text"]
+
+
+def test_groot_n1_7_alternate_vl_dit_matches_oss_reference():
+    """Run the LeRobot DiT with native OSS weights and identical inputs."""
+
+    pytest.importorskip("diffusers")
+
+    fixture = torch.load(_fixture_path("alternate_vl_dit_small.pt"), map_location="cpu", weights_only=True)
+    model = AlternateVLDiT(
+        output_dim=8,
+        num_attention_heads=2,
+        attention_head_dim=4,
+        num_layers=4,
+        dropout=0.0,
+        final_dropout=False,
+        max_num_positional_embeddings=16,
+        compute_dtype=torch.float32,
+        interleave_self_attention=True,
+        cross_attention_dim=6,
+    ).eval()
+    model.load_state_dict(fixture["state_dict"], strict=True)
+
+    actual = model(
+        hidden_states=fixture["hidden_states"],
+        encoder_hidden_states=fixture["encoder_hidden_states"],
+        timestep=fixture["timestep"],
+        image_mask=fixture["image_mask"],
+        backbone_attention_mask=fixture["backbone_attention_mask"],
+    )
+
+    torch.testing.assert_close(actual, fixture["output"], atol=1e-6, rtol=1e-6)
+
+
+def _state_decode_reference():
+    fixture = np.load(_fixture_path("state_and_action_decode.npz"))
+    raw_stats = {
+        "state": {
+            "single_arm": {"q01": fixture["state_single_arm_q01"], "q99": fixture["state_single_arm_q99"]},
+            "gripper": {"q01": fixture["state_gripper_q01"], "q99": fixture["state_gripper_q99"]},
+        },
+        "action": {
+            "single_arm": {"q01": fixture["action_single_arm_q01"], "q99": fixture["action_single_arm_q99"]},
+            "gripper": {"q01": fixture["action_gripper_q01"], "q99": fixture["action_gripper_q99"]},
+        },
+        "relative_action": {
+            "single_arm": {
+                "min": fixture["relative_single_arm_min"],
+                "max": fixture["relative_single_arm_max"],
+            },
+        },
+    }
+    for modality_stats in raw_stats.values():
+        for entry in modality_stats.values():
+            for key, value in entry.items():
+                if isinstance(value, np.ndarray):
+                    entry[key] = value.tolist()
+    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},
+            ],
+        },
+    }
+    state_min = np.concatenate((fixture["state_single_arm_q01"], fixture["state_gripper_q01"]))
+    state_max = np.concatenate((fixture["state_single_arm_q99"], fixture["state_gripper_q99"]))
+    pack_step = GrootN17PackInputsStep(
+        normalize_min_max=True,
+        stats={OBS_STATE: {"min": state_min, "max": state_max}},
+        raw_stats=raw_stats,
+        modality_config=modality_config,
+        use_percentiles=True,
+    )
+    raw_state = np.concatenate((fixture["state_single_arm"], fixture["state_gripper"]), axis=-1)
+    transition = {
+        TransitionKey.OBSERVATION: {OBS_STATE: torch.from_numpy(raw_state)},
+        TransitionKey.COMPLEMENTARY_DATA: {},
+    }
+    packed = pack_step(transition)
+    return fixture, raw_stats, modality_config, pack_step, packed
+
+
+def test_groot_n1_7_state_normalization_matches_oss_checkpoint_reference():
+    fixture, _raw_stats, _modality_config, _pack_step, packed = _state_decode_reference()
+    expected = np.concatenate(
+        (fixture["normalized_state_single_arm"], fixture["normalized_state_gripper"]), axis=-1
+    )
+
+    actual = packed[TransitionKey.OBSERVATION]["state"][:, 0, :6]
+
+    torch.testing.assert_close(actual, torch.from_numpy(expected), atol=1e-6, rtol=1e-6)
+
+
+def test_groot_n1_7_relative_action_decode_matches_oss_checkpoint_reference():
+    fixture, raw_stats, modality_config, pack_step, _packed = _state_decode_reference()
+    decode_step = GrootN17ActionDecodeStep(
+        env_action_dim=6,
+        raw_stats=raw_stats,
+        modality_config=modality_config,
+        use_percentiles=True,
+        use_relative_action=True,
+        pack_step=pack_step,
+    )
+    decoded = decode_step({TransitionKey.ACTION: torch.from_numpy(fixture["normalized_action"])})[
+        TransitionKey.ACTION
+    ]
+    expected = np.concatenate((fixture["decoded_single_arm"], fixture["decoded_gripper"]), axis=-1).astype(
+        np.float32
+    )
+
+    torch.testing.assert_close(decoded, torch.from_numpy(expected), atol=1e-5, rtol=1e-5)
+
+
+def test_groot_n1_7_qwen_backbone_matches_oss_checkpoint_reference():
+    """Compare the actual 3B checkpoint backbone when explicitly enabled."""
+
+    checkpoint = os.environ.get("GROOT_N17_PARITY_CHECKPOINT")
+    if checkpoint is None:
+        pytest.skip("Set GROOT_N17_PARITY_CHECKPOINT to run the 3B OSS Qwen parity test.")
+    if not torch.cuda.is_available():
+        pytest.skip("The 3B OSS Qwen parity test requires CUDA.")
+
+    pytest.importorskip("transformers")
+
+    from transformers.feature_extraction_utils import BatchFeature
+
+    fixture = torch.load(_fixture_path("qwen_backbone_so101.pt"), map_location="cpu", weights_only=True)
+    model = GR00TN17.from_pretrained(checkpoint).to(device="cuda", dtype=torch.bfloat16).eval()
+    backbone_input = BatchFeature(
+        data={
+            key.removeprefix("input."): value.to("cuda")
+            for key, value in fixture.items()
+            if key.startswith("input.")
+        }
+    )
+
+    with torch.inference_mode():
+        actual = model.backbone(backbone_input)
+
+    feature_error = (
+        actual.backbone_features.cpu().float() - fixture["output.backbone_features"].float()
+    ).abs()
+    # Native OSS and LeRobot use different Torch/Transformers/Flash-Attention releases.
+    # Require the measured BF16 accumulation envelope while rejecting structural drift.
+    assert feature_error.mean().item() <= 0.04
+    assert feature_error.max().item() <= 2.0
+    torch.testing.assert_close(
+        actual.backbone_attention_mask.cpu(), fixture["output.backbone_attention_mask"]
+    )
+    torch.testing.assert_close(actual.image_mask.cpu(), fixture["output.image_mask"])

tests/policies/groot/test_groot_state_dropout.py

@@ -0,0 +1,100 @@
+#!/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.
+
+"""Isaac-GR00T N1.7 raw-state dropout training contract.
+
+Isaac-GR00T zeroes the entire proprioceptive state of a sample with probability
+``state_dropout_prob`` (configured in the checkpoint's processor sidecar) during
+training only. Baseline LeRobot kept the processor deterministic, so this
+regularization never activated. These tests pin the train/eval split.
+"""
+
+import torch
+
+from lerobot.policies.groot.processor_groot import GrootN17PackInputsStep
+from lerobot.types import TransitionKey
+from lerobot.utils.constants import OBS_STATE
+
+
+def _make_transition():
+    return {
+        TransitionKey.OBSERVATION: {OBS_STATE: torch.tensor([[1.0, 2.0], [3.0, 4.0]])},
+        TransitionKey.COMPLEMENTARY_DATA: {"task": ["Move", "Move"]},
+    }
+
+
+def test_groot_n1_7_training_applies_raw_state_dropout_before_encoder():
+    step = GrootN17PackInputsStep(
+        max_state_dim=4,
+        max_action_dim=4,
+        normalize_min_max=False,
+        training=True,
+        state_dropout_prob=1.0,
+    )
+
+    output = step(_make_transition())
+
+    expected = torch.zeros(2, 1, 4)
+    torch.testing.assert_close(output[TransitionKey.OBSERVATION]["state"], expected)
+
+
+def test_groot_n1_7_training_state_dropout_is_disabled_under_no_grad():
+    step = GrootN17PackInputsStep(
+        max_state_dim=4,
+        max_action_dim=4,
+        normalize_min_max=False,
+        training=True,
+        state_dropout_prob=1.0,
+    )
+
+    with torch.no_grad():
+        output = step(_make_transition())
+
+    expected = torch.tensor([[[1.0, 2.0, 0.0, 0.0]], [[3.0, 4.0, 0.0, 0.0]]])
+    torch.testing.assert_close(output[TransitionKey.OBSERVATION]["state"], expected)
+
+
+def test_groot_n1_7_eval_mode_state_dropout_is_inactive():
+    step = GrootN17PackInputsStep(
+        max_state_dim=4,
+        max_action_dim=4,
+        normalize_min_max=False,
+        training=False,
+        state_dropout_prob=1.0,
+    )
+
+    output = step(_make_transition())
+
+    expected = torch.tensor([[[1.0, 2.0, 0.0, 0.0]], [[3.0, 4.0, 0.0, 0.0]]])
+    torch.testing.assert_close(output[TransitionKey.OBSERVATION]["state"], expected)
+
+
+def test_groot_n1_7_pack_step_serializes_dropout_prob_but_not_training_mode():
+    step = GrootN17PackInputsStep(
+        max_state_dim=4,
+        max_action_dim=4,
+        normalize_min_max=False,
+        training=True,
+        state_dropout_prob=0.2,
+    )
+
+    serialized = step.get_config()
+    restored = GrootN17PackInputsStep(**serialized)
+
+    assert "training" not in serialized
+    assert serialized["state_dropout_prob"] == 0.2
+    assert restored.training is False
+    assert restored.state_dropout_prob == 0.2

tests/policies/groot/test_groot_train_random_crop.py

@@ -0,0 +1,169 @@
+#!/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.
+
+"""Isaac-GR00T N1.7 train-time random crop contract (crop geometry only).
+
+Isaac-GR00T crops a random ``crop_fraction`` window during training and the
+deterministic center window at eval, replaying the sampled window across all
+camera views of a sample (gr00t/data/transform/video.py, n1.5-release onward:
+"If mode is 'train', return a random crop transform. If mode is 'eval', return
+a center crop transform."). This mirrors LeRobot's own Diffusion/VQBeT
+``crop_is_random`` pattern. Color jitter is intentionally out of scope here.
+"""
+
+import random
+
+import numpy as np
+import torch
+
+from lerobot.policies.groot.processor_groot import (
+    GrootN17VLMEncodeStep,
+    _transform_n1_7_image_for_vlm_albumentations,
+)
+
+
+def _structured_image(h=480, w=640):
+    yy, xx = np.mgrid[0:h, 0:w]
+    return np.stack([(xx * 255 / w), (yy * 255 / h), ((xx + yy) * 255 / (h + w))], axis=-1).astype(np.uint8)
+
+
+def test_crop_position_none_is_bitexact_center_crop():
+    """crop_position=None must remain byte-identical to the pre-change eval path."""
+    img = _structured_image()
+    ref = _transform_n1_7_image_for_vlm_albumentations(
+        img,
+        image_crop_size=None,
+        image_target_size=[256, 256],
+        shortest_image_edge=256,
+        crop_fraction=0.95,
+    )
+    out = _transform_n1_7_image_for_vlm_albumentations(
+        img,
+        image_crop_size=None,
+        image_target_size=[256, 256],
+        shortest_image_edge=256,
+        crop_fraction=0.95,
+        crop_position=None,
+    )
+    np.testing.assert_array_equal(ref, out)
+
+
+def test_crop_position_center_matches_center_crop():
+    img = _structured_image()
+    center = _transform_n1_7_image_for_vlm_albumentations(
+        img,
+        image_crop_size=None,
+        image_target_size=[256, 256],
+        shortest_image_edge=256,
+        crop_fraction=0.95,
+        crop_position=None,
+    )
+    explicit = _transform_n1_7_image_for_vlm_albumentations(
+        img,
+        image_crop_size=None,
+        image_target_size=[256, 256],
+        shortest_image_edge=256,
+        crop_fraction=0.95,
+        crop_position=(0.5, 0.5),
+    )
+    # int-floor center vs rounded positional center may differ by <=1 px of grid
+    assert center.shape == explicit.shape
+    diff = np.abs(center.astype(np.int16) - explicit.astype(np.int16))
+    assert diff.mean() < 3.0
+
+
+def test_crop_position_corners_differ_from_center():
+    img = _structured_image()
+
+    def crop_at(position):
+        return _transform_n1_7_image_for_vlm_albumentations(
+            img,
+            image_crop_size=None,
+            image_target_size=[256, 256],
+            shortest_image_edge=256,
+            crop_fraction=0.95,
+            crop_position=position,
+        )
+
+    center = crop_at(None)
+    tl = crop_at((0.0, 0.0))
+    br = crop_at((1.0, 1.0))
+    assert not np.array_equal(center, tl)
+    assert not np.array_equal(tl, br)
+
+
+def _video(img, views=2):
+    return np.stack([img] * views, axis=0).reshape(1, 1, views, *img.shape)
+
+
+def _step(training):
+    return GrootN17VLMEncodeStep(
+        image_target_size=[256, 256],
+        shortest_image_edge=256,
+        crop_fraction=0.95,
+        use_albumentations=True,
+        training=training,
+    )
+
+
+def test_training_crop_replays_one_window_across_views():
+    video = _video(_structured_image())
+    frames = _step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0]
+    np.testing.assert_array_equal(np.asarray(frames[0]), np.asarray(frames[1]))
+
+
+def test_training_crop_differs_from_eval_center_crop():
+    video = _video(_structured_image())
+    random.seed(3)  # a draw that is not the exact center
+    train_frame = np.asarray(
+        _step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0][0]
+    )
+    eval_frame = np.asarray(
+        _step(training=False)._build_sample_images(video, batch_size=1, target_device=None)[0][0]
+    )
+    assert not np.array_equal(train_frame, eval_frame)
+
+
+def test_training_crop_is_disabled_under_no_grad():
+    video = _video(_structured_image())
+    with torch.no_grad():
+        no_grad_frame = np.asarray(
+            _step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0][0]
+        )
+    eval_frame = np.asarray(
+        _step(training=False)._build_sample_images(video, batch_size=1, target_device=None)[0][0]
+    )
+    np.testing.assert_array_equal(no_grad_frame, eval_frame)
+
+
+def test_training_mode_is_not_serialized():
+    step = _step(training=True)
+    serialized = step.get_config()
+    assert "training" not in serialized
+    restored = GrootN17VLMEncodeStep(**serialized)
+    assert restored.training is False
+
+
+def test_training_crop_respects_global_seed():
+    video = _video(_structured_image())
+
+    def draw():
+        random.seed(11)
+        return np.asarray(
+            _step(training=True)._build_sample_images(video, batch_size=1, target_device=None)[0][0]
+        )
+
+    np.testing.assert_array_equal(draw(), draw())

tests/policies/groot/test_groot_training_optim_contract.py

@@ -0,0 +1,125 @@
+#!/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.
+
+"""Isaac-GR00T N1.7 optimizer/scheduler/precision training contract.
+
+Pins the LeRobot GR00T fine-tuning recipe to the native Isaac-GR00T contract:
+AdamW(lr=1e-4, betas=(0.9, 0.999), eps=1e-8, weight_decay=1e-5, grad clip 1.0),
+HF cosine schedule with ~5% warmup over the actual update count, FP32 master
+parameters under BF16 autocast, transformers-style weight-decay grouping, the
+frozen LM-head weight tie, and episode-tail exclusion for incomplete chunks.
+"""
+
+import pytest
+import torch
+
+from lerobot.optim.schedulers import DiffuserSchedulerConfig
+from lerobot.policies.groot.configuration_groot import GrootConfig
+from lerobot.policies.groot.groot_n1_7 import _tie_unused_qwen_lm_head
+from lerobot.policies.groot.modeling_groot import GrootPolicy
+
+
+def test_groot_n1_7_optimizer_matches_isaac_training_contract():
+    optimizer = GrootConfig().get_optimizer_preset()
+
+    assert optimizer.lr == pytest.approx(1e-4)
+    assert optimizer.betas == pytest.approx((0.9, 0.999))
+    assert optimizer.eps == pytest.approx(1e-8)
+    assert optimizer.weight_decay == pytest.approx(1e-5)
+    assert optimizer.grad_clip_norm == pytest.approx(1.0)
+
+
+def test_groot_n1_7_sampler_excludes_incomplete_action_tails():
+    config = GrootConfig(chunk_size=16, n_action_steps=16)
+
+    assert len(config.action_delta_indices) == 16
+    assert config.drop_n_last_frames == 15
+
+
+def test_groot_n1_7_scheduler_matches_isaac_hf_cosine_contract():
+    pytest.importorskip("diffusers", reason="the scheduler preset requires the `groot` extra (diffusers)")
+    config = GrootConfig(max_steps=20_000)
+    scheduler_config = config.get_scheduler_preset()
+
+    assert isinstance(scheduler_config, DiffuserSchedulerConfig)
+    assert scheduler_config.name == "cosine"
+    assert scheduler_config.num_warmup_steps == 1_000
+
+    parameter = torch.nn.Parameter(torch.ones(()))
+    optimizer = torch.optim.AdamW([parameter], lr=config.optimizer_lr)
+    scheduler = scheduler_config.build(optimizer, num_training_steps=20_000)
+    lr_factor = scheduler.lr_lambdas[0]
+
+    assert lr_factor(0) == pytest.approx(0.0)
+    assert lr_factor(1_000) == pytest.approx(1.0)
+    assert lr_factor(10_500) == pytest.approx(0.5)
+    assert lr_factor(20_000) == pytest.approx(0.0, abs=1e-12)
+
+
+def test_groot_n1_7_scheduler_rounds_fractional_warmup_up_like_transformers():
+    scheduler_config = GrootConfig(max_steps=777).get_scheduler_preset()
+
+    assert scheduler_config.num_warmup_steps == 39
+
+
+def test_groot_n1_7_model_parameters_use_fp32_checkpoint_and_optimizer_precision():
+    module = torch.nn.Module()
+    module.trainable = torch.nn.Parameter(torch.ones(3, dtype=torch.bfloat16))
+    module.frozen = torch.nn.Parameter(torch.ones(3, dtype=torch.bfloat16), requires_grad=False)
+
+    GrootPolicy._cast_model_parameters_to_fp32(module)
+
+    assert module.trainable.dtype == torch.float32
+    assert module.frozen.dtype == torch.float32
+
+
+def test_groot_n1_7_ties_unused_qwen_lm_head_to_frozen_input_embeddings():
+    class DummyQwen(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.embed_tokens = torch.nn.Embedding(7, 3)
+            self.lm_head = torch.nn.Linear(3, 7, bias=False)
+
+        def get_input_embeddings(self):
+            return self.embed_tokens
+
+    model = DummyQwen()
+    _tie_unused_qwen_lm_head(model)
+
+    assert model.lm_head.weight is model.embed_tokens.weight
+    assert len(list(model.parameters())) == 1
+
+
+def test_groot_n1_7_optimizer_groups_match_transformers_weight_decay_rules():
+    pytest.importorskip(
+        "transformers", reason="weight-decay grouping requires the `groot` extra (transformers)"
+    )
+    module = torch.nn.Module()
+    module.linear = torch.nn.Linear(3, 2)
+    module.norm = torch.nn.LayerNorm(2)
+    module.frozen = torch.nn.Parameter(torch.ones(1), requires_grad=False)
+
+    groups = GrootPolicy._build_weight_decay_parameter_groups(module)
+
+    assert len(groups) == 2
+    assert "weight_decay" not in groups[0]
+    assert groups[1]["weight_decay"] == 0.0
+    assert groups[0]["params"] == [module.linear.weight]
+    assert {id(parameter) for parameter in groups[1]["params"]} == {
+        id(module.linear.bias),
+        id(module.norm.weight),
+        id(module.norm.bias),
+    }

tests/policies/groot/test_groot_vs_original.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+# 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.
@@ -14,431 +14,194 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Test script to verify Groot policy integration with LeRobot vs the original implementation, only meant to be run locally!"""
+"""Parity test: original NVIDIA GR00T N1.7 vs the GR00T N1.7 integration in LeRobot.
+
+Verifies that the self-contained LeRobot reimplementation of the GR00T N1.7 action
+head + Qwen3-VL backbone produces the SAME raw model output (``action_pred``, the
+normalized flow-matching prediction before any action decoding) as NVIDIA's original
+``gr00t`` package, given byte-identical pre-processed inputs and the same
+flow-matching seed. The comparison is parametrized over every embodiment tag present
+in the checkpoint.
+
+To keep the comparison fair, the original outputs + the exact collated inputs are
+produced once per embodiment in the original ``gr00t`` env via the companion script
+``utils/dump_original_n1_7.py`` (in the ``utils`` package next to this file) and saved
+to per-tag ``.npz`` files.
+This test discovers those artifacts, replays the identical inputs through the LeRobot
+model, and compares.
+
+This test is LOCAL-only and skips on CI, when ``gr00t``-side prerequisites are not
+present, or when no artifact has been generated. By default it looks for artifacts in
+``<this dir>/artifacts/``; override with ``GROOT_N1_7_PARITY_DIR``. See the
+"Original-vs-LeRobot parity test" section of ``src/lerobot/policies/groot/README.md``
+for the full run procedure.
+"""
 
-import gc
 import os
-from copy import deepcopy
-from typing import Any
+from pathlib import Path
 
 import numpy as np
 import pytest
 import torch
 
-from lerobot.policies.groot.configuration_groot import GrootConfig
-from lerobot.policies.groot.modeling_groot import GrootPolicy
-from lerobot.policies.groot.processor_groot import make_groot_pre_post_processors
-from lerobot.processor import PolicyProcessorPipeline
-from lerobot.types import PolicyAction
-
-pytest.importorskip("gr00t")
-pytest.importorskip("transformers")
-
 pytestmark = pytest.mark.skipif(
     os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
-    reason="This test requires local Groot installation and is not meant for CI",
+    reason="Requires a local GR00T N1.7 checkpoint + pre-generated artifacts; not for CI.",
 )
 
+from lerobot.policies.groot.configuration_groot import GROOT_N1_7  # noqa: E402,F401
 
-from gr00t.data.dataset import ModalityConfig  # noqa: E402
-from gr00t.data.embodiment_tags import EmbodimentTag  # noqa: E402
-from gr00t.data.transform.base import ComposedModalityTransform  # noqa: E402
-from gr00t.model.policy import Gr00tPolicy  # noqa: E402
+SEED = 42
+DEVICE = os.environ.get("GROOT_PARITY_DEVICE", "cuda" if torch.cuda.is_available() else "cpu")
+ATOL = float(os.environ.get("GROOT_PARITY_ATOL", "1e-3"))
+RTOL = float(os.environ.get("GROOT_PARITY_RTOL", "1e-3"))
 
-# GR1 humanoid dimensions (from pretrained model metadata)
-# The actual GR1 robot has 44 dimensions for both state and action
-# GR00TTransform will pad state to 64 and truncate action to 32
-DUMMY_STATE_DIM = 44
-DUMMY_ACTION_DIM = 44
-DUMMY_ACTION_HORIZON = 16
-IMAGE_SIZE = 256
-DEVICE = "cpu"
-MODEL_PATH = "nvidia/GR00T-N1.5-3B"
-
-GR1_BODY_PARTS = {
-    "left_arm": 7,
-    "left_hand": 6,
-    "left_leg": 6,
-    "neck": 3,
-    "right_arm": 7,
-    "right_hand": 6,
-    "right_leg": 6,
-    "waist": 3,
-}
+# Artifact filenames are original_n1_7_<embodiment_tag>.npz
+_ARTIFACT_PREFIX = "original_n1_7_"
+_ARTIFACT_SUFFIX = ".npz"
 
 
-def cleanup_memory():
-    """Clean up GPU/MPS memory to prevent OOM errors between tests."""
-    print("\nCleaning up memory...")
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
-    if torch.backends.mps.is_available():
-        torch.mps.empty_cache()
-    print("Memory cleanup complete.")
+def _artifact_dir() -> Path:
+    """Directory holding the per-embodiment .npz artifacts.
 
-
-def set_seed_all(seed: int):
-    """Set random seed for all RNG sources to ensure reproducibility."""
-    import random
-
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed(seed)
-        torch.cuda.manual_seed_all(seed)
-
-    # Set deterministic behavior
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-    torch.use_deterministic_algorithms(True, warn_only=True)
-
-
-def instantiate_lerobot_groot(
-    from_pretrained: bool = False,
-    model_path: str = MODEL_PATH,
-) -> tuple[
-    GrootPolicy,
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """Instantiate LeRobot Groot policy with preprocessor and postprocessor."""
-    if from_pretrained:
-        policy = GrootPolicy.from_pretrained(
-            pretrained_name_or_path=model_path,
-            strict=False,
-        )
-        policy.config.embodiment_tag = "gr1"
-    else:
-        config = GrootConfig(
-            base_model_path=model_path,
-            n_action_steps=DUMMY_ACTION_HORIZON,
-            chunk_size=DUMMY_ACTION_HORIZON,
-            image_size=[IMAGE_SIZE, IMAGE_SIZE],
-            device=DEVICE,
-            embodiment_tag="gr1",
-        )
-        policy = GrootPolicy(config)
-
-    policy.to(DEVICE)
-    policy.config.device = DEVICE
-
-    preprocessor, postprocessor = make_groot_pre_post_processors(
-        config=policy.config,
-        dataset_stats=None,  # Pass None for dataset_stats to disable normalization (original GR00T doesn't normalize)
-    )
-
-    return (policy, preprocessor, postprocessor)
-
-
-def instantiate_original_groot(
-    from_pretrained: bool = False,
-    model_path: str = MODEL_PATH,
-):
-    """Instantiate original Groot policy from NVIDIA's implementation."""
-    from gr00t.data.transform.concat import ConcatTransform
-    from gr00t.data.transform.state_action import StateActionToTensor
-    from gr00t.data.transform.video import VideoToNumpy, VideoToTensor
-    from gr00t.model.transforms import GR00TTransform
-
-    video_keys = ["video.ego_view"]
-    state_keys = [
-        "state"
-    ]  # Important: Use single concatenated "state" key (not split body parts) to match preprocessing
-    action_keys = [
-        "action.left_arm",
-        "action.right_arm",
-        "action.left_hand",
-        "action.right_hand",
-        "action.left_leg",
-        "action.right_leg",
-        "action.neck",
-        "action.waist",
-    ]
-    language_keys = ["annotation.human.action.task_description"]
-
-    modality_config = {
-        "video": ModalityConfig(
-            delta_indices=[0],  # Current frame only
-            modality_keys=video_keys,
-        ),
-        "state": ModalityConfig(
-            delta_indices=[0],
-            modality_keys=state_keys,
-        ),
-        "action": ModalityConfig(
-            delta_indices=list(range(DUMMY_ACTION_HORIZON)),
-            modality_keys=action_keys,
-        ),
-        "language": ModalityConfig(
-            delta_indices=[0],
-            modality_keys=language_keys,
-        ),
-    }
-
-    modality_transform = ComposedModalityTransform(
-        transforms=[
-            VideoToTensor(apply_to=video_keys),
-            VideoToNumpy(apply_to=video_keys),  # Convert to numpy (GR00TTransform expects numpy arrays)
-            # State is already a single concatenated key, so no StateActionToTensor needed
-            # Convert action from numpy to tensor
-            StateActionToTensor(apply_to=action_keys),
-            # Concatenate only video and actions (state is already single key)
-            ConcatTransform(
-                video_concat_order=video_keys,
-                state_concat_order=[],  # Empty:state is already single key
-                action_concat_order=action_keys,
-            ),
-            GR00TTransform(
-                max_state_dim=64,
-                max_action_dim=32,
-                state_horizon=1,
-                action_horizon=DUMMY_ACTION_HORIZON,
-                training=False,
-            ),
-        ]
-    )
-
-    policy = Gr00tPolicy(
-        model_path=model_path,
-        embodiment_tag=EmbodimentTag.GR1,
-        modality_config=modality_config,
-        modality_transform=modality_transform,
-        device=DEVICE,
-    )
-
-    return policy, modality_config, modality_transform
-
-
-def create_dummy_data(device=DEVICE):
-    """Create dummy data for testing both implementations."""
-    batch_size = 2
-    prompt = "Pick up the red cube and place it in the bin"
-    state = torch.randn(batch_size, DUMMY_STATE_DIM, dtype=torch.float32, device=device)
-
-    batch = {
-        "observation.state": state,
-        "action": torch.randn(
-            batch_size,
-            DUMMY_ACTION_HORIZON,
-            DUMMY_ACTION_DIM,
-            dtype=torch.float32,
-            device=device,  # Action ground truth (for training)
-        ),
-        "observation.images.ego_view": torch.rand(
-            batch_size,
-            3,
-            IMAGE_SIZE,
-            IMAGE_SIZE,
-            dtype=torch.float32,
-            device=device,  # Images in [0, 1] range as expected by LeRobot
-        ),
-        "task": [prompt for _ in range(batch_size)],
-    }
-
-    return batch
-
-
-def convert_lerobot_to_original_format(batch, modality_config):
-    """Convert LeRobot batch format to original Groot format.
-
-    The original Groot expects observations in this format:
-    {
-        "video.<camera_name>": np.ndarray (T, H, W, C) or (B, T, H, W, C)
-        "state.<state_component>": np.ndarray (T, D) or (B, T, D)
-        "action.<action_component>": np.ndarray (T, D) or (B, T, D)
-        "annotation.<annotation_type>": str or list[str]
-    }
+    Self-contained by default: a sibling ``artifacts/`` directory next to this test.
+    Override with ``GROOT_N1_7_PARITY_DIR`` (e.g. to point at a scratch location).
+    The directory is read-only here -- it is populated by ``utils/dump_original_n1_7.py``
+    run in the original gr00t environment; the test never creates it.
     """
-    # Original Groot expects (T, H, W, C) format for images
-    # LeRobot has (B, C, H, W) format, so we need to convert
-    observation = {}
-
-    for img_key in ["ego_view"]:
-        lerobot_key = f"observation.images.{img_key}"
-        if lerobot_key in batch:
-            img = batch[lerobot_key]
-            # Convert from (B, C, H, W) to (B, T=1, H, W, C)
-            img_np = img.permute(0, 2, 3, 1).unsqueeze(1).cpu().numpy()
-            # Convert [0, 1] to [0, 255] uint8 as expected by original
-            img_np = (img_np * 255).astype(np.uint8)
-            observation[f"video.{img_key}"] = img_np
-
-    # Important: The Original's GR00TTransform expects "state" as (B, T, D), not split body parts
-    if "observation.state" in batch:
-        state = batch["observation.state"]
-        state_np = state.unsqueeze(1).cpu().numpy()  # (B, 1, D)
-        observation["state"] = state_np
-
-    if "action" in batch:
-        action = batch["action"]
-        action_np = action.cpu().numpy()
-
-        start_idx = 0
-        for part_name, part_dim in GR1_BODY_PARTS.items():
-            end_idx = start_idx + part_dim
-            observation[f"action.{part_name}"] = action_np[:, :, start_idx:end_idx]
-            start_idx = end_idx
-
-    if "task" in batch:
-        task_list = batch["task"]
-        # GR00TTransform expects language with (B, T) shape for batched data
-        # Create a (B, T=1) array where each element is the string directly
-        bsz = len(task_list)
-        task_array = np.empty((bsz, 1), dtype=object)
-        for i in range(bsz):
-            task_array[i, 0] = task_list[i]  # Assign string directly to each (i, 0) position
-        observation["annotation.human.action.task_description"] = task_array
-
-    return observation
+    env = os.environ.get("GROOT_N1_7_PARITY_DIR")
+    if env:
+        return Path(env)
+    return Path(__file__).resolve().parent / "artifacts"
 
 
-def test_groot_original_vs_lerobot_pretrained():
-    """Test Groot original implementation vs LeRobot implementation with pretrained weights."""
-    print("Test: Groot Original vs LeRobot with Pretrained Weights (Inference)")
+def _discover_artifacts() -> list[tuple[str, Path]]:
+    """Return [(embodiment_tag, npz_path), ...] for every dumped artifact."""
+    d = _artifact_dir()
+    if not d.is_dir():
+        return []
+    out = []
+    for p in sorted(d.glob(f"{_ARTIFACT_PREFIX}*{_ARTIFACT_SUFFIX}")):
+        tag = p.name[len(_ARTIFACT_PREFIX) : -len(_ARTIFACT_SUFFIX)]
+        out.append((tag, p))
+    return out
 
-    set_seed_all(42)
 
-    lerobot_policy, lerobot_preprocessor, lerobot_postprocessor = instantiate_lerobot_groot(
-        from_pretrained=True
+def _resolve_checkpoint() -> str:
+    env = os.environ.get("GROOT_N1_7_LIBERO_CKPT")
+    if env:
+        if not Path(env).exists():
+            pytest.skip(f"GROOT_N1_7_LIBERO_CKPT={env} does not exist")
+        return env
+    try:
+        from huggingface_hub import snapshot_download
+
+        root = snapshot_download(
+            "nvidia/GR00T-N1.7-LIBERO",
+            local_files_only=True,
+            allow_patterns=["libero_10/*"],
+        )
+    except Exception as exc:  # noqa: BLE001
+        pytest.skip(f"GR00T N1.7 LIBERO checkpoint not available locally: {exc}")
+    ckpt = Path(root) / "libero_10"
+    if not (ckpt / "config.json").exists():
+        pytest.skip(f"GR00T N1.7 LIBERO checkpoint incomplete at {ckpt}")
+    return str(ckpt)
+
+
+def _load_artifact(path: Path):
+    data = np.load(path, allow_pickle=True)
+    original_action = torch.from_numpy(data["action_pred"]).float()
+    dtypes = dict(zip(data["meta_keys"].tolist(), data["meta_dtypes"].tolist(), strict=False))
+    inputs = {}
+    for key in data.files:
+        if not key.startswith("in::"):
+            continue
+        name = key[4:]
+        arr = data[key]
+        t = torch.from_numpy(np.asarray(arr))
+        declared = dtypes.get(key, "")
+        if "int" in declared or "long" in declared:
+            t = t.long()
+        inputs[name] = t
+    return original_action, inputs
+
+
+def _unflatten(inputs: dict[str, torch.Tensor]) -> dict:
+    """Rebuild the nested model-input dict from dot-prefixed flat keys."""
+    nested: dict = {}
+    for dotted, value in inputs.items():
+        parts = dotted.split(".")
+        cur = nested
+        for p in parts[:-1]:
+            cur = cur.setdefault(p, {})
+        cur[parts[-1]] = value
+    return nested.get("inputs", nested)
+
+
+@pytest.fixture(scope="module")
+def lerobot_model():
+    """Load the LeRobot GR00T N1.7 model once (fp32 + SDPA) and reuse across tags."""
+    ckpt = _resolve_checkpoint()
+    from lerobot.policies.groot.groot_n1_7 import GR00TN17
+
+    model = GR00TN17.from_pretrained(
+        ckpt,
+        tune_llm=False,
+        tune_visual=False,
+        tune_projector=False,
+        tune_diffusion_model=False,
+        tune_vlln=False,
+        transformers_loading_kwargs={"trust_remote_code": True},
     )
-    original_policy, modality_config, modality_transform = instantiate_original_groot(from_pretrained=True)
+    # fp32 + SDPA on both sides: bf16 + differing attention kernels otherwise introduce
+    # ~1e-2 numerical noise unrelated to the implementations.
+    model.compute_dtype = "float32"
+    model.config.compute_dtype = model.compute_dtype
+    model.to(device=DEVICE, dtype=torch.float32)
+    model.eval()
+    return model
 
-    batch = create_dummy_data()
-    batch_lerobot = deepcopy(batch)
 
-    print("\n[LeRobot] Running inference...")
-    lerobot_policy.eval()
-    batch_lerobot_processed = lerobot_preprocessor(batch_lerobot)
+_ARTIFACTS = _discover_artifacts()
 
-    # Important: Reset seed immediately before inference to ensure identical RNG state
-    torch.manual_seed(42)
 
-    with torch.no_grad():
-        lerobot_actions = lerobot_policy.select_action(batch_lerobot_processed)
+@pytest.mark.skipif(
+    not _ARTIFACTS,
+    reason=(
+        "No GR00T N1.7 parity artifacts found. Generate them first in the original gr00t "
+        "env:\n  .venv-original/bin/python tests/policies/groot/utils/dump_original_n1_7.py "
+        "--ckpt <ckpt> --out-dir tests/policies/groot/artifacts --device cuda"
+    ),
+)
+@pytest.mark.parametrize("embodiment_tag,artifact", _ARTIFACTS, ids=[t for t, _ in _ARTIFACTS])
+def test_groot_get_action_parity(embodiment_tag, artifact, lerobot_model):
+    """Raw model.get_action(action_pred) parity per embodiment: original vs LeRobot."""
+    original_action, flat_inputs = _load_artifact(artifact)
+    model_inputs = _unflatten(flat_inputs)
 
-    print("\n[Original] Running inference...")
-    original_policy.model.eval()
-    observation = convert_lerobot_to_original_format(batch, modality_config)
-    original_obs_transformed = modality_transform(deepcopy(observation))
+    # Align the flow-matching RNG exactly as the producer did (seed right before sampling).
+    torch.manual_seed(SEED)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(SEED)
+    with torch.inference_mode():
+        out = lerobot_model.get_action(model_inputs)
+    lerobot_action = out["action_pred"].float().cpu()
 
-    # Important: Reset seed immediately before inference to ensure identical RNG state
-    torch.manual_seed(42)
+    t = min(original_action.shape[1], lerobot_action.shape[1])
+    d = min(original_action.shape[2], lerobot_action.shape[2])
+    original_action = original_action[:, :t, :d]
+    lerobot_action = lerobot_action[:, :t, :d]
 
-    with torch.no_grad():
-        original_model_output = original_policy.model.get_action(original_obs_transformed)
-        original_actions_raw = original_model_output["action_pred"]  # [2, 16, 32]
-    # Take first timestep
-    original_actions = original_actions_raw[:, 0, :].to(lerobot_actions.device).to(lerobot_actions.dtype)
-
-    print("Action Comparison:")
-    diff = lerobot_actions - original_actions
-    abs_diff = torch.abs(diff)
-
-    for batch_idx in range(lerobot_actions.shape[0]):
-        print(f"\n{'=' * 60}")
-        print(f"Batch {batch_idx}")
-        print(f"{'=' * 60}")
-        print(f"{'Idx':<5} {'LeRobot':<14} {'Original':<14} {'Difference':<14}")
-        print("-" * 60)
-        for action_idx in range(lerobot_actions.shape[1]):
-            lr_val = lerobot_actions[batch_idx, action_idx].item()
-            orig_val = original_actions[batch_idx, action_idx].item()
-            diff_val = abs(lr_val - orig_val)
-            sign = "+" if (lr_val - orig_val) > 0 else "-"
-            print(f"{action_idx:<5} {lr_val:>13.6f} {orig_val:>13.6f} {sign}{diff_val:>12.6f}")
-
-    max_diff = abs_diff.max().item()
-    tolerance = 0.001
-    assert torch.allclose(lerobot_actions, original_actions, atol=tolerance), (
-        f"Actions differ by more than tolerance ({tolerance}): max diff = {max_diff:.6f}"
+    diff = torch.abs(lerobot_action - original_action)
+    max_diff = diff.max().item()
+    print(
+        f"\n[{embodiment_tag}] shapes lerobot={tuple(lerobot_action.shape)} "
+        f"original={tuple(original_action.shape)}  "
+        f"max|diff|={max_diff:.6e}  mean|diff|={diff.mean().item():.6e}"
     )
-    print(f"\nSuccess: Actions match within tolerance ({tolerance})!")
 
-    del lerobot_policy, lerobot_preprocessor, lerobot_postprocessor
-    del original_policy, modality_config, modality_transform
-    del batch, batch_lerobot, observation
-    cleanup_memory()
-
-
-def test_groot_forward_pass_comparison():
-    """Test forward pass comparison between LeRobot and Original Groot implementations."""
-    print("Test: Forward Pass Comparison (Training Mode)")
-
-    set_seed_all(42)
-
-    lerobot_policy, lerobot_preprocessor, lerobot_postprocessor = instantiate_lerobot_groot(
-        from_pretrained=True
+    assert torch.allclose(lerobot_action, original_action, atol=ATOL, rtol=RTOL), (
+        f"GR00T N1.7 raw action_pred differs for embodiment '{embodiment_tag}' beyond "
+        f"atol={ATOL}, rtol={RTOL}: max|diff|={max_diff:.6e}"
     )
-    original_policy, modality_config, modality_transform = instantiate_original_groot(from_pretrained=True)
-
-    batch = create_dummy_data()
-    lerobot_policy.eval()
-    original_policy.model.eval()
-
-    print("\n[LeRobot] Running forward pass...")
-    batch_lerobot = deepcopy(batch)
-    batch_lerobot_processed = lerobot_preprocessor(batch_lerobot)
-
-    set_seed_all(42)
-    with torch.no_grad():
-        lerobot_loss, lerobot_metrics = lerobot_policy.forward(batch_lerobot_processed)
-
-    print(f"  Loss: {lerobot_loss.item():.6f}")
-
-    print("\n[Original] Running forward pass...")
-    observation = convert_lerobot_to_original_format(batch, modality_config)
-    transformed_obs = modality_transform(observation)
-
-    if "action" not in transformed_obs:
-        action_for_forward = batch_lerobot_processed["action"]
-        action_mask_for_forward = batch_lerobot_processed["action_mask"]
-
-        # Match action horizon if needed
-        if action_for_forward.shape[1] != original_policy.model.action_horizon:
-            if action_for_forward.shape[1] < original_policy.model.action_horizon:
-                pad_size = original_policy.model.action_horizon - action_for_forward.shape[1]
-                last_action = action_for_forward[:, -1:, :]
-                padding = last_action.repeat(1, pad_size, 1)
-                action_for_forward = torch.cat([action_for_forward, padding], dim=1)
-
-                mask_padding = torch.zeros(
-                    action_mask_for_forward.shape[0],
-                    pad_size,
-                    action_mask_for_forward.shape[2],
-                    dtype=action_mask_for_forward.dtype,
-                    device=action_mask_for_forward.device,
-                )
-                action_mask_for_forward = torch.cat([action_mask_for_forward, mask_padding], dim=1)
-            else:
-                action_for_forward = action_for_forward[:, : original_policy.model.action_horizon, :]
-                action_mask_for_forward = action_mask_for_forward[
-                    :, : original_policy.model.action_horizon, :
-                ]
-
-        transformed_obs["action"] = action_for_forward
-        transformed_obs["action_mask"] = action_mask_for_forward
-
-    set_seed_all(42)
-    with torch.no_grad():
-        original_outputs = original_policy.model.forward(transformed_obs)
-
-    original_loss = original_outputs["loss"]
-    print(f"  Loss: {original_loss.item():.6f}")
-
-    loss_diff = abs(lerobot_loss.item() - original_loss.item())
-    loss_rel_diff = loss_diff / (abs(original_loss.item()) + 1e-8) * 100
-
-    print("\nLoss Values:")
-    print(f"  LeRobot: {lerobot_loss.item():.6f}")
-    print(f"  Original: {original_loss.item():.6f}")
-    print(f"  Absolute difference: {loss_diff:.6f}")
-    print(f"  Relative difference: {loss_rel_diff:.2f}%")
-
-    del lerobot_policy, lerobot_preprocessor, lerobot_postprocessor
-    del original_policy, modality_config, modality_transform
-    del batch, batch_lerobot, observation, transformed_obs
-    cleanup_memory()

tests/policies/groot/utils/dump_original_n1_7.py

@@ -0,0 +1,212 @@
+#!/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.
+"""Producer (run in the ORIGINAL gr00t env): dump original GR00T N1.7 outputs + inputs.
+
+The original NVIDIA ``gr00t`` package pins ``transformers==4.57.3`` (py3.10) and its
+model-config dataclasses are incompatible with the ``transformers==5.x`` that the
+LeRobot GR00T N1.7 integration requires. The two implementations therefore cannot be
+imported in the same Python process. To keep the parity comparison FAIR, we run the
+original model in its native env here and serialize, PER EMBODIMENT TAG:
+
+  * the exact pre-processed/collated model inputs (so the LeRobot side consumes the
+    byte-identical tensors -- same image preprocessing, tokenization, normalization),
+  * the random seed used right before the flow-matching sampler,
+  * the raw ``action_pred`` tensor returned by ``model.get_action`` (normalized space,
+    before any per-implementation action decoding).
+
+Inputs are built GENERICALLY from the checkpoint metadata (no per-tag hardcoding):
+state keys + dims come from ``statistics.json``; video + language keys come from the
+processor's per-embodiment modality configs. This lets us test many embodiment tags
+from the SAME checkpoint and confirm the LeRobot integration is not overfit to
+``libero_sim``.
+
+The companion pytest (run in the LeRobot env) loads each .npz, replays the identical
+inputs + seed through the LeRobot GR00T N1.7 model, and asserts the outputs match.
+
+Usage:
+    .venv-original/bin/python tests/policies/groot/utils/dump_original_n1_7.py \
+        --ckpt <path-to-GR00T-N1.7-LIBERO/libero_10> \
+        --out-dir tests/policies/groot/artifacts \
+        [--tags libero_sim,oxe_droid_relative_eef_relative_joint,...] \
+        [--device cuda] [--seed 42]
+
+If --tags is omitted, every embodiment present in the checkpoint statistics is dumped.
+"""
+
+import argparse
+import json
+import os
+from pathlib import Path
+
+import numpy as np
+import torch
+
+IMAGE_SIZE = 256
+BATCH_SIZE = 2
+PROMPT = "pick up the black bowl and place it on the plate"
+
+
+def load_statistics(ckpt: str) -> dict:
+    with open(os.path.join(ckpt, "statistics.json")) as f:
+        return json.load(f)
+
+
+def make_observation(seed: int, video_keys, lang_key, state_spec):
+    """Build a dummy observation dict generically from the embodiment metadata."""
+    rng = np.random.default_rng(seed)
+    video = {
+        k: rng.integers(0, 256, (BATCH_SIZE, 1, IMAGE_SIZE, IMAGE_SIZE, 3), dtype=np.uint8)
+        for k in video_keys
+    }
+    # One ndarray per state key, shape (B, T=1, key_dim); dim taken from statistics.
+    # Keys with dim 0 (e.g. disabled eef on some embodiments) are still emitted as
+    # present-but-empty so the processor's state transform finds every expected key.
+    state = {k: rng.standard_normal((BATCH_SIZE, 1, dim)).astype(np.float32) for k, dim in state_spec}
+    language = {lang_key: [[PROMPT] for _ in range(BATCH_SIZE)]}
+    return {"video": video, "state": state, "language": language}
+
+
+def dump_one_tag(policy, fair_model, tag, modality_cfg, state_spec, args, out_path):
+    from gr00t.data.types import MessageType
+
+    video_keys = modality_cfg["video"].modality_keys
+    lang_key = modality_cfg["language"].modality_keys[0]
+    observation = make_observation(args.seed, video_keys, lang_key, state_spec)
+
+    # Point the policy preprocessing at this embodiment (mirrors Gr00tPolicy.__init__).
+    policy.embodiment_tag = type(policy.embodiment_tag)(tag)
+    policy.modality_configs = {
+        k: v for k, v in policy.processor.get_modality_configs()[tag].items() if k != "rl_info"
+    }
+    policy.language_key = policy.modality_configs["language"].modality_keys[0]
+
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    unbatched = policy._unbatch_observation(observation)
+    processed = []
+    for obs in unbatched:
+        vla = policy._to_vla_step_data(obs)
+        processed.append(policy.processor([{"type": MessageType.EPISODE_STEP.value, "content": vla}]))
+    collated = policy.collate_fn(processed)
+
+    def to_dev(x):
+        if isinstance(x, torch.Tensor) and torch.is_floating_point(x):
+            return x.to(args.device, torch.float32)
+        if isinstance(x, torch.Tensor):
+            return x.to(args.device)
+        if isinstance(x, dict):
+            return {k: to_dev(v) for k, v in x.items()}
+        return x
+
+    collated = {k: to_dev(v) for k, v in collated.items()}
+
+    torch.manual_seed(args.seed)
+    with torch.inference_mode():
+        out = fair_model.get_action(**collated)
+    action_pred = out["action_pred"].float().cpu().numpy()
+
+    flat, meta = {}, {}
+
+    def flatten(prefix, obj):
+        if isinstance(obj, torch.Tensor):
+            arr = obj.float().cpu().numpy() if torch.is_floating_point(obj) else obj.cpu().numpy()
+            flat[f"in::{prefix}"] = arr
+            meta[f"in::{prefix}"] = str(obj.dtype)
+        elif isinstance(obj, dict):
+            for k, v in obj.items():
+                flatten(f"{prefix}.{k}" if prefix else k, v)
+        elif isinstance(obj, (list, tuple)):
+            flat[f"in::{prefix}"] = np.array(obj, dtype=object)
+        else:
+            flat[f"in::{prefix}"] = np.array(obj)
+
+    flatten("", collated)
+
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    np.savez(
+        out_path,
+        action_pred=action_pred,
+        seed=np.array(args.seed),
+        device=np.array(args.device),
+        embodiment_tag=np.array(tag),
+        meta_keys=np.array(list(meta.keys()), dtype=object),
+        meta_dtypes=np.array(list(meta.values()), dtype=object),
+        **flat,
+    )
+    print(f"[{tag}] action_pred {action_pred.shape} -> {out_path.name} ({os.path.getsize(out_path)} B)")
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--out-dir", required=True, help="directory for per-tag .npz files")
+    ap.add_argument("--tags", default="", help="comma-separated embodiment tags (default: all in stats)")
+    ap.add_argument("--device", default="cuda")
+    ap.add_argument("--seed", type=int, default=42)
+    args = ap.parse_args()
+
+    from gr00t.policy.gr00t_policy import Gr00tPolicy
+    from transformers import AutoConfig, AutoModel
+
+    stats = load_statistics(args.ckpt)
+    requested = [t.strip() for t in args.tags.split(",") if t.strip()] or list(stats.keys())
+
+    # Load the policy once (for its processor/preprocessing) on any valid tag.
+    bootstrap_tag = "libero_sim" if "libero_sim" in stats else requested[0]
+    policy = Gr00tPolicy(embodiment_tag=bootstrap_tag, model_path=args.ckpt, device=args.device)
+    all_modality = policy.processor.get_modality_configs()
+
+    # Load a FAIR model (SDPA + fp32) once and reuse across tags. Otherwise the
+    # original checkpoint default (flash_attention_2 + bf16) introduces kernel/rounding
+    # noise vs the LeRobot env (which has no flash_attn and runs SDPA).
+    cfg = AutoConfig.from_pretrained(args.ckpt, trust_remote_code=True)
+    cfg.use_flash_attention = False
+    cfg.load_bf16 = False
+    fair_model = AutoModel.from_pretrained(args.ckpt, config=cfg, trust_remote_code=True)
+    fair_model.to(device=args.device, dtype=torch.float32)
+    fair_model.eval()
+
+    out_dir = Path(args.out_dir)
+    done, skipped = [], []
+    for tag in requested:
+        if tag not in stats or tag not in all_modality:
+            print(f"[skip] {tag}: not present in checkpoint statistics/modality configs")
+            skipped.append(tag)
+            continue
+        state_spec = [(k, len(v["min"])) for k, v in stats[tag]["state"].items()]
+        try:
+            dump_one_tag(
+                policy,
+                fair_model,
+                tag,
+                all_modality[tag],
+                state_spec,
+                args,
+                out_dir / f"original_n1_7_{tag}.npz",
+            )
+            done.append(tag)
+        except Exception as exc:  # noqa: BLE001
+            print(f"[fail] {tag}: {type(exc).__name__}: {exc}")
+            skipped.append(tag)
+
+    print(f"\nDumped {len(done)} tags: {done}")
+    if skipped:
+        print(f"Skipped/failed {len(skipped)} tags: {skipped}")
+
+
+if __name__ == "__main__":
+    main()

tests/policies/lingbot_va/test_configuration_lingbot_va.py

@@ -0,0 +1,78 @@
+#!/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.
+
+from __future__ import annotations
+
+import pytest
+
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.policies.lingbot_va.configuration_lingbot_va import LingBotVAConfig
+from lerobot.utils.constants import ACTION, OBS_IMAGES
+
+
+def make_config(**overrides) -> LingBotVAConfig:
+    kwargs = {"device": "cpu"}
+    kwargs.update(overrides)
+    return LingBotVAConfig(**kwargs)
+
+
+def test_registered_in_choice_registry() -> None:
+    assert "lingbot_va" in PreTrainedConfig.get_known_choices()
+    assert PreTrainedConfig.get_choice_class("lingbot_va") is LingBotVAConfig
+
+
+def test_type_property() -> None:
+    assert make_config().type == "lingbot_va"
+
+
+def test_chunk_size_and_action_steps() -> None:
+    cfg = make_config(frame_chunk_size=4, action_per_frame=4)
+    assert cfg.chunk_size == 16
+    assert cfg.n_action_steps == 16
+    assert cfg.action_delta_indices == list(range(16))
+    assert cfg.observation_delta_indices == list(range(16))
+    assert cfg.reward_delta_indices is None
+
+
+def test_optimizer_and_scheduler_presets() -> None:
+    cfg = make_config()
+    opt = cfg.get_optimizer_preset()
+    assert opt.lr == cfg.optimizer_lr
+    sched = cfg.get_scheduler_preset()
+    assert sched.num_warmup_steps == cfg.scheduler_warmup_steps
+
+
+def test_validate_features_sets_action_feature() -> None:
+    cfg = make_config()
+    cfg.input_features = {f"{OBS_IMAGES}.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128))}
+    cfg.output_features = {}
+    cfg.validate_features()
+    assert ACTION in cfg.output_features
+    assert cfg.output_features[ACTION].shape == (len(cfg.used_action_channel_ids),)
+
+
+def test_validate_features_no_visual_raises() -> None:
+    cfg = make_config()
+    cfg.input_features = {}
+    cfg.output_features = {}
+    with pytest.raises(ValueError, match="at least one visual input feature"):
+        cfg.validate_features()
+
+
+def test_invalid_attn_mode_raises() -> None:
+    with pytest.raises(ValueError, match="attn_mode"):
+        make_config(attn_mode="banana")

tests/policies/lingbot_va/test_factory.py

@@ -0,0 +1,38 @@
+#!/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.
+
+from __future__ import annotations
+
+import pytest
+
+from lerobot.policies.factory import make_policy_config
+from lerobot.policies.lingbot_va.configuration_lingbot_va import LingBotVAConfig
+
+
+def test_make_policy_config_returns_lingbot_va() -> None:
+    cfg = make_policy_config("lingbot_va", device="cpu")
+    assert isinstance(cfg, LingBotVAConfig)
+
+
+def test_get_policy_class_resolves_lazily() -> None:
+    # Importing the policy class pulls in diffusers (Wan2.2 stack); skip if unavailable.
+    pytest.importorskip("diffusers")
+    pytest.importorskip("transformers")
+    from lerobot.policies.factory import get_policy_class
+
+    cls = get_policy_class("lingbot_va")
+    assert cls.name == "lingbot_va"
+    assert cls.config_class is LingBotVAConfig

tests/policies/lingbot_va/test_modules.py

@@ -0,0 +1,128 @@
+#!/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.
+
+"""Unit tests for the vendored LingBot-VA helper code (scheduler + grid utilities)."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+pytest.importorskip("diffusers")  # the model code lives in modeling_lingbot_va, which imports diffusers
+
+from lerobot.policies.lingbot_va.modeling_lingbot_va import FlowMatchScheduler
+from lerobot.policies.lingbot_va.utils import data_seq_to_patch, get_mesh_id
+
+
+def test_flow_match_scheduler_timesteps_monotone_decreasing() -> None:
+    sch = FlowMatchScheduler(shift=5.0, sigma_min=0.0, extra_one_step=True)
+    sch.set_timesteps(20)
+    assert sch.timesteps.shape == (20,)
+    diffs = sch.timesteps[1:] - sch.timesteps[:-1]
+    assert torch.all(diffs <= 0)  # decreasing
+
+
+def test_flow_match_scheduler_step_preserves_shape() -> None:
+    sch = FlowMatchScheduler(shift=5.0, sigma_min=0.0, extra_one_step=True)
+    sch.set_timesteps(20)
+    sample = torch.zeros(1, 48, 4, 8, 16)
+    out = sch.step(torch.ones_like(sample), sch.timesteps[0], sample)
+    assert out.shape == sample.shape
+
+
+def test_flow_match_scheduler_add_noise() -> None:
+    sch = FlowMatchScheduler(shift=5.0, sigma_min=0.0, extra_one_step=True)
+    sch.set_timesteps(20)
+    sample = torch.randn(1, 48, 4, 8, 16)
+    noise = torch.randn_like(sample)
+    noisy = sch.add_noise(sample, noise, sch.timesteps[:4], t_dim=2)
+    assert noisy.shape == sample.shape
+
+
+def test_get_mesh_id_latent_shape() -> None:
+    grid = get_mesh_id(4, 8, 16, 0, 1, 0)
+    assert grid.shape == (4, 4 * 8 * 16)  # (f, h, w, stream) x tokens
+
+
+def test_get_mesh_id_action_shape() -> None:
+    grid = get_mesh_id(4, 4, 1, 1, 1, 0, action=True)
+    assert grid.shape == (4, 4 * 4 * 1)
+    # Action rows for h/w are sentinel -1.
+    assert torch.all(grid[1] < 0)
+    assert torch.all(grid[2] < 0)
+
+
+def test_data_seq_to_patch_roundtrip_shape() -> None:
+    b, f, h, w, c = 1, 4, 8, 16, 48
+    seq = torch.arange(b * f * h * w * c, dtype=torch.float32).reshape(b, f * h * w, c)
+    out = data_seq_to_patch((1, 2, 2), seq, f, h, w, batch_size=b)
+    assert out.shape == (b, c, f, h, w)
+
+
+def test_training_step_reduces_loss_tiny_flex() -> None:
+    """End-to-end single training step (flow-matching loss -> backward -> AdamW) on a tiny config.
+
+    Exercises the flex-attention training path; requires a CUDA GPU with flex-attention support.
+    """
+    if not torch.cuda.is_available():
+        import pytest
+
+        pytest.skip("training step test requires a CUDA GPU (flex-attention)")
+
+    from lerobot.configs.types import FeatureType, PolicyFeature
+    from lerobot.policies.lingbot_va.configuration_lingbot_va import LingBotVAConfig
+    from lerobot.policies.lingbot_va.modeling_lingbot_va import LingBotVAPolicy
+    from lerobot.utils.constants import ACTION, OBS_IMAGES
+
+    cfg = LingBotVAConfig(
+        attn_mode="flex",
+        dtype="bfloat16",
+        in_channels=16,
+        out_channels=16,
+        action_dim=8,
+        text_dim=32,
+        freq_dim=64,
+        ffn_dim=64,
+        num_attention_heads=2,
+        attention_head_dim=24,
+        num_layers=2,
+        frame_chunk_size=2,
+        action_per_frame=4,
+        used_action_channel_ids=[0, 1, 2, 3],
+        obs_cam_keys=[f"{OBS_IMAGES}.image"],
+        device="cuda",
+    )
+    cfg.input_features = {f"{OBS_IMAGES}.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 64, 64))}
+    cfg.output_features = {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(4,))}
+    cfg.validate_features()
+
+    policy = LingBotVAPolicy(cfg).to("cuda")
+    policy.train()
+    opt = torch.optim.AdamW(policy.get_optim_params(), lr=1e-4)
+
+    b, fc, apf = 1, cfg.frame_chunk_size, cfg.action_per_frame
+    latents = torch.randn(b, cfg.in_channels, fc, 4, 4, device="cuda", dtype=torch.bfloat16)
+    actions = torch.randn(b, cfg.action_dim, fc, apf, 1, device="cuda", dtype=torch.bfloat16)
+    amask = torch.zeros(cfg.action_dim, device="cuda")
+    amask[cfg.used_action_channel_ids] = 1.0
+    actions_mask = amask.view(1, -1, 1, 1, 1).expand_as(actions)
+    text_emb = torch.randn(b, cfg.max_sequence_length, cfg.text_dim, device="cuda", dtype=torch.bfloat16)
+
+    loss, metrics = policy.training_loss_from_streams(latents, actions, actions_mask, text_emb)
+    assert torch.isfinite(loss) and {"latent_loss", "action_loss"} <= set(metrics)
+    loss.backward()
+    assert any(p.grad is not None and torch.isfinite(p.grad).all() for p in policy.get_optim_params())
+    opt.step()

tests/policies/lingbot_va/test_processor.py

@@ -0,0 +1,88 @@
+#!/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.
+
+from __future__ import annotations
+
+import torch
+
+from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.policies.lingbot_va.configuration_lingbot_va import LingBotVAConfig
+from lerobot.policies.lingbot_va.processor_lingbot_va import make_lingbot_va_pre_post_processors
+from lerobot.processor import PolicyProcessorPipeline, UnnormalizerProcessorStep
+from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
+from lerobot.utils.constants import (
+    ACTION,
+    OBS_IMAGES,
+    POLICY_POSTPROCESSOR_DEFAULT_NAME,
+    POLICY_PREPROCESSOR_DEFAULT_NAME,
+)
+
+
+def _make_config() -> LingBotVAConfig:
+    cfg = LingBotVAConfig(device="cpu")
+    cfg.input_features = {f"{OBS_IMAGES}.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128))}
+    cfg.output_features = {}
+    cfg.validate_features()
+    return cfg
+
+
+def test_make_pre_post_processors_names_and_steps() -> None:
+    cfg = _make_config()
+    pre, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats=None)
+    assert pre.name == POLICY_PREPROCESSOR_DEFAULT_NAME
+    assert post.name == POLICY_POSTPROCESSOR_DEFAULT_NAME
+    # Actions are unnormalized by the standard built-in quantile unnormalizer.
+    assert any(isinstance(s, UnnormalizerProcessorStep) for s in post.steps)
+
+
+def test_freshly_built_postprocessor_is_identity() -> None:
+    # Without action stats the quantile unnormalizer is a no-op (identity passthrough): the real
+    # per-benchmark q01/q99 are restored from the saved checkpoint on load, not hardcoded here.
+    cfg = _make_config()
+    _, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats=None)
+    normed = torch.tensor([[0.3, -0.5, 1.0, -1.0, 0.0, 0.7, -0.2]])
+    assert torch.allclose(post(normed), normed, atol=1e-6)
+
+
+def test_postprocessor_quantile_unnormalization() -> None:
+    # QUANTILES unnormalize maps [-1, 1] -> [q01, q99]: -1 -> q01, +1 -> q99.
+    cfg = _make_config()
+    q01 = [-1.0, -0.5, 0.0, -1.0, -1.0, -1.0, -1.0]
+    q99 = [1.0, 0.5, 2.0, 1.0, 1.0, 1.0, 1.0]
+    stats = {ACTION: {"q01": q01, "q99": q99}}
+    _, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats=stats)
+    out_lo = post(torch.full((1, 7), -1.0))
+    out_hi = post(torch.full((1, 7), 1.0))
+    assert torch.allclose(out_lo, torch.tensor(q01).unsqueeze(0), atol=1e-4)
+    assert torch.allclose(out_hi, torch.tensor(q99).unsqueeze(0), atol=1e-4)
+
+
+def test_postprocessor_stats_survive_save_load(tmp_path) -> None:
+    # Regression guard for the Hub mechanism: the q01/q99 stats live in the saved post-processor
+    # state and must round-trip through save_pretrained / from_pretrained.
+    cfg = _make_config()
+    q01 = [-0.6, -0.8, -0.9, -0.1, -0.15, -0.25, -1.0]
+    q99 = [0.9, 0.85, 0.9, 0.17, 0.18, 0.34, 1.0]
+    _, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats={ACTION: {"q01": q01, "q99": q99}})
+    post.save_pretrained(tmp_path)
+    loaded = PolicyProcessorPipeline.from_pretrained(
+        tmp_path,
+        config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json",
+        to_transition=policy_action_to_transition,
+        to_output=transition_to_policy_action,
+    )
+    out = loaded(torch.full((1, 7), -1.0))
+    assert torch.allclose(out, torch.tensor(q01).unsqueeze(0), atol=1e-4)

tests/policies/molmoact2/test_molmoact2.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-
 # Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -35,24 +33,27 @@ pytest.importorskip("scipy")
 from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.policies import get_policy_class, make_policy_config
 from lerobot.policies.molmoact2 import (
-    configuration_molmoact2 as molmoact2_config,
     modeling_molmoact2 as molmoact2_modeling,
     processor_molmoact2 as molmoact2_processor,
 )
-from lerobot.policies.molmoact2.configuration_molmoact2 import (
-    MolmoAct2Config,
-    MolmoAct2CosineDecayWithWarmupSchedulerConfig,
-    infer_molmoact2_max_sequence_length,
+from lerobot.policies.molmoact2.configuration_molmoact2 import MolmoAct2Config
+from lerobot.policies.molmoact2.modeling_molmoact2 import (
+    MolmoAct2Policy,
+    _apply_action_chunk_padding_mask,
+    _apply_action_dim_padding_mask,
+    _combine_rollout_seeds,
 )
-from lerobot.policies.molmoact2.modeling_molmoact2 import MolmoAct2Policy
 from lerobot.policies.molmoact2.processor_molmoact2 import (
+    MolmoAct2ActionFrameTransformStep,
     MolmoAct2ClampNormalizedProcessorStep,
     MolmoAct2MaskedNormalizerProcessorStep,
     MolmoAct2MaskedUnnormalizerProcessorStep,
     MolmoAct2PackInputsProcessorStep,
+    MolmoAct2StateFrameTransformStep,
     _add_gripper_masks_to_stats,
     _build_discrete_state_string,
     _normalize_question_text,
+    infer_molmoact2_max_sequence_length,
     make_molmoact2_pre_post_processors,
 )
 from lerobot.policies.rtc.configuration_rtc import RTCConfig
@@ -71,34 +72,38 @@ def test_molmoact2_policy_registration():
     assert cfg.per_episode_seed is False
     assert cfg.eval_seed is None
     assert cfg.normalize_language is True
-    assert cfg.get_scheduler_preset().num_decay_steps is None
+    assert cfg.get_scheduler_preset().num_decay_steps == 100_000
     assert cfg.action_delta_indices == list(range(cfg.chunk_size))
     assert get_policy_class("molmoact2") is MolmoAct2Policy
 
 
 def test_molmoact2_checkpoint_download_ignores_remote_python(monkeypatch):
+    import huggingface_hub
+
     download_kwargs = {}
 
     def fake_snapshot_download(**kwargs):
         download_kwargs.update(kwargs)
         return "/tmp/downloaded-molmoact2"
 
-    monkeypatch.setattr(molmoact2_config, "snapshot_download", fake_snapshot_download)
+    monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
 
-    checkpoint_location = molmoact2_config._resolve_checkpoint_location("allenai/MolmoAct2")
+    checkpoint_location = molmoact2_modeling._resolve_checkpoint_location("allenai/MolmoAct2")
 
     assert checkpoint_location == "/tmp/downloaded-molmoact2"
     assert download_kwargs["ignore_patterns"] == ["*.py", "*.pyc", "__pycache__/*"]
 
 
-def test_molmoact2_scheduler_decay_steps_auto_match_training_steps():
+def test_molmoact2_scheduler_auto_scales_to_training_steps():
+    from lerobot.optim import CosineDecayWithWarmupSchedulerConfig
+
     param = torch.nn.Parameter(torch.ones(()))
     optimizer = torch.optim.AdamW([param], lr=0.001)
-    config = MolmoAct2CosineDecayWithWarmupSchedulerConfig(
+    config = CosineDecayWithWarmupSchedulerConfig(
         peak_lr=0.01,
         decay_lr=0.001,
         num_warmup_steps=10,
-        num_decay_steps=None,
+        num_decay_steps=100_000,
     )
 
     scheduler = config.build(optimizer, num_training_steps=100)
@@ -123,9 +128,7 @@ def test_molmoact2_rollout_generator_uses_eval_seed_per_task():
         batch_size=3,
         device=torch.device("cpu"),
     )
-    expected_first = torch.Generator().manual_seed(
-        MolmoAct2Policy._combine_rollout_seeds(first_seed=1000, batch_size=3)
-    )
+    expected_first = torch.Generator().manual_seed(_combine_rollout_seeds(first_seed=1000, batch_size=3))
     assert torch.allclose(torch.rand(4, generator=first), torch.rand(4, generator=expected_first))
 
     policy.reset()
@@ -134,9 +137,7 @@ def test_molmoact2_rollout_generator_uses_eval_seed_per_task():
         batch_size=3,
         device=torch.device("cpu"),
     )
-    expected_second = torch.Generator().manual_seed(
-        MolmoAct2Policy._combine_rollout_seeds(first_seed=1003, batch_size=3)
-    )
+    expected_second = torch.Generator().manual_seed(_combine_rollout_seeds(first_seed=1003, batch_size=3))
     assert torch.allclose(torch.rand(4, generator=second), torch.rand(4, generator=expected_second))
 
     policy.reset()
@@ -145,9 +146,7 @@ def test_molmoact2_rollout_generator_uses_eval_seed_per_task():
         batch_size=3,
         device=torch.device("cpu"),
     )
-    expected_new_task = torch.Generator().manual_seed(
-        MolmoAct2Policy._combine_rollout_seeds(first_seed=1000, batch_size=3)
-    )
+    expected_new_task = torch.Generator().manual_seed(_combine_rollout_seeds(first_seed=1000, batch_size=3))
     assert torch.allclose(torch.rand(4, generator=new_task), torch.rand(4, generator=expected_new_task))
 
 
@@ -537,36 +536,26 @@ def test_train_action_expert_only_requires_continuous_action_mode():
 
 
 def test_molmoact2_sequence_length_is_inferred_from_fixed_token_budget():
-    cfg = MolmoAct2Config(
-        action_mode="both",
-        chunk_size=10,
-        n_action_steps=10,
-        image_keys=["observation.images.image", "observation.images.wrist_image"],
-        input_features={OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(8,))},
-        output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(7,))},
-    )
-
-    assert cfg.max_sequence_length is None
-    assert cfg.inferred_max_sequence_length() == 640
-    assert cfg.inferred_max_sequence_length(include_discrete_action=False) == 576
     assert (
         infer_molmoact2_max_sequence_length(
-            num_images=2,
-            state_dim=8,
-            action_dim=7,
-            action_horizon=30,
-            include_discrete_action=True,
+            num_images=2, state_dim=8, action_dim=7, action_horizon=10, include_discrete_action=True
+        )
+        == 640
+    )
+    assert (
+        infer_molmoact2_max_sequence_length(
+            num_images=2, state_dim=8, action_dim=7, action_horizon=10, include_discrete_action=False
+        )
+        == 576
+    )
+    assert (
+        infer_molmoact2_max_sequence_length(
+            num_images=2, state_dim=8, action_dim=7, action_horizon=30, include_discrete_action=True
         )
         == 768
     )
 
 
-def test_molmoact2_sequence_length_override_is_preserved():
-    cfg = MolmoAct2Config(max_sequence_length=1024)
-
-    assert cfg.inferred_max_sequence_length(num_images=2, state_dim=8, action_dim=7) == 1024
-
-
 def test_train_action_expert_only_freezes_non_action_expert_params():
     class DummyBackbone(torch.nn.Module):
         def __init__(self):
@@ -939,6 +928,39 @@ def test_question_normalization_matches_release_prompt_style():
     )
 
 
+def test_joint_frame_transform_round_trip():
+    signs = [1.0, -1.0, 1.0, 1.0, 1.0, 1.0]
+    offsets = [0.0, 90.0, 90.0, 0.0, 0.0, 0.0]
+    original_state = torch.tensor([[10.0, -90.0, -120.0, 30.0, 0.0, -45.0]])
+
+    state_step = MolmoAct2StateFrameTransformStep(joint_signs=signs, joint_offsets=offsets)
+    action_step = MolmoAct2ActionFrameTransformStep(joint_signs=signs, joint_offsets=offsets)
+
+    transition = {
+        TransitionKey.OBSERVATION: {OBS_STATE: original_state.clone()},
+    }
+    transformed = state_step(transition)
+    model_state = transformed[TransitionKey.OBSERVATION][OBS_STATE]
+
+    action_transition = {TransitionKey.ACTION: model_state.clone()}
+    recovered = action_step(action_transition)
+    recovered_state = recovered[TransitionKey.ACTION]
+
+    assert torch.allclose(recovered_state, original_state)
+
+
+def test_joint_frame_transform_noop_when_none():
+    state_step = MolmoAct2StateFrameTransformStep(joint_signs=None, joint_offsets=None)
+    action_step = MolmoAct2ActionFrameTransformStep(joint_signs=None, joint_offsets=None)
+    state = torch.tensor([[10.0, -90.0, -120.0]])
+
+    state_transition = {TransitionKey.OBSERVATION: {OBS_STATE: state}}
+    assert state_step(state_transition) is state_transition
+
+    action_transition = {TransitionKey.ACTION: state}
+    assert action_step(action_transition) is action_transition
+
+
 def test_action_padding_marks_only_real_dimensions():
     step = object.__new__(MolmoAct2PackInputsProcessorStep)
     step.max_action_dim = 32
@@ -963,7 +985,7 @@ def test_action_dim_padding_loss_reduces_like_old_trainer():
         ]
     )
 
-    reduced = MolmoAct2Policy._apply_action_dim_padding_mask(loss, action_dim_is_pad)
+    reduced = _apply_action_dim_padding_mask(loss, action_dim_is_pad)
 
     expected = torch.stack(
         [
@@ -979,7 +1001,7 @@ def test_action_chunk_padding_keeps_old_mean_denominator():
     loss = torch.ones(1, 2, 4, 3)
     action_horizon_is_pad = torch.tensor([[False, False, True, True]])
 
-    masked = MolmoAct2Policy._apply_action_chunk_padding_mask(loss, action_horizon_is_pad)
+    masked = _apply_action_chunk_padding_mask(loss, action_horizon_is_pad)
 
     assert masked.mean().item() == 0.5
 

tests/policies/vla_jepa/conftest.py

@@ -8,7 +8,6 @@ from types import SimpleNamespace
 import numpy as np
 import pytest
 import torch
-from PIL import Image
 from torch import Tensor, nn
 
 from lerobot.configs.types import FeatureType, PolicyFeature
@@ -191,7 +190,7 @@ class _FakeQwenInterface(nn.Module):
 
     def build_inputs(
         self,
-        images: list[list[Image.Image]],
+        images: list[list[Tensor]],
         instructions: list[str],
         action_prompt: str,
         embodied_prompt: str,
@@ -214,12 +213,13 @@ class _FakeQwenInterface(nn.Module):
         }
 
     @staticmethod
-    def tensor_to_pil(image_tensor: Tensor) -> Image.Image:
-        image = image_tensor.detach().cpu()
-        if image.ndim == 3 and image.shape[0] in (1, 3):
-            image = image.permute(1, 2, 0)
-        image = (image.float().clamp(0, 1) * 255).to(torch.uint8).numpy()
-        return Image.fromarray(image)
+    def to_pixel_values(image_tensor: Tensor) -> Tensor:
+        image = image_tensor.detach().float()
+        if image.shape[-3] == 1:
+            repeats = [1] * image.ndim
+            repeats[-3] = 3
+            image = image.repeat(*repeats)
+        return image
 
 
 class _FakeVideoEncoder(nn.Module):
@@ -242,12 +242,14 @@ class _FakeVideoEncoder(nn.Module):
 
 
 class _FakeVideoProcessor:
-    def __call__(self, videos, return_tensors: str) -> dict[str, Tensor]:
+    def __call__(self, videos, return_tensors: str, device=None, **kwargs) -> dict[str, Tensor]:
         assert return_tensors == "pt"
         if isinstance(videos, list):
             pixel_values = torch.stack([torch.as_tensor(v) for v in videos])
         else:
             pixel_values = torch.as_tensor(videos).unsqueeze(0)
+        if device is not None:
+            pixel_values = pixel_values.to(device)
         return {"pixel_values_videos": pixel_values}
 
 

tests/policies/vla_jepa/test_vla_jepa.py

@@ -211,40 +211,42 @@ def test_reset_clears_action_queue(patch_vla_jepa_external_models: None) -> None
 
 
 def test_prepare_model_inputs_training_format(patch_vla_jepa_external_models: None) -> None:
-    from PIL import Image
-
     policy = VLAJEPAPolicy(make_config())
-    examples = policy._prepare_model_inputs(make_train_batch())
+    inputs = policy._prepare_model_inputs(make_train_batch())
 
-    assert len(examples) == BATCH_SIZE
-    for ex in examples:
-        assert set(ex) >= {"image", "video", "lang", "action", "state"}
-        assert len(ex["image"]) == 1 and isinstance(ex["image"][0], Image.Image)
-        assert ex["video"].ndim == 5 and ex["video"].dtype == np.uint8  # [V,T,H,W,C]
-        assert ex["action"].shape == (ACTION_HORIZON, ACTION_DIM)
-        assert ex["state"].shape == (1, STATE_DIM)
+    assert set(inputs) >= {"images", "instructions", "videos", "actions", "state"}
+    # images: per-sample, per-view [C, H, W] float tensors (kept as a list for Qwen messages)
+    assert len(inputs["images"]) == BATCH_SIZE and len(inputs["images"][0]) == 1
+    img = inputs["images"][0][0]
+    assert isinstance(img, torch.Tensor) and img.dtype == torch.float32 and img.ndim == 3
+    assert len(inputs["instructions"]) == BATCH_SIZE
+    # videos: batched [B, V, T, C, H, W] float
+    assert inputs["videos"].ndim == 6 and inputs["videos"].shape[0] == BATCH_SIZE
+    assert inputs["videos"].dtype == torch.float32
+    assert inputs["actions"].shape == (BATCH_SIZE, ACTION_HORIZON, ACTION_DIM)
+    assert inputs["state"].shape == (BATCH_SIZE, 1, STATE_DIM)
 
 
 def test_prepare_model_inputs_inference_omits_action(patch_vla_jepa_external_models: None) -> None:
     policy = VLAJEPAPolicy(make_config())
-    for ex in policy._prepare_model_inputs(make_inference_batch()):
-        assert "action" not in ex
-        assert "image" in ex and "video" in ex and "lang" in ex
+    inputs = policy._prepare_model_inputs(make_inference_batch())
+    assert "actions" not in inputs and "action_is_pad" not in inputs
+    assert {"images", "instructions", "state"} <= set(inputs)
 
 
 def test_prepare_model_inputs_missing_task_uses_default(patch_vla_jepa_external_models: None) -> None:
     policy = VLAJEPAPolicy(make_config())
     batch = make_inference_batch()
     del batch["task"]
-    examples = policy._prepare_model_inputs(batch)
-    assert all(isinstance(ex["lang"], str) and len(ex["lang"]) > 0 for ex in examples)
+    instructions = policy._prepare_model_inputs(batch)["instructions"]
+    assert all(isinstance(s, str) and len(s) > 0 for s in instructions)
 
 
 def test_prepare_model_inputs_string_task_broadcast(patch_vla_jepa_external_models: None) -> None:
     policy = VLAJEPAPolicy(make_config())
     batch = make_inference_batch()
     batch["task"] = "open the drawer"
-    assert all(ex["lang"] == "open the drawer" for ex in policy._prepare_model_inputs(batch))
+    assert policy._prepare_model_inputs(batch)["instructions"] == ["open the drawer"] * BATCH_SIZE
 
 
 def test_prepare_model_inputs_no_state_omitted(patch_vla_jepa_external_models: None) -> None:
@@ -253,7 +255,7 @@ def test_prepare_model_inputs_no_state_omitted(patch_vla_jepa_external_models: N
     policy = VLAJEPAPolicy(make_config())
     batch = make_inference_batch()
     del batch[OBS_STATE]
-    assert all("state" not in ex for ex in policy._prepare_model_inputs(batch))
+    assert "state" not in policy._prepare_model_inputs(batch)
 
 
 # ---------------------------------------------------------------------------
@@ -446,14 +448,14 @@ def test_postprocessor_applied_after_predict_action_chunk(
     """
     from lerobot.policies.vla_jepa.processor_vla_jepa import make_vla_jepa_pre_post_processors
 
-    raw_actions = np.zeros((BATCH_SIZE, ACTION_HORIZON, ACTION_DIM), dtype=np.float32)
+    raw_actions = torch.zeros((BATCH_SIZE, ACTION_HORIZON, ACTION_DIM), dtype=torch.float32)
 
     cfg = make_config()
     cfg.clip_normalized_actions = False
     cfg.binarize_gripper_action = False
     policy = VLAJEPAPolicy(cfg)
     policy.eval()
-    monkeypatch.setattr(policy.model, "predict_action", lambda *a, **kw: raw_actions.copy())
+    monkeypatch.setattr(policy.model, "predict_action", lambda *a, **kw: raw_actions.clone())
 
     dataset_stats = _make_dataset_stats()
     _, postprocessor = make_vla_jepa_pre_post_processors(cfg, dataset_stats)
@@ -564,9 +566,9 @@ def test_single_view_is_duplicated_for_world_model(patch_vla_jepa_external_model
     original_processor = policy.model.video_processor
 
     class _CapturingProcessor:
-        def __call__(self, videos: list, return_tensors: str) -> dict:
+        def __call__(self, videos: list, return_tensors: str, **kwargs) -> dict:
             captured_videos.extend(videos)
-            return original_processor(videos=videos, return_tensors=return_tensors)
+            return original_processor(videos=videos, return_tensors=return_tensors, **kwargs)
 
     policy.model.video_processor = _CapturingProcessor()
     policy.forward(_make_multiview_train_batch(num_views=1))
@@ -587,9 +589,9 @@ def test_excess_views_trimmed_for_world_model(patch_vla_jepa_external_models: No
     original_processor = policy.model.video_processor
 
     class _CapturingProcessor:
-        def __call__(self, videos: list, return_tensors: str) -> dict:
+        def __call__(self, videos: list, return_tensors: str, **kwargs) -> dict:
             captured_videos.extend(videos)
-            return original_processor(videos=videos, return_tensors=return_tensors)
+            return original_processor(videos=videos, return_tensors=return_tensors, **kwargs)
 
     policy.model.video_processor = _CapturingProcessor()
     policy.forward(_make_multiview_train_batch(num_views=3))