add tests and docs for the gripper, clipping and view-count changes

docs/source/vla_jepa.mdx

@@ -30,10 +30,19 @@ Only Qwen + the action head are used. The world model is not needed at inference
 
 Available presets via `action_model_type`:
 
-| Preset  | Hidden dim | Heads | Head dim |
-| ------- | ---------- | ----- | -------- |
-| `DiT-B` | 768        | 12    | 64       |
-| `DiT-L` | 1536       | 32    | 48       |
+| Preset  | Heads | Head dim |
+| ------- | ----- | -------- |
+| `DiT-B` | 12    | 64       |
+| `DiT-L` | 32    | 48       |
+
+The preset only sets the attention geometry, and each entry can be overridden by
+`action_num_heads` / `action_attention_head_dim`. Two widths follow from it:
+
+- the DiT's **internal** width is `heads x head_dim` (768 for `DiT-B`), derived rather than configured;
+- the DiT's **output** width, and the width of the action-decoder and state-encoder MLPs, is
+  `action_hidden_size` (default 1024).
+
+So `DiT-B` runs a 768-wide transformer that projects to 1024. The two are independent.
 
 ### World model details
 
@@ -74,10 +83,27 @@ Key parameters in `VLAJEPAConfig`:
 | `num_inference_timesteps` | 4       | Euler integration steps for action denoising                                                                                                                                    |
 | `freeze_qwen`             | `False` | Freeze the Qwen3-VL backbone and only train the action head                                                                                                                     |
 | `reinit_modules`          | `None`  | Key prefixes allowed to be randomly re-initialised on load (for cross-embodiment transfer, see [Fine-tuning on a different embodiment](#fine-tuning-on-a-different-embodiment)) |
-| `gripper_dim`             | 6       | Index of the gripper dimension in the action vector (e.g. 6 for a 7-DoF arm with gripper as the last joint)                                                                     |
-| `gripper_threshold`       | 0.5     | Threshold used by `pre_snap_gripper_action` and `binarize_gripper_action` to binarize the gripper dimension                                                                     |
-| `pre_snap_gripper_action` | `True`  | Snap the gripper dim to {0, 1} before unnormalization. Set to `False` for robots without a binary gripper                                                                       |
-| `binarize_gripper_action` | `True`  | Binarize the gripper dim to {-1, 1} after unnormalization. Set to `False` for robots without a binary gripper                                                                   |
+| `resize_images_to`        | `None`  | `(height, width)` every camera frame is resized to before the Qwen3-VL vision tower. `None` keeps the native resolution, and Qwen3-VL's patch count grows with it, so a 720x1280 camera can exhaust GPU memory. The published checkpoints use `[224, 224]` |
+| `gripper_dim`             | 6       | Index of the gripper dimension in the action vector. Ignored when `gripper_joint_names` matches a dataset action name                                                            |
+| `gripper_joint_names`     | `["gripper"]` | Action-dimension names identifying the gripper; the matched index wins over `gripper_dim`                                                                                  |
+| `gripper_threshold`       | 0.5     | Threshold used by `pre_snap_gripper_action` and `binarize_gripper_action`. Note `binarize` runs *after* unnormalization, so this is compared against the gripper's physical value |
+| `pre_snap_gripper_action` | `False` | Snap the gripper dim to {0, 1} before unnormalization. LIBERO-specific, see below                                                                                                |
+| `binarize_gripper_action` | `False` | Binarize the gripper dim to {-1, 1} after unnormalization. LIBERO-specific, see below                                                                                            |
+| `clip_normalized_actions` | `True`  | Clip normalized actions to [-1, 1] before unnormalizing. Only applied when `ACTION` uses `MIN_MAX`; ignored (with a warning) under `MEAN_STD`, where it would truncate at 1 sigma |
+| `world_model_num_views`   | `None`  | Camera views the world-model predictor is built for. Baked into checkpoint shapes. `None` falls back to `jepa_tubelet_size`, which is what the published checkpoints encode       |
+
+<Tip warning={true}>
+
+`pre_snap_gripper_action` and `binarize_gripper_action` are a port of the starVLA LIBERO eval
+loop and are only correct for LIBERO's action convention. `pre_snap` writes {0, 1} into
+*normalized* space, the unnormalizer maps those to the midpoint and the max, and `binarize` then
+compares that **physical** value against `gripper_threshold` (0.5). For a gripper measured in
+degrees, mm or [0, 100], both values land above the threshold and the commanded gripper becomes a
+constant. They default to `False` for that reason; enable them only for LIBERO-style setups, and
+set `gripper_threshold` in the gripper's own units if you do. The processor factory warns when the
+dataset stats show the range cannot work.
+
+</Tip>
 
 ---
 
@@ -187,14 +213,20 @@ lerobot-eval \
 
 ## Fine-tuning on datasets with a different number of cameras
 
-The pretrained world model predictor was trained with `embed_dim = jepa_tubelet_size × 1024` (default `jepa_tubelet_size=2`).
+The pretrained world model predictor was trained with `embed_dim = world_model_num_views × 1024`, i.e. two camera views.
+
+<Tip>
+
+This view count used to be read from `jepa_tubelet_size`, which also names the JEPA encoder's *temporal* tubelet size. `world_model_num_views` is the field for it now; leaving it at `None` falls back to `jepa_tubelet_size` so the published checkpoints keep loading unchanged.
+
+</Tip>
 
 **Default behaviour — view padding / trimming (no action required)**
 
-When fine-tuning from `VLA-JEPA-Pretrain` the model automatically adjusts the number of views fed to the world model to match `jepa_tubelet_size`:
+When fine-tuning from `VLA-JEPA-Pretrain` the model automatically adjusts the number of views fed to the world model to match `world_model_num_views`:
 
 - **Single-view datasets (e.g. BridgeV2):** the single-view latent is duplicated to produce a two-view world-model input, preserving the JEPA self-supervised signal without any weight mismatch.
-- **>2-view datasets (e.g. DROID with 3 views):** all views are passed to the Qwen backbone (for richer context), but only the first `jepa_tubelet_size` views (one wrist + one third-person, following the configured view order) are used for the world model.
+- **>2-view datasets (e.g. DROID with 3 views):** all views are passed to the Qwen backbone (for richer context), but only the first `world_model_num_views` views (one wrist + one third-person, following the configured view order) are used for the world model.
 
 **Option 1 — Disable the world model**
 
@@ -210,7 +242,7 @@ lerobot-train \
 
 **Option 2 — Reinitialize the predictor input projection**
 
-If you want to change `jepa_tubelet_size` to a value other than 2, load the checkpoint with `strict=False` and reinitialize `model.video_predictor.predictor_embed` for the new `embed_dim`. All other predictor block weights (attention, MLP, norm, output projection) are camera-count-agnostic and can be reused from the pretrained checkpoint.
+If you want to change `world_model_num_views` to a value other than 2, load the checkpoint with `strict=False` and reinitialize `model.video_predictor.predictor_embed` for the new `embed_dim`. All other predictor block weights (attention, MLP, norm, output projection) are camera-count-agnostic and can be reused from the pretrained checkpoint.
 
 ---
 

tests/policies/vla_jepa/conftest.py

@@ -1,4 +1,19 @@
 #!/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.
+
 """Shared fixtures and helpers for VLA-JEPA tests."""
 
 from __future__ import annotations

tests/policies/vla_jepa/test_action_head.py

@@ -1,5 +1,19 @@
 #!/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

tests/policies/vla_jepa/test_configuration_vla_jepa.py

@@ -1,5 +1,19 @@
 #!/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

tests/policies/vla_jepa/test_processor_vla_jepa.py

@@ -26,6 +26,7 @@ equivalence of the two paths on the Qwen image path.
 
 from __future__ import annotations
 
+import logging
 from copy import deepcopy
 
 import pytest
@@ -43,14 +44,19 @@ from conftest import (  # noqa: E402
     make_train_batch,
 )
 
-from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature  # noqa: E402
+from lerobot.configs.types import (  # noqa: E402
+    FeatureType,
+    NormalizationMode,
+    PipelineFeatureType,
+    PolicyFeature,
+)
 from lerobot.policies.vla_jepa.modeling_vla_jepa import VLAJEPAPolicy  # noqa: E402
 from lerobot.policies.vla_jepa.processor_vla_jepa import (  # noqa: E402
     ImagePrepProcessorStep,
     make_vla_jepa_pre_post_processors,
 )
-from lerobot.processor import ProcessorStepRegistry  # noqa: E402
-from lerobot.utils.constants import OBS_IMAGES, OBS_STATE  # noqa: E402
+from lerobot.processor import PolicyProcessorPipeline, ProcessorStepRegistry  # noqa: E402
+from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE  # noqa: E402
 
 RESIZE = (IMAGE_SIZE // 2, IMAGE_SIZE // 2)  # (4, 4)
 IMG_KEY = f"{OBS_IMAGES}.laptop"
@@ -180,3 +186,131 @@ def test_image_prep_matches_model_qwen_path(patch_vla_jepa_external_models: None
     for views_a, views_b in zip(imgs_a, imgs_b, strict=True):
         for a, b in zip(views_a, views_b, strict=True):
             assert torch.equal(a, b)
+
+
+# ---------------------------------------------------------------------------
+# Gripper post-step serialization and the normalization-mode coupling
+# ---------------------------------------------------------------------------
+
+
+def _stats_for(cfg, gripper_min: float = 0.0, gripper_max: float = 1.0):
+    """Dataset stats matching `cfg`'s features, with a settable gripper range."""
+    action_min = torch.zeros(cfg.action_dim)
+    action_max = torch.ones(cfg.action_dim)
+    gripper = cfg.resolved_gripper_dim
+    if gripper < cfg.action_dim:
+        action_min[gripper] = gripper_min
+        action_max[gripper] = gripper_max
+    stats = {
+        ACTION: {
+            "min": action_min,
+            "max": action_max,
+            "mean": torch.zeros(cfg.action_dim),
+            "std": torch.ones(cfg.action_dim),
+        }
+    }
+    for key, feat in cfg.input_features.items():
+        stats[key] = {
+            "min": torch.zeros(feat.shape),
+            "max": torch.ones(feat.shape),
+            "mean": torch.zeros(feat.shape),
+            "std": torch.ones(feat.shape),
+        }
+    return stats
+
+
+def _find(pipeline, class_name: str):
+    return next((s for s in pipeline.steps if type(s).__name__ == class_name), None)
+
+
+def test_gripper_steps_survive_a_save_load_round_trip(tmp_path):
+    """gripper_dim/gripper_threshold must not silently revert to the class defaults.
+
+    Both steps used to inherit `get_config() -> {}`, so a reloaded pipeline came back at
+    `gripper_dim=6, threshold=0.5` no matter what the training config said.
+    """
+    cfg = make_config(action_dim=7)
+    cfg.gripper_dim = 5
+    cfg.gripper_threshold = 0.25
+    cfg.pre_snap_gripper_action = True
+    cfg.binarize_gripper_action = True
+
+    _, post = make_vla_jepa_pre_post_processors(cfg, _stats_for(cfg))
+    post.save_pretrained(str(tmp_path))
+
+    reloaded = PolicyProcessorPipeline.from_pretrained(
+        str(tmp_path), config_filename="policy_postprocessor.json", overrides={}
+    )
+    for class_name in ("PreSnapGripperProcessorStep", "BinarizeGripperProcessorStep"):
+        step = _find(reloaded, class_name)
+        assert step is not None, class_name
+        assert (step.gripper_dim, step.threshold) == (5, 0.25), class_name
+
+
+def test_gripper_steps_default_off():
+    """The LIBERO-specific gripper steps are opt-in, since they assume LIBERO's units."""
+    cfg = make_config(action_dim=7)
+    assert not cfg.pre_snap_gripper_action
+    assert not cfg.binarize_gripper_action
+    _, post = make_vla_jepa_pre_post_processors(cfg, _stats_for(cfg))
+    assert _find(post, "PreSnapGripperProcessorStep") is None
+    assert _find(post, "BinarizeGripperProcessorStep") is None
+
+
+def test_gripper_dim_out_of_range_raises():
+    """An out-of-range gripper index used to make both steps no-op silently."""
+    cfg = make_config(action_dim=7)
+    cfg.gripper_dim = 7
+    cfg.pre_snap_gripper_action = True
+    with pytest.raises(ValueError, match="out of range"):
+        cfg.validate_features()
+
+
+def test_gripper_dim_resolved_from_action_names():
+    cfg = make_config(action_dim=4)
+    cfg.gripper_dim = 6
+    cfg.action_feature_names = ["shoulder.pos", "elbow.pos", "gripper.pos", "wrist.pos"]
+    assert cfg.resolved_gripper_dim == 2
+
+
+def test_physical_range_mismatch_warns(caplog):
+    """A gripper in degrees would be pinned to a constant by the 0.5 threshold."""
+    cfg = make_config(action_dim=7)
+    cfg.gripper_dim = 6
+    cfg.pre_snap_gripper_action = True
+    cfg.binarize_gripper_action = True
+    with caplog.at_level(logging.WARNING):
+        make_vla_jepa_pre_post_processors(cfg, _stats_for(cfg, gripper_min=0.0, gripper_max=90.0))
+    assert "looks misconfigured" in caplog.text
+
+
+def test_libero_style_gripper_range_does_not_warn(caplog):
+    cfg = make_config(action_dim=7)
+    cfg.gripper_dim = 6
+    cfg.pre_snap_gripper_action = True
+    cfg.binarize_gripper_action = True
+    with caplog.at_level(logging.WARNING):
+        make_vla_jepa_pre_post_processors(cfg, _stats_for(cfg, gripper_min=-1.0, gripper_max=1.0))
+    assert "looks misconfigured" not in caplog.text
+
+
+def test_action_clipping_is_skipped_unless_min_max(caplog):
+    """Clipping to [-1, 1] is a range bound under MIN_MAX but a 1-sigma truncation under MEAN_STD."""
+    cfg = make_config(action_dim=7)
+    assert cfg.clip_normalized_actions
+    _, post = make_vla_jepa_pre_post_processors(cfg, _stats_for(cfg))
+    assert _find(post, "ClipActionsProcessorStep") is not None
+
+    cfg.normalization_mapping = {**cfg.normalization_mapping, "ACTION": NormalizationMode.MEAN_STD}
+    with caplog.at_level(logging.WARNING):
+        _, post = make_vla_jepa_pre_post_processors(cfg, _stats_for(cfg))
+    assert _find(post, "ClipActionsProcessorStep") is None
+    assert "clip_normalized_actions" in caplog.text
+
+
+def test_observation_delta_indices_collapse_without_world_model():
+    """Only the world model reads frames past index 0; asking for more decodes video for nothing."""
+    assert make_config(num_video_frames=8).observation_delta_indices == list(range(8))
+    cfg = make_config(num_video_frames=8)
+    cfg.enable_world_model = False
+    assert cfg.observation_delta_indices == [0]

tests/policies/vla_jepa/test_vla_jepa.py

@@ -1,5 +1,19 @@
 #!/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 os

tests/policies/vla_jepa/test_world_model.py

@@ -1,5 +1,19 @@
 #!/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