diff --git a/docs/source/robomme.mdx b/docs/source/robomme.mdx index cb9e6bc16..c93764b5f 100644 --- a/docs/source/robomme.mdx +++ b/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 diff --git a/examples/robomme/smolvla_visual_memory_ablation.sh b/examples/robomme/smolvla_visual_memory_ablation.sh index 437ebffba..09d03b485 100755 --- a/examples/robomme/smolvla_visual_memory_ablation.sh +++ b/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}" diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py index 38991a665..03d7ef2bc 100644 --- a/src/lerobot/configs/default.py +++ b/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 diff --git a/src/lerobot/datasets/sampler.py b/src/lerobot/datasets/sampler.py index aee6ce46d..9cdd85372 100644 --- a/src/lerobot/datasets/sampler.py +++ b/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 diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 0859a7fd8..7fcd406b6 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/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, diff --git a/tests/datasets/test_sampler.py b/tests/datasets/test_sampler.py index dbe2eb7f0..de65229fa 100644 --- a/tests/datasets/test_sampler.py +++ b/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( {