Merge branch 'main' into fix/evo1-repad-normalizer-on-load

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>

.github/workflows/documentation.yml

@@ -55,7 +55,7 @@ jobs:
       github.repository == 'huggingface/lerobot'
     permissions:
       contents: read
-    uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@2430c1ec91d04667414e2fa31ecfc36c153ea391  # main
+    uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@e60a538eea9817ab312196d0d233604b01697265  # main
     with:
       commit_sha: ${{ github.sha }}
       package: lerobot
@@ -78,7 +78,7 @@ jobs:
     permissions:
       contents: read
       pull-requests: write
-    uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@2430c1ec91d04667414e2fa31ecfc36c153ea391  # main
+    uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@e60a538eea9817ab312196d0d233604b01697265  # main
     with:
       commit_sha: ${{ github.event.pull_request.head.sha }}
       pr_number: ${{ github.event.number }}

docs/source/groot.mdx

@@ -162,11 +162,11 @@ Preliminary LeRobot integration results (GR00T-LeRobot, `eval.n_episodes >= 50`
 
 | Suite            | Success rate | Checkpoint                                                                                                    |
 | ---------------- | -----------: | ------------------------------------------------------------------------------------------------------------- |
-| LIBERO Spatial   |          91% | [nvidia/gr00t17-lerobot-libero_spatial-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_spatial-640) |
-| LIBERO Object    |          81% | [nvidia/gr00t17-lerobot-libero_object-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_object-640)   |
-| LIBERO Goal      |          97% | [nvidia/gr00t17-lerobot-libero_goal-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_goal-640)       |
-| LIBERO 10 (Long) |          84% | [nvidia/gr00t17-lerobot-libero_10-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_10-640)           |
-| **Average**      |   **88.25%** |                                                                                                               |
+| LIBERO Spatial   |          95% | [nvidia/gr00t17-lerobot-libero_spatial-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_spatial-640) |
+| LIBERO Object    |         100% | [nvidia/gr00t17-lerobot-libero_object-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_object-640)   |
+| LIBERO Goal      |          98% | [nvidia/gr00t17-lerobot-libero_goal-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_goal-640)       |
+| LIBERO 10 (Long) |          93% | [nvidia/gr00t17-lerobot-libero_10-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_10-640)           |
+| **Average**      |    **96.5%** |                                                                                                               |
 
 ```bash
 export MODEL_ID=your_trained_model_on_huggingface

src/lerobot/datasets/compute_stats.py

@@ -519,13 +519,6 @@ def compute_episode_stats(
         if features[key]["dtype"] in {"string", "language"}:
             continue
 
-        # Features with zero-width shapes are skipped (no data to compute stats on)
-        if any(d == 0 for d in features[key].get("shape", ())):
-            logging.debug(
-                f"Skipping statistics computation for feature '{key}' with a zero-width shape {features[key]['shape']}."
-            )
-            continue
-
         if features[key]["dtype"] in ["image", "video"]:
             ep_ft_array = sample_images(data)
             axes_to_reduce = (0, 2, 3)

src/lerobot/datasets/feature_utils.py

@@ -67,9 +67,9 @@ def get_hf_features_from_features(features: dict) -> datasets.Features:
         elif ft["shape"] == (1,):
             hf_features[key] = datasets.Value(dtype=ft["dtype"])
         elif len(ft["shape"]) == 1:
-            # pyarrow rejects fixed-size lists of length 0, so use a variable length list instead
-            length = ft["shape"][0] if ft["shape"][0] > 0 else -1
-            hf_features[key] = datasets.Sequence(length=length, feature=datasets.Value(dtype=ft["dtype"]))
+            hf_features[key] = datasets.Sequence(
+                length=ft["shape"][0], feature=datasets.Value(dtype=ft["dtype"])
+            )
         elif len(ft["shape"]) == 2:
             hf_features[key] = datasets.Array2D(shape=ft["shape"], dtype=ft["dtype"])
         elif len(ft["shape"]) == 3:

src/lerobot/policies/evo1/processor_evo1.py

@@ -302,6 +302,33 @@ def _pad_evo1_stats(
     return padded_stats
 
 
+def _refresh_evo1_normalization_steps(
+    config: Evo1Config,
+    preprocessor: PolicyProcessorPipeline,
+    postprocessor: PolicyProcessorPipeline,
+) -> None:
+    """Re-pad checkpoint-loaded (un)normalizer stats/features to EVO1's fixed widths.
+
+    Loading a checkpoint injects the raw dataset stats (unpadded to max_state_dim/max_action_dim)
+    into the (un)normalizer via the generic override path in make_pre_post_processors. Those stats
+    and their declared features must be re-padded/reshaped to EVO1's fixed widths, otherwise
+    normalization fails against the padded state/action tensors (e.g. state padded to 24 vs. 8-dim
+    LIBERO stats). Padding is a no-op when stats are already at the target width.
+    """
+    normalization_features = _evo1_normalization_features(config)
+    action_features = _evo1_action_features(config)
+    for step in preprocessor.steps:
+        if isinstance(step, NormalizerProcessorStep):
+            step.features = normalization_features
+            step.stats = _pad_evo1_stats(config, step.stats)
+            step.to(device=step.device, dtype=step.dtype)
+    for step in postprocessor.steps:
+        if isinstance(step, UnnormalizerProcessorStep):
+            step.features = action_features
+            step.stats = _pad_evo1_stats(config, step.stats)
+            step.to(device=step.device, dtype=step.dtype)
+
+
 def reconcile_evo1_processors(
     config: Evo1Config,
     preprocessor: PolicyProcessorPipeline,
@@ -309,16 +336,19 @@ def reconcile_evo1_processors(
 ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
     """Reconcile checkpoint-loaded pipelines with the current EVO1 config.
 
-    Two things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
-    (converters are plain functions and are never serialized), and eval-time CLI overrides of the
-    action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`). This
-    restores the converter and rebuilds the action step from the current config so those overrides
-    take effect.
+    Three things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
+    (converters are plain functions and are never serialized), eval-time CLI overrides of the
+    action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`), and the
+    (un)normalizer stats/features when the generic override path injects raw, unpadded dataset
+    stats. This restores the converter, re-pads the normalization stats to EVO1's fixed widths, and
+    rebuilds the action step from the current config so those overrides take effect.
     """
     # Pipelines reloaded from a checkpoint come back with the default batch converter, which drops
     # non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1.
     preprocessor.to_transition = evo1_batch_to_transition
 
+    _refresh_evo1_normalization_steps(config, preprocessor, postprocessor)
+
     action_step = Evo1ActionProcessorStep(
         action_dim=_evo1_action_dim(config),
         binarize_gripper=config.binarize_gripper,

tests/datasets/test_compute_stats.py

@@ -13,7 +13,6 @@
 # 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 logging
 from unittest.mock import patch
 
 import numpy as np
@@ -688,28 +687,6 @@ def test_compute_episode_stats_string_features_skipped():
     assert "q01" in stats["action"]
 
 
-def test_compute_episode_stats_zero_width_features_skipped(caplog):
-    """Test that features with a zero-width dim (e.g. shape=(0,)) are skipped with a debug log."""
-    episode_data = {
-        "empty": np.zeros((100, 0), dtype=np.float32),  # Zero-width feature
-        "action": np.random.normal(0, 1, (100, 5)),
-    }
-    features = {
-        "empty": {"dtype": "float32", "shape": (0,)},
-        "action": {"dtype": "float32", "shape": (5,)},
-    }
-
-    with caplog.at_level(logging.DEBUG):
-        stats = compute_episode_stats(episode_data, features)
-
-    # Zero-width features should be skipped with a debug log, others computed as usual
-    assert "empty" not in stats
-    assert "empty" in caplog.text
-    assert "action" in stats
-    assert "q01" in stats["action"]
-    assert stats["action"]["mean"].shape == (5,)
-
-
 def test_aggregate_feature_stats_with_quantiles():
     """Test aggregating feature stats that include quantiles."""
     stats_ft_list = [

tests/datasets/test_datasets.py

@@ -1804,11 +1804,3 @@ def test_episode_filter_unknown_key_raises(tmp_path, lerobot_dataset_factory):
             root=dataset.root,
             episode_filter=lambda ep: ep["not_a_real_field"] > 0,
         )
-
-
-def test_get_hf_features_zero_width_feature_does_not_raise_on_from_dict():
-    import datasets
-
-    features = {"empty": {"dtype": "float32", "shape": (0,), "names": ["empty"]}}
-    hf_features = get_hf_features_from_features(features)
-    datasets.Dataset.from_dict({"empty": [[], []]}, features=hf_features)

tests/policies/evo1/test_evo1.py

@@ -496,6 +496,60 @@ def test_evo1_processor_save_load_round_trip_applies_config_overrides(tmp_path):
     assert "embodiment_id" in processed
 
 
+def test_reconcile_evo1_processors_repads_overridden_stats(tmp_path):
+    """Loading a checkpoint and injecting raw (unpadded) dataset stats must be re-padded.
+
+    Regression test: lerobot-train passes the raw dataset stats as normalizer/unnormalizer
+    overrides when resuming from a checkpoint (e.g. stage2 from a stage1 checkpoint). Those stats
+    are at the dataset dims (e.g. LIBERO state=8/action=7), but EVO1 pads state/action to
+    max_state_dim/max_action_dim before normalization, so reconcile_evo1_processors must re-pad the
+    stats or normalization crashes with a shape mismatch.
+    """
+    config = make_config()
+    preprocessor, postprocessor = make_evo1_pre_post_processors(config, dataset_stats=make_stats())
+    preprocessor.save_pretrained(tmp_path)
+    postprocessor.save_pretrained(tmp_path)
+
+    # Reload with the generic override path injecting raw, unpadded dataset stats.
+    raw_stats = make_stats()
+    loaded_pre = PolicyProcessorPipeline.from_pretrained(
+        tmp_path,
+        config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json",
+        overrides={"normalizer_processor": {"stats": raw_stats}},
+        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",
+        overrides={"unnormalizer_processor": {"stats": raw_stats}},
+        to_transition=policy_action_to_transition,
+        to_output=transition_to_policy_action,
+    )
+
+    # Sanity: the override really injected unpadded stats before reconciliation.
+    normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
+    assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (STATE_DIM,)
+
+    loaded_pre, loaded_post = reconcile_evo1_processors(config, loaded_pre, loaded_post)
+
+    normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
+    unnormalizer = next(step for step in loaded_post.steps if isinstance(step, UnnormalizerProcessorStep))
+    assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (MAX_STATE_DIM,)
+    assert normalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
+    assert unnormalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
+
+    # Normalizing a padded state must not raise (this is the exact runtime path that crashed).
+    processed = loaded_pre(
+        {
+            "task": "pick the block",
+            OBS_STATE: torch.zeros(STATE_DIM),
+            f"{OBS_IMAGES}.front": torch.rand(3, 16, 16),
+        }
+    )
+    assert processed[OBS_STATE].shape == (1, MAX_STATE_DIM)
+
+
 def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch):
     monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
     policy = modeling_evo1.Evo1Policy(make_config())