Merge branch 'main' into fix/zero-shaped-features

Per review, log a warning when compute_episode_stats skips a feature with a
zero-width shape, so users know stats were intentionally not computed for it.

.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@e60a538eea9817ab312196d0d233604b01697265  # main
+    uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@2430c1ec91d04667414e2fa31ecfc36c153ea391  # 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@e60a538eea9817ab312196d0d233604b01697265  # main
+    uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@2430c1ec91d04667414e2fa31ecfc36c153ea391  # 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   |          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%** |                                                                                                               |
+| 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%** |                                                                                                               |
 
 ```bash
 export MODEL_ID=your_trained_model_on_huggingface

src/lerobot/datasets/compute_stats.py

@@ -519,6 +519,13 @@ 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:
-            hf_features[key] = datasets.Sequence(
-                length=ft["shape"][0], feature=datasets.Value(dtype=ft["dtype"])
-            )
+            # 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"]))
         elif len(ft["shape"]) == 2:
             hf_features[key] = datasets.Array2D(shape=ft["shape"], dtype=ft["dtype"])
         elif len(ft["shape"]) == 3:

src/lerobot/policies/pi0_fast/modeling_pi0_fast.py

@@ -613,14 +613,15 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
         device = tokens.device
         lm_head = self.paligemma_with_expert.paligemma.lm_head
 
-        # NOTE (bug 2 fix): do NOT append a second <bos> here. The language tokens
-        # already begin with <bos> (standard PaliGemma prefix "[image] <bos> prompt \n").
-        # Appending another <bos> right before decoding pushes the checkpoint into a
-        # bos->bos attractor and yields degenerate generation. Generate directly after
-        # the prompt instead.
+        # add bos token after tokens
+        bos_token = torch.full(
+            (bsize, 1), self._paligemma_tokenizer.bos_token_id, dtype=torch.long, device=device
+        )
+        tokens = torch.cat([tokens, bos_token], dim=1)
+        masks = torch.cat([masks, torch.ones((bsize, 1), dtype=torch.bool, device=device)], dim=1)
 
         # 1. Initial Embedding (matches training prefix)
-        # prefix_embs will include [Images, Language Prompt]
+        # prefix_embs will include [Images, Language Prompt, BOS]
         prefix_embs, prefix_pad_masks, prefix_att_masks, total_t_images, _ = self.embed_prefix_fast(
             images, img_masks, tokens, masks, fast_action_tokens=None, fast_action_masks=None
         )
@@ -708,13 +709,14 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
         # --- 1. PREFILL PHASE ---
         # Process Images + Text Prompt + BOS token once to populate the KV cache.
 
-        # NOTE (bug 2 fix): do NOT append a second <bos> here. The language tokens
-        # already begin with <bos> (standard PaliGemma prefix "[image] <bos> prompt \n").
-        # A second <bos> right before decoding causes degenerate bos->bos generation.
-        tokens_in = tokens
-        masks_in = masks
+        # Add BOS token to the prompt
+        bos_token = torch.full(
+            (bsize, 1), self._paligemma_tokenizer.bos_token_id, dtype=torch.long, device=device
+        )
+        tokens_in = torch.cat([tokens, bos_token], dim=1)
+        masks_in = torch.cat([masks, torch.ones((bsize, 1), dtype=torch.bool, device=device)], dim=1)
 
-        # Embed prefix [Images, Language]
+        # Embed prefix [Images, Language, BOS]
         # fast_action_tokens=None means we are just embedding the condition (images+text)
         prefix_embs, prefix_pad_masks, prefix_att_masks, total_t_images, _ = self.embed_prefix_fast(
             images, img_masks, tokens_in, masks_in, fast_action_tokens=None, fast_action_masks=None

src/lerobot/processor/tokenizer_processor.py

@@ -476,12 +476,11 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
             if tokens.dim() > 1:
                 tokens = tokens.flatten()
 
-            # NOTE (bug 2 fix): do NOT prepend a <bos> to the action target. The prompt
-            # already carries the leading <bos>; a second one before "Action:" mismatches
-            # the generation-time prefix (see sample_actions_fast*) and drives degenerate
-            # bos->bos decoding. Target is "Action: <fast tokens> |".
+            bos_id = self._paligemma_tokenizer.bos_token_id
+            # add bos
             tokens = torch.cat(
                 [
+                    torch.tensor([bos_id], device=action.device),
                     torch.tensor(
                         self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False),
                         device=action.device,

tests/datasets/test_compute_stats.py

@@ -13,6 +13,7 @@
 # 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
@@ -687,6 +688,28 @@ 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,3 +1804,11 @@ 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)