fix(robomme): train only on execution targets

docs/source/robomme.mdx

@@ -110,16 +110,18 @@ dataset = LeRobotDataset("lerobot/robomme")
 
 ### Dataset features
 
-| Feature            | Shape         | Description                     |
-| ------------------ | ------------- | ------------------------------- |
-| `image`            | (256, 256, 3) | Front camera RGB                |
-| `wrist_image`      | (256, 256, 3) | Wrist camera RGB                |
-| `actions`          | (8,)          | Joint angles + gripper          |
-| `state`            | (8,)          | Joint positions + gripper state |
-| `simple_subgoal`   | str           | High-level language annotation  |
-| `grounded_subgoal` | str           | Grounded language annotation    |
-| `episode_index`    | int           | Episode ID                      |
-| `frame_index`      | int           | Frame within episode            |
+| Feature            | Shape         | Description                      |
+| ------------------ | ------------- | -------------------------------- |
+| `image`            | (256, 256, 3) | Front camera RGB                 |
+| `wrist_image`      | (256, 256, 3) | Wrist camera RGB                 |
+| `actions`          | (8,)          | Joint angles + gripper           |
+| `state`            | (8,)          | Joint positions + gripper state  |
+| `simple_subgoal`   | str           | High-level language annotation   |
+| `grounded_subgoal` | str           | Grounded language annotation     |
+| `exec_start_idx`   | scalar        | First action-execution frame     |
+| `is_demo`          | bool          | Video-demonstration prefix frame |
+| `episode_index`    | int           | Episode ID                       |
+| `frame_index`      | int           | Frame within episode             |
 
 ### Feature key alignment (training)
 
@@ -132,9 +134,17 @@ uv run lerobot-train \
   --policy.path=lerobot/smolvla_base \
   --policy.empty_cameras=1 \
   --dataset.repo_id=lerobot/robomme \
+  --dataset.training_target_start_feature=exec_start_idx \
   '--rename_map={"image":"observation.images.camera1","wrist_image":"observation.images.camera2","state":"observation.state","actions":"action"}'
 ```
 
+RoboMME episodes begin with video-demonstration/context frames that should remain available to a
+memory policy but should not be sampled as action-training targets. Setting
+`dataset.training_target_start_feature=exec_start_idx` starts target sampling at each episode's
+execution boundary while preserving earlier frames for temporal observation deltas. In the published
+training split this keeps 476,857 execution targets out of 768,897 total frames. One execution-target
+epoch therefore takes `ceil(476857 / effective_batch_size)` optimizer steps.
+
 At evaluation time the environment wrapper has already converted observations to canonical keys. Only the two camera suffixes need to be aligned with the checkpoint:
 
 ```bash

examples/robomme/smolvla_visual_memory_ablation.sh

@@ -4,8 +4,11 @@
 
 set -euo pipefail
 
-STEPS="${STEPS:-30000}"
 BATCH_SIZE="${BATCH_SIZE:-4}"
+# The published training split has 476,857 execution frames. By default, train on one
+# execution-frame epoch; set STEPS explicitly to use a different optimizer-step budget.
+TARGET_SAMPLES="${TARGET_SAMPLES:-476857}"
+STEPS="${STEPS:-$(((TARGET_SAMPLES + BATCH_SIZE - 1) / BATCH_SIZE))}"
 SEED="${SEED:-1000}"
 OUTPUT_ROOT="${OUTPUT_ROOT:-outputs/robomme-smolvla-mem-ablation}"
 WANDB_ENABLE="${WANDB_ENABLE:-false}"
@@ -22,6 +25,7 @@ COMMON_TRAIN_ARGS=(
   --policy.freeze_vision_encoder=false
   --policy.train_expert_only=false
   --dataset.repo_id=lerobot/robomme
+  --dataset.training_target_start_feature=exec_start_idx
   '--rename_map={"image":"observation.images.camera1","wrist_image":"observation.images.camera2","state":"observation.state","actions":"action"}'
   --batch_size="${BATCH_SIZE}"
   --steps="${STEPS}"

src/lerobot/configs/default.py

@@ -33,6 +33,10 @@ class DatasetConfig:
     # looked up under $HF_LEROBOT_HOME/repo_id and Hub downloads use a revision-safe cache under $HF_LEROBOT_HOME/hub.
     root: str | None = None
     episodes: list[int] | None = None
+    # Optional per-frame feature containing the episode-relative index of the first frame that may be
+    # sampled as a training target. Earlier frames remain in the dataset and can still be loaded through
+    # temporal observation deltas. This is useful for datasets with demonstration/context prefixes.
+    training_target_start_feature: str | None = None
     image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
     revision: str | None = None
     use_imagenet_stats: bool = True

src/lerobot/datasets/sampler.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 import logging
 import math
-from collections.abc import Iterator
+from collections.abc import Iterator, Sequence
 
 import numpy as np
 import torch
@@ -49,7 +49,7 @@ class EpisodeAwareSampler:
         dataset_from_indices: list[int],
         dataset_to_indices: list[int],
         episode_indices_to_use: list | None = None,
-        drop_n_first_frames: int = 0,
+        drop_n_first_frames: int | Sequence[int] = 0,
         drop_n_last_frames: int = 0,
         shuffle: bool = False,
         seed: int = 0,
@@ -60,13 +60,12 @@ class EpisodeAwareSampler:
             dataset_from_indices: Start index of each episode in the dataset.
             dataset_to_indices: End index of each episode in the dataset.
             episode_indices_to_use: Episode indices to use; None means all.
-            drop_n_first_frames: Frames to drop from the start of each episode.
+            drop_n_first_frames: Frames to drop from the start of each episode. An integer applies the
+                same offset to every episode; a sequence supplies one offset per episode.
             drop_n_last_frames: Frames to drop from the end of each episode.
             shuffle: Whether to shuffle the indices.
             seed: Seed the permutation is derived from (together with the epoch).
         """
-        if drop_n_first_frames < 0:
-            raise ValueError(f"drop_n_first_frames must be >= 0, got {drop_n_first_frames}")
         if drop_n_last_frames < 0:
             raise ValueError(f"drop_n_last_frames must be >= 0, got {drop_n_last_frames}")
 
@@ -78,12 +77,27 @@ class EpisodeAwareSampler:
                 f"got {len(from_indices)} and {len(to_indices)}"
             )
 
+        if isinstance(drop_n_first_frames, int):
+            first_frame_offsets = np.full(len(from_indices), drop_n_first_frames, dtype=np.int64)
+        else:
+            first_frame_offsets = np.asarray(drop_n_first_frames, dtype=np.int64)
+            if first_frame_offsets.shape != from_indices.shape:
+                raise ValueError(
+                    "drop_n_first_frames must be an integer or have one value per episode; "
+                    f"got {len(first_frame_offsets)} values for {len(from_indices)} episodes"
+                )
+        if np.any(first_frame_offsets < 0):
+            raise ValueError(
+                "drop_n_first_frames must be >= 0, got "
+                f"{first_frame_offsets[first_frame_offsets < 0].tolist()}"
+            )
+
         used = np.ones(len(from_indices), dtype=bool)
         if episode_indices_to_use is not None:
             used = np.zeros(len(from_indices), dtype=bool)
             used[np.asarray(episode_indices_to_use, dtype=np.int64)] = True
 
-        starts = from_indices + drop_n_first_frames
+        starts = from_indices + first_frame_offsets
         lengths = to_indices - drop_n_last_frames - starts
         for episode_idx in np.flatnonzero(used & (lengths <= 0)):
             logger.warning(
@@ -91,7 +105,7 @@ class EpisodeAwareSampler:
                 "drop_n_last_frames=%d removes all frames. Skipping.",
                 episode_idx,
                 to_indices[episode_idx] - from_indices[episode_idx],
-                drop_n_first_frames,
+                first_frame_offsets[episode_idx],
                 drop_n_last_frames,
             )
         used &= lengths > 0

src/lerobot/scripts/lerobot_train.py

@@ -455,10 +455,46 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         to_indices = dataset.meta.episodes["dataset_to_index"]
         seed = cfg.seed if cfg.seed is not None else 0
 
+        drop_n_first_frames: int | list[int] = 0
+        target_start_feature = cfg.dataset.training_target_start_feature
+        if target_start_feature is not None:
+            if not hasattr(dataset, "hf_dataset"):
+                raise ValueError(
+                    "dataset.training_target_start_feature is only supported for a single map-style "
+                    "LeRobotDataset."
+                )
+            if target_start_feature not in dataset.hf_dataset.column_names:
+                raise ValueError(
+                    f"Training target start feature {target_start_feature!r} is not present in the dataset. "
+                    f"Available features: {dataset.hf_dataset.column_names}"
+                )
+            start_column = dataset.hf_dataset.data.column(target_start_feature)
+            absolute_to_relative_idx = dataset.absolute_to_relative_idx
+            episode_indices = dataset.episodes if dataset.episodes is not None else range(len(from_indices))
+            drop_n_first_frames = [0] * len(from_indices)
+            for episode_index in episode_indices:
+                absolute_index = int(from_indices[episode_index])
+                relative_index = (
+                    absolute_to_relative_idx[absolute_index]
+                    if absolute_to_relative_idx is not None
+                    else absolute_index
+                )
+                drop_n_first_frames[episode_index] = int(start_column[relative_index].as_py())
+            logging.info(
+                "Restricting training targets with %s: keeping %d of %d frames",
+                target_start_feature,
+                sum(
+                    int(to_indices[index]) - int(from_indices[index]) - drop_n_first_frames[index]
+                    for index in episode_indices
+                ),
+                sum(int(to_indices[index]) - int(from_indices[index]) for index in episode_indices),
+            )
+
         sampler = EpisodeAwareSampler(
             from_indices,
             to_indices,
             episode_indices_to_use=dataset.episodes,
+            drop_n_first_frames=drop_n_first_frames,
             drop_n_last_frames=getattr(active_cfg, "drop_n_last_frames", 0),
             shuffle=True,
             seed=seed,

tests/datasets/test_sampler.py

@@ -60,6 +60,26 @@ def test_drop_n_first_frames():
     assert list(sampler) == [1, 4, 5]
 
 
+def test_drop_different_number_of_first_frames_per_episode():
+    sampler = EpisodeAwareSampler(
+        dataset_from_indices=[0, 3, 5],
+        dataset_to_indices=[3, 5, 9],
+        drop_n_first_frames=[1, 0, 3],
+    )
+    assert sampler.indices == [1, 2, 3, 4, 8]
+    assert len(sampler) == 5
+
+
+def test_drop_first_frames_sequence_must_match_episode_count():
+    with pytest.raises(ValueError, match="one value per episode"):
+        EpisodeAwareSampler([0, 3], [3, 6], drop_n_first_frames=[1])
+
+
+def test_drop_first_frames_sequence_must_be_non_negative():
+    with pytest.raises(ValueError, match="must be >= 0"):
+        EpisodeAwareSampler([0, 3], [3, 6], drop_n_first_frames=[1, -1])
+
+
 def test_drop_n_last_frames():
     dataset = Dataset.from_dict(
         {