diff --git a/docs/source/language_and_recipes.mdx b/docs/source/language_and_recipes.mdx index 4181dbe34..037dc4182 100644 --- a/docs/source/language_and_recipes.mdx +++ b/docs/source/language_and_recipes.mdx @@ -108,6 +108,7 @@ own binding plus a matching image block, e.g. ```yaml ask_vqa_top: + route: vqa bindings: vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.top)" vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.top)" @@ -127,7 +128,9 @@ ask_vqa_top: } ``` -Add one such sub-recipe per camera the dataset records. +Add one such sub-recipe per camera the dataset records. The explicit +`route: vqa` marker makes a matching sparse VQA annotation take precedence +over normal weighted blend selection; component names are purely descriptive. ## Layer 3 — training format @@ -141,7 +144,20 @@ sample["target_message_indices"] The renderer does not apply a tokenizer chat template. Policy processors decide how to serialize the messages for their backbone, which keeps the same dataset usable across SmolVLA, Pi0.5, and any future VLM that expects OpenAI-style chat messages. +## Blends + +Blend recipes select one weighted sub-recipe deterministically from the sample index. +`recipes/subtask_mem.yaml` trains the compact core blend — high-level subtask prediction, low-level execution, and memory. `recipes/subtask_mem_vqa_speech.yaml` is the fuller variant that also adds VQA and spoken interjection responses. + +`recipes/subtask_joint.yaml` demonstrates joint sequence training rather than a +weighted blend. For the same sample, its assistant subtask is supervised with +text cross-entropy on the `low_level` stream while action prediction remains +active, matching the joint setup from the π0.5 paper. Enable +`--policy.joint_subtask_conditioning=true` to use that subtask conditioning at inference. + ## Graceful absence -If both language columns are missing, `None`, or empty, `RenderMessagesStep` is a no-op. -If an event-scoped branch is selected on a frame without the required event row, rendering returns `None`, allowing a loader to retry another sample. +If both language columns are missing, `None`, or empty, `RenderMessagesStep` uses +the task string as low-level supervision when available and otherwise leaves the +sample unchanged. For an annotated sample, if no recipe branch applies and no +task fallback exists, rendering returns `None`, allowing a loader to retry another sample. diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py index 0734d0bc2..6cbda669a 100644 --- a/src/lerobot/configs/default.py +++ b/src/lerobot/configs/default.py @@ -14,6 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging from dataclasses import dataclass, field from lerobot.transforms import ImageTransformsConfig @@ -21,6 +22,8 @@ from lerobot.utils.import_utils import get_safe_default_video_backend from .video import DEFAULT_DEPTH_UNIT, DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT +logger = logging.getLogger(__name__) + @dataclass class DatasetConfig: @@ -36,6 +39,8 @@ 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 + # Episode indices to drop (e.g. corrupt or heterogeneous ones). Applied on top of `episodes`. + exclude_episodes: list[int] | None = None image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig) revision: str | None = None use_imagenet_stats: bool = True @@ -75,6 +80,14 @@ class DatasetConfig: if len(self.episodes) != len(set(self.episodes)): duplicates = sorted({ep for ep in self.episodes if self.episodes.count(ep) > 1}) raise ValueError(f"Episode indices contain duplicates: {duplicates}") + if self.exclude_episodes is not None: + negative_episodes = [episode for episode in self.exclude_episodes if episode < 0] + if negative_episodes: + logger.warning( + "Ignoring negative exclude_episodes entries: %s", + negative_episodes, + ) + self.exclude_episodes = [episode for episode in self.exclude_episodes if episode >= 0] @dataclass diff --git a/src/lerobot/configs/recipe.py b/src/lerobot/configs/recipe.py index 28e5a0db3..9f41feada 100644 --- a/src/lerobot/configs/recipe.py +++ b/src/lerobot/configs/recipe.py @@ -23,6 +23,7 @@ from typing import Any, Literal, get_args MessageRole = Literal["user", "assistant", "system", "tool"] MessageStream = Literal["high_level", "low_level"] +RecipeRoute = Literal["vqa"] DEFAULT_BINDINGS = { "subtask": "active_at(t, style=subtask)", @@ -40,6 +41,7 @@ discovery (here) and rendered-message substitution (in ``language_render``).""" _VALID_ROLES = frozenset(get_args(MessageRole)) _VALID_STREAMS = frozenset(get_args(MessageStream)) +_VALID_ROUTES = frozenset(get_args(RecipeRoute)) @dataclass @@ -78,7 +80,7 @@ class MessageTurn: raise ValueError(f"Unsupported message stream: {self.stream!r}") if self.content is None and self.tool_calls_from is None: raise ValueError("MessageTurn.content is required unless tool_calls_from is set.") - if self.content is not None and not isinstance(self.content, (str, list)): + if self.content is not None and not isinstance(self.content, str | list): raise TypeError("MessageTurn.content must be a string, a list of HF-style blocks, or None.") if isinstance(self.content, list): for block in self.content: @@ -99,13 +101,16 @@ class TrainingRecipe: A recipe is either a *message recipe* (``messages`` plus optional ``bindings``) or a *blend recipe* (``blend`` mapping names to weighted - sub-recipes). ``weight`` is only meaningful inside a blend. + sub-recipes). ``weight`` and ``route`` are only meaningful inside a blend; + ``route: vqa`` gives sparse VQA annotations priority over normal weighted + selection. """ messages: list[MessageTurn] | None = None bindings: dict[str, str] | None = None blend: dict[str, TrainingRecipe] | None = None weight: float | None = None + route: RecipeRoute | None = None def __post_init__(self) -> None: """Validate that exactly one of ``messages`` or ``blend`` is set.""" @@ -113,6 +118,10 @@ class TrainingRecipe: raise ValueError("TrainingRecipe must set only one of messages or blend.") if self.messages is None and self.blend is None: raise ValueError("TrainingRecipe must set one of messages or blend.") + if self.route is not None and self.route not in _VALID_ROUTES: + raise ValueError(f"Unsupported recipe route: {self.route!r}") + if self.blend is not None and self.route is not None: + raise ValueError("TrainingRecipe.route may only be set on a message recipe inside a blend.") if self.messages is not None: self._validate_message_recipe() @@ -147,8 +156,9 @@ class TrainingRecipe: return cls.from_dict(data) def _validate_message_recipe(self) -> None: - """Ensure every templated binding is known and at least one turn is a target.""" - assert self.messages is not None + """Validate bindings and require text or low-level action supervision.""" + if self.messages is None: + raise ValueError("Cannot validate a message recipe without messages.") known_bindings = set(DEFAULT_BINDINGS) | set(self.bindings or {}) | {"task"} for turn in self.messages: @@ -156,12 +166,19 @@ class TrainingRecipe: if missing: raise ValueError(f"MessageTurn references unknown binding(s): {sorted(missing)}") - if not any(turn.target for turn in self.messages): - raise ValueError("Message recipes must contain at least one target turn.") + has_target = any(turn.target for turn in self.messages) + has_low_level = any(turn.stream == "low_level" for turn in self.messages) + if not (has_target or has_low_level): + raise ValueError( + "Message recipes must contain at least one supervised turn — " + "either ``target: true`` (text CE) or ``stream: low_level`` " + "(flow/action loss)." + ) def _validate_blend_recipe(self) -> None: """Ensure each blend component is a non-empty, weighted message recipe.""" - assert self.blend is not None + if self.blend is None: + raise ValueError("Cannot validate a blend recipe without blend components.") if not self.blend: raise ValueError("Blend recipes must contain at least one component.") diff --git a/src/lerobot/configs/recipes/subtask.yaml b/src/lerobot/configs/recipes/subtask.yaml new file mode 100644 index 000000000..c90ca8f78 --- /dev/null +++ b/src/lerobot/configs/recipes/subtask.yaml @@ -0,0 +1,16 @@ +# Predicts subtasks from tasks and trains subtask-conditioned action flow without memory or plans. +# Requires `subtask` annotations; samples with missing `if_present` bindings do not render. + +blend: + + high_level_subtask: + weight: 0.30 + messages: + - {role: user, content: "${task}", stream: high_level} + - {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask} + + low_level_execution: + weight: 0.70 + messages: + # The low-level stream trains action flow on the generated or annotated subtask. + - {role: user, content: "${subtask}", stream: low_level, if_present: subtask} diff --git a/src/lerobot/configs/recipes/subtask_joint.yaml b/src/lerobot/configs/recipes/subtask_joint.yaml new file mode 100644 index 000000000..5258ccf23 --- /dev/null +++ b/src/lerobot/configs/recipes/subtask_joint.yaml @@ -0,0 +1,13 @@ +# Paper-style joint sequence (pi0.5 §IV-B): one sample supervises the subtask +# text with CE and, because the assistant turn is part of the prefix, conditions +# the FAST and flow action losses on the same annotated subtask in one forward. +# The supervised span is attended causally; the action losses see task + subtask. +# +# Pair with `--policy.joint_subtask_conditioning=true` at inference so the flow +# prefix reproduces this layout (task turn with state + causal generated subtask). +# Samples without a `subtask` annotation fall back to a plain task-prompt +# low-level sample via `if_present`. + +messages: + - {role: user, content: "${task}", stream: low_level} + - {role: assistant, content: "${subtask}", stream: low_level, target: true, if_present: subtask} diff --git a/src/lerobot/configs/recipes/subtask_mem.yaml b/src/lerobot/configs/recipes/subtask_mem.yaml new file mode 100644 index 000000000..d12fe5009 --- /dev/null +++ b/src/lerobot/configs/recipes/subtask_mem.yaml @@ -0,0 +1,30 @@ +# Trains subtask prediction, subtask-conditioned action flow, and memory updates without plans. +# Requires `subtask` and `memory`; missing `if_present` bindings skip the affected sub-recipe. + +blend: + + high_level_subtask: + weight: 0.25 + messages: + - {role: user, content: "${task}", stream: high_level} + - {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask} + + low_level_execution: + weight: 0.60 + messages: + # The low-level stream trains action flow on the generated or annotated subtask. + - {role: user, content: "${subtask}", stream: low_level, if_present: subtask} + + memory_update: + # `active_at` densifies sparse boundaries while preserving the prior-memory/subtask mapping. + # Inference controls update timing through `subtask_change` events. + weight: 0.15 + bindings: + prior_memory: "nth_prev(style=memory, offset=1)" + current_memory: "active_at(t, style=memory)" + completed_subtask: "nth_prev(style=subtask, offset=1)" + messages: + - {role: user, content: "${task}", stream: high_level} + - {role: assistant, content: "Previous memory: ${prior_memory}", stream: high_level, if_present: prior_memory} + - {role: user, content: "Completed subtask: ${completed_subtask}", stream: high_level, if_present: completed_subtask} + - {role: assistant, content: "${current_memory}", stream: high_level, target: true, if_present: current_memory} diff --git a/src/lerobot/configs/recipes/subtask_mem_vqa_speech.yaml b/src/lerobot/configs/recipes/subtask_mem_vqa_speech.yaml new file mode 100644 index 000000000..847093c84 --- /dev/null +++ b/src/lerobot/configs/recipes/subtask_mem_vqa_speech.yaml @@ -0,0 +1,72 @@ +# Adds memory, spoken interjection responses, and camera-grounded VQA to subtask/action training. +# Missing optional annotations skip only their sub-recipe; `say` tool calls tokenize as `...`. + +blend: + + high_level_subtask: + weight: 0.25 + messages: + - {role: user, content: "${task}", stream: high_level} + - {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask} + + low_level_execution: + weight: 0.40 + messages: + # The low-level stream trains action flow on the generated or annotated subtask. + - {role: user, content: "${subtask}", stream: low_level, if_present: subtask} + + memory_update: + # `active_at` densifies sparse boundaries while preserving the prior-memory/subtask mapping. + # Inference controls update timing through `subtask_change` events. + weight: 0.10 + bindings: + prior_memory: "nth_prev(style=memory, offset=1)" + current_memory: "active_at(t, style=memory)" + completed_subtask: "nth_prev(style=subtask, offset=1)" + messages: + - {role: user, content: "${task}", stream: high_level} + - {role: assistant, content: "Previous memory: ${prior_memory}", stream: high_level, if_present: prior_memory} + - {role: user, content: "Completed subtask: ${completed_subtask}", stream: high_level, if_present: completed_subtask} + - {role: assistant, content: "${current_memory}", stream: high_level, target: true, if_present: current_memory} + + user_interjection_response: + weight: 0.10 + bindings: + interjection: "emitted_at(t, style=interjection)" + speech: "emitted_at(t, role=assistant, tool_name=say)" + messages: + - {role: user, content: "${task}", stream: high_level} + - {role: user, content: "${interjection}", stream: high_level, if_present: interjection} + # The assistant target is a `say` tool call flattened to a `...` marker. + - {role: assistant, stream: high_level, target: true, if_present: speech, tool_calls_from: speech} + + # Each camera uses a separate VQA sub-recipe for view-specific binding. + ask_vqa_top: + weight: 0.075 + route: vqa + bindings: + vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.front)" + vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.front)" + messages: + - role: user + stream: high_level + if_present: vqa_query + content: + - {type: image, feature: observation.images.front} + - {type: text, text: "${vqa_query}"} + - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa} + + ask_vqa_wrist: + weight: 0.075 + route: vqa + bindings: + vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.wrist)" + vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.wrist)" + messages: + - role: user + stream: high_level + if_present: vqa_query + content: + - {type: image, feature: observation.images.wrist} + - {type: text, text: "${vqa_query}"} + - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa} diff --git a/src/lerobot/data_processing/sarm_annotations/subtask_annotation.py b/src/lerobot/data_processing/sarm_annotations/subtask_annotation.py index b26257d44..c97860020 100644 --- a/src/lerobot/data_processing/sarm_annotations/subtask_annotation.py +++ b/src/lerobot/data_processing/sarm_annotations/subtask_annotation.py @@ -76,7 +76,7 @@ import torch from pydantic import BaseModel, Field from transformers import AutoProcessor, Qwen3VLMoeForConditionalGeneration -from lerobot.datasets import LeRobotDataset +from lerobot.datasets import LeRobotDataset, resolve_episode_indices # Pydantic Models for SARM Subtask Annotation @@ -1049,7 +1049,10 @@ def main(): torch_dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[args.dtype] # Determine episodes - episode_indices = args.episodes or list(range(dataset.meta.total_episodes)) + resolved_episodes = resolve_episode_indices(args.episodes, dataset.meta.total_episodes) + episode_indices = ( + resolved_episodes if resolved_episodes is not None else list(range(dataset.meta.total_episodes)) + ) existing_annotations = load_annotations_from_dataset(dataset.root, prefix="sparse") if args.skip_existing: diff --git a/src/lerobot/datasets/__init__.py b/src/lerobot/datasets/__init__.py index 7715a115e..7bee46815 100644 --- a/src/lerobot/datasets/__init__.py +++ b/src/lerobot/datasets/__init__.py @@ -52,7 +52,7 @@ from .pipeline_features import aggregate_pipeline_dataset_features, create_initi from .pyav_utils import check_video_encoder_parameters_pyav, detect_available_encoders_pyav from .sampler import EpisodeAwareSampler, compute_sampler_state from .streaming_dataset import StreamingLeRobotDataset -from .utils import DEFAULT_EPISODES_PATH, create_lerobot_dataset_card +from .utils import DEFAULT_EPISODES_PATH, create_lerobot_dataset_card, resolve_episode_indices from .video_utils import VideoEncodingManager # NOTE: Low-level I/O functions (cast_stats_to_numpy, get_parquet_file_size_in_mb, etc.) @@ -97,6 +97,7 @@ __all__ = [ "reencode_dataset", "remove_feature", "resolve_delta_timestamps", + "resolve_episode_indices", "safe_stop_image_writer", "split_dataset", "write_stats", diff --git a/src/lerobot/datasets/dataset_reader.py b/src/lerobot/datasets/dataset_reader.py index f4e1f6a31..0c1c0c914 100644 --- a/src/lerobot/datasets/dataset_reader.py +++ b/src/lerobot/datasets/dataset_reader.py @@ -39,6 +39,7 @@ from .io_utils import ( hf_transform_to_torch, load_nested_dataset, ) +from .utils import resolve_episode_indices from .video_utils import decode_video_frames @@ -83,7 +84,7 @@ class DatasetReader: """ self._meta = meta self.root = root - self.episodes = episodes + self.episodes = resolve_episode_indices(episodes, meta.total_episodes) self._tolerance_s = tolerance_s self._video_backend = video_backend if image_transforms is not None and not callable(image_transforms): @@ -163,10 +164,34 @@ class DatasetReader: def _load_hf_dataset(self) -> datasets.Dataset: """hf_dataset contains all the observations, states, actions, rewards, etc.""" features = get_hf_features_from_features(self._meta.features) + self._validate_language_columns_declared(features) hf_dataset = load_nested_dataset(self.root / "data", features=features, episodes=self.episodes) hf_dataset.set_transform(hf_transform_to_torch) return hf_dataset + def _validate_language_columns_declared(self, features: datasets.Features) -> None: + """Require language columns stored in Parquet to be declared in metadata.""" + # Leave empty datasets to fail through the normal loading path. + try: + sample = next((self.root / "data").glob("*/*.parquet")) + except StopIteration: + return + + from pyarrow import parquet as _pq # noqa: PLC0415 + + # LeRobot shards are schema-uniform, so one schema represents the dataset. + schema_names = set(_pq.read_schema(sample).names) + from .language import LANGUAGE_COLUMNS # noqa: PLC0415 + + missing = sorted(set(LANGUAGE_COLUMNS) & schema_names - set(features)) + if missing: + raise ValueError( + f"Dataset Parquet files contain language feature(s) missing from metadata: {missing}. " + "Metadata must describe the stored data; add the entries returned by " + "lerobot.datasets.language.language_feature_info() to meta/info.json['features'] " + "or rerun the annotation pipeline's metadata synchronization." + ) + def _check_cached_episodes_sufficient(self) -> bool: """Check if the cached dataset contains all requested episodes and their video files.""" if self.hf_dataset is None or len(self.hf_dataset) == 0: diff --git a/src/lerobot/datasets/factory.py b/src/lerobot/datasets/factory.py index a727bf924..0388c2caa 100644 --- a/src/lerobot/datasets/factory.py +++ b/src/lerobot/datasets/factory.py @@ -29,6 +29,7 @@ from .dataset_metadata import LeRobotDatasetMetadata from .lerobot_dataset import LeRobotDataset from .multi_dataset import MultiLeRobotDataset from .streaming_dataset import StreamingLeRobotDataset +from .utils import resolve_episode_indices def resolve_delta_timestamps( @@ -90,6 +91,9 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas repo_type=cfg.dataset.repo_type, ) delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, ds_meta) + episodes = resolve_episode_indices( + cfg.dataset.episodes, ds_meta.total_episodes, cfg.dataset.exclude_episodes + ) if not cfg.dataset.streaming: if cfg.dataset.repo_type == "bucket": raise ValueError( @@ -98,7 +102,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas dataset = LeRobotDataset( cfg.dataset.repo_id, root=cfg.dataset.root, - episodes=cfg.dataset.episodes, + episodes=episodes, delta_timestamps=delta_timestamps, image_transforms=image_transforms, revision=cfg.dataset.revision, @@ -111,7 +115,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas dataset = StreamingLeRobotDataset( cfg.dataset.repo_id, root=cfg.dataset.root, - episodes=cfg.dataset.episodes, + episodes=episodes, delta_timestamps=delta_timestamps, image_transforms=image_transforms, revision=cfg.dataset.revision, diff --git a/src/lerobot/datasets/language_render.py b/src/lerobot/datasets/language_render.py index 999fa19ad..4d4afa58f 100644 --- a/src/lerobot/datasets/language_render.py +++ b/src/lerobot/datasets/language_render.py @@ -162,14 +162,32 @@ def render_sample( task: str | None = None, dataset_ctx: Any | None = None, ) -> RenderedMessages | None: - """Render the chat-style messages for a single dataset sample. + """Render recipe-defined messages and supervision for one dataset sample. - Resolves the recipe's bindings against ``persistent`` and ``events`` rows - at frame timestamp ``t``, then expands the recipe's message templates. - Returns ``None`` if the resolved sample contains no target message. + Resolves bindings against ``persistent`` and ``events`` at frame timestamp + ``t``. Blend recipes first route matching sparse VQA annotations, then use + deterministic weighted selection for the remaining samples. Returns + ``None`` when the selected recipe provides no text or low-level action + supervision for this sample. """ persistent_rows = _normalize_rows(persistent or []) event_rows = _normalize_rows(events or []) + + # Route sparse VQA frames to a matching view-specific component before weighted selection. + # This avoids dropping annotated frames or selecting VQA without annotations. + if recipe.blend is not None: + vqa_rendered = _render_vqa_if_present( + recipe, + persistent=persistent_rows, + events=event_rows, + t=t, + sample_idx=sample_idx, + task=task, + dataset_ctx=dataset_ctx, + ) + if vqa_rendered is not None: + return vqa_rendered + selected_recipe = _select_recipe(recipe, sample_idx) bindings = _resolve_bindings( selected_recipe, @@ -183,6 +201,58 @@ def render_sample( return _render_message_recipe(selected_recipe, bindings) +def _render_vqa_if_present( + recipe: TrainingRecipe, + *, + persistent: Sequence[LanguageRow], + events: Sequence[LanguageRow], + t: float, + sample_idx: int, + task: str | None, + dataset_ctx: Any | None, +) -> RenderedMessages | None: + """Render a matching VQA component, or return ``None`` for normal selection. + + Multiple matching views are selected deterministically by relative weight. + """ + if recipe.blend is None: + return None + renderable: list[tuple[float, RenderedMessages]] = [] + for component in recipe.blend.values(): + if component.route != "vqa": + continue + bindings = _resolve_bindings( + component, + persistent=persistent, + events=events, + t=t, + sample_idx=sample_idx, + task=task, + dataset_ctx=dataset_ctx, + ) + rendered = _render_message_recipe(component, bindings) + if rendered is not None: + if component.weight is None: + raise ValueError("Routed VQA blend components must define a weight.") + renderable.append((component.weight, rendered)) + + if not renderable: + return None + if len(renderable) == 1: + return renderable[0][1] + + # Choose among matching cameras by their validated positive relative weights. + total = sum(weight for weight, _ in renderable) + digest = hashlib.blake2b(f"vqa:{sample_idx}".encode(), digest_size=8).digest() + draw = int.from_bytes(digest, "big") / 2**64 * total + cumulative = 0.0 + for weight, rendered in renderable: + cumulative += weight + if draw < cumulative: + return rendered + return renderable[-1][1] + + def _select_recipe(recipe: TrainingRecipe, sample_idx: int) -> TrainingRecipe: """Pick a deterministic blend component for ``sample_idx`` (or return ``recipe``).""" if recipe.blend is None: @@ -201,7 +271,8 @@ def _select_recipe(recipe: TrainingRecipe, sample_idx: int) -> TrainingRecipe: cumulative += component.weight or 0.0 if draw < cumulative: return component - assert last_component is not None + if last_component is None: + raise ValueError("Blend recipes must contain at least one component.") return last_component @@ -321,7 +392,8 @@ def _render_message_recipe( bindings: dict[str, LanguageRow | str | None], ) -> RenderedMessages | None: """Expand ``recipe.messages`` into rendered chat messages using ``bindings``.""" - assert recipe.messages is not None + if recipe.messages is None: + raise ValueError("Cannot render a blend recipe as a message recipe.") messages: list[dict[str, Any]] = [] streams: list[str | None] = [] target_indices: list[int] = [] @@ -346,7 +418,9 @@ def _render_message_recipe( if turn.target: target_indices.append(message_idx) - if not target_indices: + # Keep samples with either text targets or low-level action supervision. + has_low_level = any(stream == "low_level" for stream in streams) + if not target_indices and not has_low_level: return None rendered = { @@ -403,14 +477,12 @@ def _validate_rendered(rendered: RenderedMessages) -> None: if len(streams) != len(messages): raise ValueError("message_streams must be aligned with messages.") - if not target_indices: - raise ValueError("Rendered samples must contain at least one target message.") + # Require text or low-level action supervision. + if not target_indices and not any(s == "low_level" for s in streams): + raise ValueError("Rendered samples must contain a target message or a low_level-stream message.") for idx in target_indices: if idx < 0 or idx >= len(messages): raise ValueError(f"Target message index {idx} is out of bounds.") - # ``stream`` is enforced non-None at MessageTurn construction time - # (see ``MessageTurn.__post_init__``), so a missing stream here would - # mean the dataclass invariant was bypassed; no need to re-check. def _nth_relative( diff --git a/src/lerobot/datasets/utils.py b/src/lerobot/datasets/utils.py index 9fde26067..2ab0c2d6d 100644 --- a/src/lerobot/datasets/utils.py +++ b/src/lerobot/datasets/utils.py @@ -18,6 +18,7 @@ import dataclasses import importlib.resources import json import logging +from collections.abc import Sequence from dataclasses import dataclass, field from pathlib import Path @@ -98,6 +99,47 @@ VIDEO_DIR = "videos" CHUNK_FILE_PATTERN = "chunk-{chunk_index:03d}/file-{file_index:03d}" IMAGE_FILE_PATTERN = "frame-{frame_index:06d}.png" + + +def resolve_episode_indices( + episodes: Sequence[int] | None, + total_episodes: int, + exclude_episodes: Sequence[int] | None = None, +) -> list[int] | None: + """Resolve an optional episode allowlist and exclusion list against dataset bounds. + + ``None`` is preserved when no filtering is requested so callers can retain + their native "all episodes" fast path. Invalid indices are ignored with a + warning, and the input order is preserved. + """ + if total_episodes < 0: + raise ValueError(f"total_episodes must be non-negative, got {total_episodes}") + + if episodes is None and not exclude_episodes: + return None + + candidates = list(range(total_episodes)) if episodes is None else list(episodes) + invalid = [episode for episode in candidates if not 0 <= episode < total_episodes] + if invalid: + logger.warning( + "Ignoring episode indices outside the dataset range [0, %d): %s", + total_episodes, + invalid, + ) + candidates = [episode for episode in candidates if 0 <= episode < total_episodes] + + excluded = set(exclude_episodes or []) + invalid_excluded = sorted(episode for episode in excluded if not 0 <= episode < total_episodes) + if invalid_excluded: + logger.warning( + "Ignoring excluded episode indices outside the dataset range [0, %d): %s", + total_episodes, + invalid_excluded, + ) + excluded = {episode for episode in excluded if 0 <= episode < total_episodes} + return [episode for episode in candidates if episode not in excluded] + + DEPTH_FILE_PATTERN = "frame-{frame_index:06d}.tiff" DEFAULT_TASKS_PATH = "meta/tasks.parquet" DEFAULT_EPISODES_PATH = EPISODES_DIR + "/" + CHUNK_FILE_PATTERN + ".parquet" diff --git a/src/lerobot/processor/batch_processor.py b/src/lerobot/processor/batch_processor.py index 27742f3b2..b91ce77e0 100644 --- a/src/lerobot/processor/batch_processor.py +++ b/src/lerobot/processor/batch_processor.py @@ -175,9 +175,6 @@ class AddBatchDimensionComplementaryDataStep(ComplementaryDataProcessorStep): if isinstance(task_index_value, Tensor) and task_index_value.dim() == 0: complementary_data["task_index"] = task_index_value.unsqueeze(0) - complementary_data.pop("language_persistent", None) - complementary_data.pop("language_events", None) - if "messages" in complementary_data: messages = complementary_data["messages"] if isinstance(messages, list) and (not messages or isinstance(messages[0], dict)): diff --git a/src/lerobot/processor/pipeline.py b/src/lerobot/processor/pipeline.py index 54036f986..0c0534572 100644 --- a/src/lerobot/processor/pipeline.py +++ b/src/lerobot/processor/pipeline.py @@ -41,7 +41,7 @@ from pathlib import Path from typing import Any, TypedDict, TypeVar, cast import torch -from huggingface_hub import hf_hub_download +from huggingface_hub import hf_hub_download, snapshot_download from safetensors.torch import load_file, save_file from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature @@ -212,6 +212,10 @@ class ProcessorStep(ABC): """ return None + def save_artifacts(self, save_directory: Path) -> dict[str, str]: + """Save non-tensor assets and map constructor arguments to relative paths.""" + return {} + def reset(self) -> None: """Resets the internal state of the processor step, if any.""" return None @@ -556,6 +560,22 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): pipeline_config = self.get_config() pipeline_state_dict = self.state_dict() + for processor_step, step_entry in zip(self.steps, pipeline_config["steps"], strict=True): + artifacts = processor_step.save_artifacts(save_directory) + if artifacts: + for config_key, relative_path in artifacts.items(): + artifact_path = Path(relative_path) + if artifact_path.is_absolute() or ".." in artifact_path.parts: + raise ValueError( + f"Processor artifact path must be relative to the checkpoint: {relative_path!r}" + ) + if not (save_directory / artifact_path).exists(): + raise FileNotFoundError( + f"Processor step did not save declared artifact '{relative_path}'" + ) + step_entry["config"][config_key] = artifact_path.as_posix() + step_entry["artifacts"] = artifacts + for state_key, step_state_dict in pipeline_state_dict.items(): state_filename = f"{state_key}.safetensors" save_file(step_state_dict, save_directory / state_filename) @@ -740,7 +760,13 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): # 3. Build steps with overrides steps, validated_overrides = cls._build_steps_with_overrides( - loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs, is_local_source + loaded_config, + overrides or {}, + model_id, + base_path, + config_filename, + hub_download_kwargs, + is_local_source, ) # 4. Validate that all overrides were used @@ -936,6 +962,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): overrides: dict[str, Any], model_id: str, base_path: Path | None, + config_filename: str, hub_download_kwargs: dict[str, Any], is_local_source: bool = False, ) -> tuple[list[ProcessorStep], set[str]]: @@ -945,6 +972,11 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): **For each step in loaded_config["steps"]**: + 0. **Artifact Resolution** (via _resolve_artifact_paths): + - Resolve declared relative artifact paths against a local checkpoint + - Download declared artifacts when loading the pipeline from the Hub + - Reject absolute paths and path traversal before step construction + 1. **Class Resolution** (via _resolve_step_class): - **If "registry_name" exists**: Look up in ProcessorStepRegistry Example: {"registry_name": "normalize_step"} -> Get registered class @@ -978,6 +1010,8 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): overrides: User-provided parameter overrides (keyed by class/registry name) model_id: The model identifier (needed for Hub state file downloads) base_path: Local directory path for finding state files + config_filename: Processor config path, used as the repository-relative + base for state files and declared artifacts. hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.) is_local_source: Whether model_id resolved to a local directory or config file. @@ -990,15 +1024,80 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): ImportError: If a step class cannot be imported or found in registry ValueError: If a step cannot be instantiated with its configuration """ + loaded_config = deepcopy(loaded_config) + cls._resolve_artifact_paths( + loaded_config, + model_id, + base_path, + config_filename, + hub_download_kwargs, + ) steps, remaining_override_keys = cls._build_steps_from_config(loaded_config, overrides) for step_instance, step_entry in zip(steps, loaded_config["steps"], strict=True): cls._load_step_state( - step_instance, step_entry, model_id, base_path, hub_download_kwargs, is_local_source + step_instance, + step_entry, + model_id, + base_path, + config_filename, + hub_download_kwargs, + is_local_source, ) return steps, remaining_override_keys + @classmethod + def _resolve_artifact_paths( + cls, + loaded_config: dict[str, Any], + model_id: str, + base_path: Path | None, + config_filename: str, + hub_download_kwargs: dict[str, Any], + ) -> None: + """Resolve declared relative processor artifacts before step construction. + + Args: + loaded_config: Mutable processor configuration containing step artifact declarations. + model_id: Local checkpoint path or Hub model identifier. + base_path: Local directory containing the resolved processor configuration. + config_filename: Processor config path, whose parent is the artifact root on the Hub. + hub_download_kwargs: Authentication, revision, and cache arguments for Hub downloads. + + Raises: + ValueError: If a declared artifact path is absolute or escapes the checkpoint. + FileNotFoundError: If a declared artifact cannot be found locally or downloaded. + """ + is_local = Path(model_id).is_dir() or Path(model_id).is_file() + + for step_entry in loaded_config["steps"]: + artifacts = step_entry.get("artifacts", {}) + for config_key, relative_path in artifacts.items(): + artifact_path = Path(relative_path) + if artifact_path.is_absolute() or ".." in artifact_path.parts: + raise ValueError( + f"Processor artifact path must be relative to the checkpoint: {relative_path!r}" + ) + + resolved_path = base_path / artifact_path if base_path is not None else artifact_path + if not resolved_path.exists() and not is_local: + repository_path = Path(config_filename).parent / artifact_path + snapshot_download( + repo_id=model_id, + repo_type="model", + allow_patterns=f"{repository_path.as_posix()}/**", + **hub_download_kwargs, + ) + + if not resolved_path.exists(): + step_name = step_entry.get("registry_name", step_entry.get("class", "unknown")) + raise FileNotFoundError( + f"Missing processor artifact '{relative_path}' for step '{step_name}' " + f"next to '{config_filename}'. Checkpoint artifacts are incomplete." + ) + step_entry["config"][config_key] = str(resolved_path) + @classmethod def _build_steps_from_config( cls, @@ -1158,6 +1257,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): step_entry: dict[str, Any], model_id: str, base_path: Path | None, + config_filename: str, hub_download_kwargs: dict[str, Any], is_local_source: bool = False, ) -> None: @@ -1198,6 +1298,8 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): step_entry: The step configuration dictionary (may contain "state_file") model_id: The model identifier (used for Hub downloads if needed) base_path: Local directory path for finding state files (None for Hub-only) + config_filename: Processor config path, whose parent is used to resolve + repository-relative state files on the Hub. hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.) is_local_source: Whether model_id resolved to a local directory or config file. @@ -1223,7 +1325,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): # Download from Hub state_path = hf_hub_download( repo_id=model_id, - filename=state_filename, + filename=(Path(config_filename).parent / state_filename).as_posix(), repo_type="model", **hub_download_kwargs, ) diff --git a/src/lerobot/processor/render_messages_processor.py b/src/lerobot/processor/render_messages_processor.py index 92ba9b8d9..4fe8128b8 100644 --- a/src/lerobot/processor/render_messages_processor.py +++ b/src/lerobot/processor/render_messages_processor.py @@ -16,9 +16,11 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import asdict, dataclass from typing import Any +import numpy as np + from lerobot.configs import PipelineFeatureType, PolicyFeature from lerobot.configs.recipe import TrainingRecipe from lerobot.datasets.language import LANGUAGE_EVENTS, LANGUAGE_PERSISTENT @@ -32,25 +34,46 @@ from .pipeline import ProcessorStep, ProcessorStepRegistry @dataclass @ProcessorStepRegistry.register(name="render_messages_processor") class RenderMessagesStep(ProcessorStep): - """Processor step that turns raw language columns into rendered chat messages. + """Turn raw language columns into recipe-defined messages and supervision. - Reads ``language_persistent`` and ``language_events`` from the transition's - complementary data, renders them through ``recipe`` at the sample timestamp, - and replaces the raw columns with the resulting ``messages`` / - ``message_streams`` / ``target_message_indices`` keys. + Reads ``language_persistent`` and ``language_events`` from complementary + data, renders them at each sample timestamp, and replaces the raw columns + with ``messages``, ``message_streams``, and ``target_message_indices``. + Batched inputs are filtered to samples with applicable supervision; samples + without language annotations use their task string as low-level supervision + when one is available. """ recipe: TrainingRecipe dataset_ctx: Any | None = None + def __post_init__(self) -> None: + if isinstance(self.recipe, dict): + self.recipe = TrainingRecipe.from_dict(self.recipe) + + def get_config(self) -> dict[str, Any]: + return {"recipe": asdict(self.recipe)} + def __call__(self, transition: EnvTransition) -> EnvTransition | None: - """Render messages for a single transition; return ``None`` to drop it.""" + """Render messages, preserving unannotated samples and dropping unmatched annotated ones.""" complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA) or {} persistent = complementary_data.get(LANGUAGE_PERSISTENT) or [] events = complementary_data.get(LANGUAGE_EVENTS) or [] if not persistent and not events: - return transition + # A dataset without language annotations remains usable: render its + # task as low-level supervision, or pass it through when no task exists. + rendered = _fallback_low_level_render(complementary_data.get("task")) + if rendered is None: + return transition + new_transition = transition.copy() + new_complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}) + new_complementary_data.update(rendered) + new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data + return new_transition + + if _is_batched_language(persistent) or _is_batched_language(events): + return self._call_batch(transition, complementary_data, persistent, events) timestamp = complementary_data.get("timestamp") if timestamp is None: @@ -67,18 +90,171 @@ class RenderMessagesStep(ProcessorStep): dataset_ctx=self.dataset_ctx, ) if rendered is None: - return None + # Language is present but this sparse frame has no applicable recipe + # branch. Keep it only when task-level action supervision is possible. + rendered = _fallback_low_level_render(complementary_data.get("task")) + if rendered is None: + return None new_transition = transition.copy() - new_complementary_data = dict(complementary_data) + new_complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}) new_complementary_data.pop(LANGUAGE_PERSISTENT, None) new_complementary_data.pop(LANGUAGE_EVENTS, None) new_complementary_data.update(rendered) new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data return new_transition + def _call_batch( + self, + transition: EnvTransition, + complementary_data: dict[str, Any], + persistent_batch: list, + events_batch: list, + ) -> EnvTransition | None: + """Render a language batch. + + Non-empty persistent and event batches must have the same size. Either + list may be empty when that language column is absent from the batch. + """ + timestamp = complementary_data.get("timestamp") + if timestamp is None: + raise KeyError("RenderMessagesStep requires sample timestamp in complementary data.") + + non_empty_batch_sizes = {len(batch) for batch in (persistent_batch, events_batch) if batch} + if len(non_empty_batch_sizes) > 1: + raise ValueError( + "Batched language columns must have equal lengths when both are non-empty, " + f"got persistent={len(persistent_batch)} and events={len(events_batch)}." + ) + batch_size = next(iter(non_empty_batch_sizes), 0) + messages: list[list[dict[str, Any]]] = [] + message_streams: list[list[str | None]] = [] + target_message_indices: list[list[int]] = [] + keep_indices: list[int] = [] + + for i in range(batch_size): + rendered = render_sample( + recipe=self.recipe, + persistent=persistent_batch[i] if i < len(persistent_batch) else [], + events=events_batch[i] if i < len(events_batch) else [], + t=_batch_value(timestamp, i), + sample_idx=int(_batch_value(complementary_data.get("index", 0), i)), + task=_batch_value(complementary_data.get("task"), i), + dataset_ctx=self.dataset_ctx, + ) + if rendered is None: + rendered = _fallback_low_level_render(_batch_value(complementary_data.get("task"), i)) + if rendered is None: + continue + keep_indices.append(i) + messages.append(rendered["messages"]) + message_streams.append(rendered["message_streams"]) + target_message_indices.append(rendered["target_message_indices"]) + + if not messages: + return None + + new_transition = ( + _select_batch_indices(transition, keep_indices, batch_size) + if len(keep_indices) != batch_size + else transition.copy() + ) + new_complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}) + new_complementary_data.pop(LANGUAGE_PERSISTENT, None) + new_complementary_data.pop(LANGUAGE_EVENTS, None) + new_complementary_data["messages"] = messages + new_complementary_data["message_streams"] = message_streams + new_complementary_data["target_message_indices"] = target_message_indices + new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data + return new_transition + def transform_features( self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: """Pass features through unchanged; rendering only touches complementary data.""" return features + + +def _is_batched_language(value: Any) -> bool: + return isinstance(value, list) and bool(value) and isinstance(value[0], list) + + +def _batch_value(value: Any, index: int) -> Any: + if value is None: + return None + if isinstance(value, list): + return value[index] + if hasattr(value, "ndim") and value.ndim > 0: + return unwrap_scalar(value[index]) + return unwrap_scalar(value) + + +def _select_batch_indices(transition: EnvTransition, indices: list[int], batch_size: int) -> EnvTransition: + selected = transition.copy() + for key in (TransitionKey.OBSERVATION, TransitionKey.COMPLEMENTARY_DATA): + data = selected.get(key) + if isinstance(data, dict): + selected[key] = { + name: _select_value(value, indices, batch_size, f"{key}.{name}") + for name, value in data.items() + } + action = selected.get(TransitionKey.ACTION) + if action is not None: + selected[TransitionKey.ACTION] = _select_value(action, indices, batch_size, str(TransitionKey.ACTION)) + return selected + + +def _select_value(value: Any, indices: list[int], batch_size: int, path: str) -> Any: + if isinstance(value, dict): + return {key: _select_value(item, indices, batch_size, f"{path}.{key}") for key, item in value.items()} + if isinstance(value, list): + if len(value) != batch_size: + raise ValueError( + f"Cannot filter batched field {path!r}: expected {batch_size} values, got {len(value)}." + ) + return [value[i] for i in indices] + if isinstance(value, np.ndarray) and value.ndim > 0: + return value[indices] + if hasattr(value, "index_select") and hasattr(value, "new_tensor") and getattr(value, "ndim", 0) > 0: + return value.index_select(0, value.new_tensor(indices).long()) + return value + + +def _fallback_low_level_render(task: Any) -> dict[str, Any] | None: + """Keep action-only samples trainable when no recipe branch matches.""" + if hasattr(task, "item"): + task = task.item() + if isinstance(task, list): + if not task: + return None + messages = [] + message_streams = [] + target_message_indices = [] + missing_indices = [] + for index, t in enumerate(task): + rendered = _fallback_low_level_render(t) + if rendered is None: + missing_indices.append(index) + continue + messages.append(rendered["messages"]) + message_streams.append(rendered["message_streams"]) + target_message_indices.append(rendered["target_message_indices"]) + if missing_indices: + if len(missing_indices) == len(task): + return None + raise ValueError( + "Batched low-level fallback requires a non-empty task for every sample; " + f"missing task at indices {missing_indices}." + ) + return { + "messages": messages, + "message_streams": message_streams, + "target_message_indices": target_message_indices, + } + if not isinstance(task, str) or not task: + return None + return { + "messages": [{"role": "user", "content": task}], + "message_streams": ["low_level"], + "target_message_indices": [], + } diff --git a/src/lerobot/processor/tokenizer_processor.py b/src/lerobot/processor/tokenizer_processor.py index 967159144..ae226b6fb 100644 --- a/src/lerobot/processor/tokenizer_processor.py +++ b/src/lerobot/processor/tokenizer_processor.py @@ -25,6 +25,7 @@ from __future__ import annotations import logging from dataclasses import dataclass, field +from pathlib import Path from typing import TYPE_CHECKING, Any import torch @@ -32,6 +33,7 @@ import torch from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature from lerobot.lerobot_types import EnvTransition, RobotObservation, TransitionKey from lerobot.utils.constants import ( + ACTION_CODE_TOKEN_MASK, ACTION_TOKEN_MASK, ACTION_TOKENS, OBS_LANGUAGE_ATTENTION_MASK, @@ -136,7 +138,7 @@ class TokenizerProcessorStep(ObservationProcessorStep): # Standardize to a list of strings for the tokenizer if isinstance(task, str): return [task] - elif isinstance(task, (list, tuple)) and all(isinstance(t, str) for t in task): + elif isinstance(task, list | tuple) and all(isinstance(t, str) for t in task): return list(task) return None @@ -293,6 +295,15 @@ class TokenizerProcessorStep(ObservationProcessorStep): return config + def save_artifacts(self, save_directory: Path) -> dict[str, str]: + """Save the tokenizer so object-provided instances reload without overrides.""" + artifact_path = Path("tokenizer") + save_pretrained = getattr(self.input_tokenizer, "save_pretrained", None) + if save_pretrained is None: + raise TypeError("Tokenizer must implement save_pretrained() to save a portable pipeline.") + save_pretrained(save_directory / artifact_path) + return {"tokenizer_name": artifact_path.as_posix()} + def transform_features( self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: @@ -349,6 +360,7 @@ class ActionTokenizerProcessorStep(ActionProcessorStep): max_action_tokens: int = 256 fast_skip_tokens: int = 128 paligemma_tokenizer_name: str = "google/paligemma-3b-pt-224" + allow_truncation: bool = True # Internal tokenizer instance (not part of the config) action_tokenizer: Any = field(default=None, init=False, repr=False) _paligemma_tokenizer: Any = field(default=None, init=False, repr=False) @@ -412,14 +424,15 @@ class ActionTokenizerProcessorStep(ActionProcessorStep): # During inference, no action is available, skip tokenization return new_transition - # Tokenize and get both tokens and mask - tokens, mask = self._tokenize_action(action) + # Tokenize and get masks for the full formatted sequence and the discrete action codes. + tokens, mask, code_mask = self._tokenize_action(action) # Store mask in complementary data complementary_data = new_transition.get(TransitionKey.COMPLEMENTARY_DATA, {}) if complementary_data is None: complementary_data = {} complementary_data[ACTION_TOKEN_MASK] = mask + complementary_data[ACTION_CODE_TOKEN_MASK] = code_mask complementary_data[ACTION_TOKENS] = tokens new_transition[TransitionKey.COMPLEMENTARY_DATA] = complementary_data return new_transition @@ -430,7 +443,7 @@ class ActionTokenizerProcessorStep(ActionProcessorStep): """ return self._paligemma_tokenizer.vocab_size - 1 - self.fast_skip_tokens - tokens - def _tokenize_action(self, action: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def _tokenize_action(self, action: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Tokenizes the action tensor and creates a mask. @@ -459,6 +472,7 @@ class ActionTokenizerProcessorStep(ActionProcessorStep): # The fast tokenizer expects action data and returns token IDs tokens_list = [] masks_list = [] + code_masks_list = [] for i in range(batch_size): # Tokenize single action (move to CPU first as tokenizer uses scipy which requires numpy) @@ -476,65 +490,82 @@ class ActionTokenizerProcessorStep(ActionProcessorStep): if tokens.dim() > 1: tokens = tokens.flatten() + action_code_tokens = self._act_tokens_to_paligemma_tokens(tokens) bos_id = self._paligemma_tokenizer.bos_token_id - # add bos + prompt_tokens = torch.tensor( + self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False), + device=action.device, + ) + end_tokens = torch.tensor(self._paligemma_tokenizer.encode("|"), device=action.device) + + code_start = 1 + len(prompt_tokens) + code_end = code_start + len(action_code_tokens) tokens = torch.cat( [ torch.tensor([bos_id], device=action.device), - torch.tensor( - self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False), - device=action.device, - ), - self._act_tokens_to_paligemma_tokens(tokens), - torch.tensor(self._paligemma_tokenizer.encode("|"), device=action.device), + prompt_tokens, + action_code_tokens, + end_tokens, ] ) + code_mask = torch.zeros(len(tokens), dtype=torch.bool, device=action.device) + code_mask[code_start:code_end] = True # Truncate or pad to max_action_tokens if len(tokens) > self.max_action_tokens: + if not self.allow_truncation: + raise ValueError( + f"FAST action sequence has {len(tokens)} tokens, exceeding " + f"max_action_tokens={self.max_action_tokens}." + ) logging.warning( f"Token length ({len(tokens)}) exceeds max length ({self.max_action_tokens}), truncating. " "Consider increasing the `max_action_tokens` in your model config if this happens frequently." ) tokens = tokens[: self.max_action_tokens] + code_mask = code_mask[: self.max_action_tokens] mask = torch.ones(self.max_action_tokens, dtype=torch.bool, device=action.device) else: + pad_len = self.max_action_tokens - len(tokens) mask = torch.cat( [ torch.ones(len(tokens), dtype=torch.bool, device=action.device), - torch.zeros( - self.max_action_tokens - len(tokens), dtype=torch.bool, device=action.device - ), + torch.zeros(pad_len, dtype=torch.bool, device=action.device), ] ) + code_mask = torch.nn.functional.pad(code_mask, (0, pad_len), value=False) # Pad tokens with zeros - tokens = torch.nn.functional.pad(tokens, (0, self.max_action_tokens - len(tokens)), value=0) + tokens = torch.nn.functional.pad(tokens, (0, pad_len), value=0) tokens_list.append(tokens) masks_list.append(mask) + code_masks_list.append(code_mask) # Stack into batched tensors tokens_batch = torch.stack(tokens_list, dim=0) # (B, max_action_tokens) masks_batch = torch.stack(masks_list, dim=0) # (B, max_action_tokens) + code_masks_batch = torch.stack(code_masks_list, dim=0) # (B, max_action_tokens) # Remove batch dimension if input was single sample if single_sample: tokens_batch = tokens_batch.squeeze(0) masks_batch = masks_batch.squeeze(0) + code_masks_batch = code_masks_batch.squeeze(0) # Move to the same device as the input if device is not None: tokens_batch = tokens_batch.to(device) masks_batch = masks_batch.to(device) + code_masks_batch = code_masks_batch.to(device) - return tokens_batch, masks_batch + return tokens_batch, masks_batch, code_masks_batch def action(self, action: torch.Tensor) -> torch.Tensor: """ This method is not used since we override __call__. Required by ActionProcessorStep ABC. """ - tokens, _ = self._tokenize_action(action) + tokens, _, _ = self._tokenize_action(action) return tokens def get_config(self) -> dict[str, Any]: @@ -550,6 +581,9 @@ class ActionTokenizerProcessorStep(ActionProcessorStep): config = { "trust_remote_code": self.trust_remote_code, "max_action_tokens": self.max_action_tokens, + "fast_skip_tokens": self.fast_skip_tokens, + "paligemma_tokenizer_name": self.paligemma_tokenizer_name, + "allow_truncation": self.allow_truncation, } # Only save tokenizer_name if it was used to create the tokenizer @@ -558,6 +592,14 @@ class ActionTokenizerProcessorStep(ActionProcessorStep): return config + def save_artifacts(self, save_directory: Path) -> dict[str, str]: + artifact_path = Path("action_tokenizer") + save_pretrained = getattr(self.action_tokenizer, "save_pretrained", None) + if save_pretrained is None: + raise TypeError("Action tokenizer must implement save_pretrained() to save a portable pipeline.") + save_pretrained(save_directory / artifact_path) + return {"action_tokenizer_name": artifact_path.as_posix()} + def transform_features( self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: diff --git a/src/lerobot/utils/collate.py b/src/lerobot/utils/collate.py index fce7e6b42..d159a9b23 100644 --- a/src/lerobot/utils/collate.py +++ b/src/lerobot/utils/collate.py @@ -22,7 +22,7 @@ from torch.utils.data._utils.collate import default_collate from lerobot.datasets.language import LANGUAGE_COLUMNS -_PYTHON_LIST_KEYS = {"messages", "message_streams", "target_message_indices"} +_PYTHON_LIST_KEYS = {"messages", "message_streams", "target_message_indices", *LANGUAGE_COLUMNS} def lerobot_collate_fn(batch: list[dict[str, Any] | None]) -> dict[str, Any] | None: diff --git a/src/lerobot/utils/constants.py b/src/lerobot/utils/constants.py index 8f735fe6d..fae865b16 100644 --- a/src/lerobot/utils/constants.py +++ b/src/lerobot/utils/constants.py @@ -26,6 +26,7 @@ OBS_IMAGES = OBS_IMAGE + "s" OBS_LANGUAGE = OBS_STR + ".language" OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens" OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask" +OBS_LANGUAGE_CAUSAL_MARKS = OBS_LANGUAGE + ".causal_marks" OBS_LANGUAGE_SUBTASK = OBS_STR + ".subtask" OBS_LANGUAGE_SUBTASK_TOKENS = OBS_LANGUAGE_SUBTASK + ".tokens" OBS_LANGUAGE_SUBTASK_ATTENTION_MASK = OBS_LANGUAGE_SUBTASK + ".attention_mask" @@ -34,6 +35,7 @@ ACTION = "action" ACTION_PREFIX = ACTION + "." ACTION_TOKENS = ACTION + ".tokens" ACTION_TOKEN_MASK = ACTION + ".token_mask" +ACTION_CODE_TOKEN_MASK = ACTION + ".code_token_mask" REWARD = "next.reward" TRUNCATED = "next.truncated" DONE = "next.done" diff --git a/tests/configs/test_default.py b/tests/configs/test_default.py index 979509ec3..e41a718f8 100644 --- a/tests/configs/test_default.py +++ b/tests/configs/test_default.py @@ -38,6 +38,13 @@ def test_dataset_config_empty_episodes_ok(): DatasetConfig(repo_id="user/repo", episodes=[]) +def test_dataset_config_ignores_negative_excluded_episodes(caplog): + config = DatasetConfig(repo_id="user/repo", exclude_episodes=[-2, 1, -1, 3]) + + assert config.exclude_episodes == [1, 3] + assert "Ignoring negative exclude_episodes entries: [-2, -1]" in caplog.text + + def test_dataset_config_bucket_streaming_ok(): DatasetConfig(repo_id="user/repo", repo_type="bucket", streaming=True) diff --git a/tests/configs/test_recipe.py b/tests/configs/test_recipe.py index b4954efbf..0d5de798d 100644 --- a/tests/configs/test_recipe.py +++ b/tests/configs/test_recipe.py @@ -29,6 +29,13 @@ def test_message_recipe_validates_unknown_binding(): ) +def test_canonical_recipe_loads(): + """The canonical PI052 blend YAML loads + validates.""" + recipe = TrainingRecipe.from_yaml(Path("src/lerobot/configs/recipes/subtask_mem_vqa_speech.yaml")) + assert recipe.blend is not None + assert sum(c.weight for c in recipe.blend.values()) == pytest.approx(1.0) + + def test_message_turn_requires_a_stream(): """Every turn must declare a stream — None is rejected at construction. @@ -81,6 +88,19 @@ def test_blend_component_weight_must_be_positive(): TrainingRecipe(blend={"a": TrainingRecipe(weight=0.0, messages=[_minimal_target_turn()])}) +def test_recipe_route_must_be_supported(): + with pytest.raises(ValueError, match="Unsupported recipe route"): + TrainingRecipe(weight=1.0, route="other", messages=[_minimal_target_turn()]) + + +def test_route_cannot_be_set_on_blend_recipe(): + with pytest.raises(ValueError, match="only be set on a message recipe"): + TrainingRecipe( + route="vqa", + blend={"a": TrainingRecipe(weight=1.0, messages=[_minimal_target_turn()])}, + ) + + def test_blend_component_must_define_messages(): # A bare TrainingRecipe(weight=1.0) would itself raise; build it without # going through __post_init__ to exercise the blend-level validator. diff --git a/tests/datasets/test_dataset_reader.py b/tests/datasets/test_dataset_reader.py index 085563bb8..fa031e2d3 100644 --- a/tests/datasets/test_dataset_reader.py +++ b/tests/datasets/test_dataset_reader.py @@ -20,6 +20,7 @@ import pytest pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") from lerobot.datasets.dataset_reader import DatasetReader +from lerobot.datasets.language import LANGUAGE_EVENTS from lerobot.utils.import_utils import get_safe_default_video_backend # ── Loading ────────────────────────────────────────────────────────── @@ -66,6 +67,22 @@ def test_try_load_returns_false_when_no_data(tmp_path): assert reader.hf_dataset is None +def test_load_rejects_language_columns_missing_from_metadata(tmp_path, lerobot_dataset_factory): + import pyarrow as pa + import pyarrow.parquet as pq + + dataset = lerobot_dataset_factory( + root=tmp_path / "ds", total_episodes=1, total_frames=10, use_videos=False + ) + parquet_path = next((dataset.root / "data").glob("*/*.parquet")) + table = pq.read_table(parquet_path) + language_events = pa.array([[] for _ in range(len(table))], type=pa.list_(pa.string())) + pq.write_table(table.append_column(LANGUAGE_EVENTS, language_events), parquet_path) + + with pytest.raises(ValueError, match=r"language feature\(s\) missing from metadata.*language_events"): + dataset.reader.load_and_activate() + + # ── Counts ─────────────────────────────────────────────────────────── diff --git a/tests/datasets/test_dataset_utils.py b/tests/datasets/test_dataset_utils.py index e8a3fc5fe..932e64df4 100644 --- a/tests/datasets/test_dataset_utils.py +++ b/tests/datasets/test_dataset_utils.py @@ -28,7 +28,12 @@ from huggingface_hub import DatasetCard import lerobot.datasets.utils as dataset_utils from lerobot.datasets.io_utils import hf_transform_to_torch -from lerobot.datasets.utils import create_lerobot_dataset_card, get_repo_versions, get_safe_version +from lerobot.datasets.utils import ( + create_lerobot_dataset_card, + get_repo_versions, + get_safe_version, + resolve_episode_indices, +) from lerobot.utils.constants import ACTION, OBS_IMAGES from lerobot.utils.feature_utils import combine_feature_dicts @@ -62,6 +67,20 @@ def test_default_parameters(): ] +def test_resolve_episode_indices_applies_allowlist_and_exclusions(): + assert resolve_episode_indices([4, 1, 3, 0], 5, [1, 4]) == [3, 0] + + +def test_resolve_episode_indices_preserves_none_without_filtering(): + assert resolve_episode_indices(None, 5) is None + + +def test_resolve_episode_indices_ignores_out_of_range_values(caplog): + assert resolve_episode_indices([-1, 0, 3, 5], 4, [-2, 3, 8]) == [0] + assert "Ignoring episode indices outside the dataset range [0, 4): [-1, 5]" in caplog.text + assert "Ignoring excluded episode indices outside the dataset range [0, 4): [-2, 8]" in caplog.text + + @pytest.mark.parametrize("token", ["hf_test_token", True, False]) def test_get_repo_versions_forwards_token(monkeypatch, token): api = Mock() diff --git a/tests/datasets/test_language_render.py b/tests/datasets/test_language_render.py index fcef41fd8..4362c3b48 100644 --- a/tests/datasets/test_language_render.py +++ b/tests/datasets/test_language_render.py @@ -343,6 +343,85 @@ def test_resolve_task_explicit_override_beats_rephrasings(): assert rendered["messages"][0]["content"] == "explicit override wins" +def test_flow_only_low_level_recipe_renders_without_target(): + """Regression: a flow-only ``low_level`` recipe has no ``target`` turn — + its supervision is the action-expert flow loss, not text-CE. It must + still render (not ``None``), otherwise every blend draw of it is dropped + and the action expert never receives a flow loss.""" + recipe = TrainingRecipe( + messages=[ + MessageTurn( + role="user", + content="${subtask}", + stream="low_level", + if_present="subtask", + ), + ], + bindings={"subtask": "active_at(t, style=subtask)"}, + ) + + rendered = render_sample( + recipe=recipe, + persistent=PERSISTENT, + events=[], + t=0.5, + sample_idx=0, + task="clean kitchen", + ) + + assert rendered is not None + assert rendered["messages"] == [{"role": "user", "content": "subtask 0"}] + assert rendered["message_streams"] == ["low_level"] + assert rendered["target_message_indices"] == [] + + +def test_vqa_frame_is_consumed_over_the_weighted_blend(): + """A frame carrying a VQA annotation renders the routed VQA sub-recipe + even when its blend weight is tiny — VQA annotations are sparse and must + never be wasted on a subtask/action draw.""" + recipe = TrainingRecipe( + blend={ + "high_level_subtask": TrainingRecipe( + weight=0.99, + messages=[ + MessageTurn(role="user", content="${task}", stream="high_level"), + MessageTurn(role="assistant", content="a subtask", stream="high_level", target=True), + ], + ), + "descriptive_top_camera_name": TrainingRecipe( + weight=0.01, + route="vqa", + bindings={ + "vqa_query": "emitted_at(t, style=vqa, role=user, camera=observation.images.top)", + "vqa": "emitted_at(t, style=vqa, role=assistant, camera=observation.images.top)", + }, + messages=[ + MessageTurn( + role="user", content="${vqa_query}", stream="high_level", if_present="vqa_query" + ), + MessageTurn( + role="assistant", + content="${vqa}", + stream="high_level", + target=True, + if_present="vqa", + ), + ], + ), + } + ) + # A frame WITH a vqa event renders VQA on every sample_idx, despite the + # routed VQA weight being only 0.01. + for sample_idx in range(20): + rendered = render_sample( + recipe=recipe, persistent=PERSISTENT, events=EVENTS_AT_1, t=1.0, sample_idx=sample_idx, task="x" + ) + assert rendered["messages"][-1]["content"] == '{"count": 2}', sample_idx + # A frame WITHOUT a vqa event falls back to the normal weighted blend. + rendered = render_sample(recipe=recipe, persistent=PERSISTENT, events=[], t=1.0, sample_idx=0, task="x") + assert rendered["messages"][-1]["content"] == "a subtask" + + def test_emitted_at_persistent_tolerates_small_timestamp_drift(): """Persistent ``emitted_at`` should match within EMITTED_AT_TOLERANCE_S so callers that derive ``t`` arithmetically (``frame_idx / fps``) still diff --git a/tests/datasets/test_sampler.py b/tests/datasets/test_sampler.py index 7a5fc0fe0..dbe2eb7f0 100644 --- a/tests/datasets/test_sampler.py +++ b/tests/datasets/test_sampler.py @@ -25,7 +25,7 @@ from datasets import Dataset # noqa: E402 from lerobot.datasets.io_utils import ( hf_transform_to_torch, ) -from lerobot.datasets.sampler import EpisodeAwareSampler +from lerobot.datasets.sampler import EpisodeAwareSampler, compute_sampler_state def calculate_episode_data_index(hf_dataset: Dataset) -> dict[str, torch.Tensor]: @@ -154,8 +154,6 @@ def test_partial_episode_drop_warns(caplog): # --- seeded (seed, epoch) shuffling, resume, and state --- -from lerobot.datasets.sampler import compute_sampler_state # noqa: E402 - EPISODE_BOUNDS = ([0, 2, 3], [2, 3, 6]) # episodes of 2, 1 and 3 frames diff --git a/tests/processor/test_render_messages_processor.py b/tests/processor/test_render_messages_processor.py index 62c02454a..c85a21649 100644 --- a/tests/processor/test_render_messages_processor.py +++ b/tests/processor/test_render_messages_processor.py @@ -4,15 +4,22 @@ import pytest pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") +import numpy as np # noqa: E402 import torch # noqa: E402 from lerobot.configs.recipe import MessageTurn, TrainingRecipe # noqa: E402 from lerobot.lerobot_types import TransitionKey # noqa: E402 from lerobot.processor.converters import create_transition # noqa: E402 -from lerobot.processor.render_messages_processor import RenderMessagesStep # noqa: E402 +from lerobot.processor.render_messages_processor import ( # noqa: E402 + RenderMessagesStep, + _fallback_low_level_render, + _select_batch_indices, +) -def test_render_messages_step_noops_without_language_columns(): +def test_render_messages_step_renders_task_fallback_without_language_columns(): + """No language columns + a task string → low-level task fallback render, + matching what the policy sees at eval time on unannotated observations.""" recipe = TrainingRecipe( messages=[ MessageTurn(role="user", content="${task}", stream="high_level"), @@ -21,6 +28,24 @@ def test_render_messages_step_noops_without_language_columns(): ) transition = create_transition(complementary_data={"task": "do it"}) + out = RenderMessagesStep(recipe)(transition) + data = out[TransitionKey.COMPLEMENTARY_DATA] + + assert data["messages"] == [{"role": "user", "content": "do it"}] + assert data["message_streams"] == ["low_level"] + assert data["target_message_indices"] == [] + assert data["task"] == "do it" + + +def test_render_messages_step_noops_without_language_columns_or_task(): + recipe = TrainingRecipe( + messages=[ + MessageTurn(role="user", content="${task}", stream="high_level"), + MessageTurn(role="assistant", content="${subtask}", stream="low_level", target=True), + ] + ) + transition = create_transition(complementary_data={}) + assert RenderMessagesStep(recipe)(transition) == transition @@ -58,3 +83,129 @@ def test_render_messages_step_renders_and_drops_raw_language(): assert data["messages"][-1]["content"] == "reach carefully" assert data["message_streams"] == ["high_level", "low_level"] assert data["target_message_indices"] == [1] + + +def test_render_messages_step_falls_back_to_low_level_task_when_recipe_misses(): + recipe = TrainingRecipe( + messages=[ + MessageTurn( + role="assistant", + content="${subtask}", + stream="high_level", + target=True, + if_present="subtask", + ), + ] + ) + transition = create_transition( + complementary_data={ + "task": "pick the cube", + "timestamp": torch.tensor(0.0), + "index": torch.tensor(7), + "language_persistent": [], + "language_events": [{"style": "unmatched", "timestamp": 0.0}], + } + ) + + out = RenderMessagesStep(recipe)(transition) + data = out[TransitionKey.COMPLEMENTARY_DATA] + + assert data["messages"] == [{"role": "user", "content": "pick the cube"}] + assert data["message_streams"] == ["low_level"] + assert data["target_message_indices"] == [] + + +def test_render_messages_step_falls_back_per_sample_in_batched_language(): + recipe = TrainingRecipe( + messages=[ + MessageTurn( + role="assistant", + content="${subtask}", + stream="high_level", + target=True, + if_present="subtask", + ), + ] + ) + transition = create_transition( + action=torch.arange(4).reshape(2, 2), + complementary_data={ + "task": ["pick the cube", "open the drawer"], + "timestamp": torch.tensor([0.0, 1.0]), + "index": torch.tensor([7, 8]), + "language_persistent": [[], []], + "language_events": [ + [{"style": "unmatched", "timestamp": 0.0}], + [{"style": "unmatched", "timestamp": 1.0}], + ], + }, + ) + + out = RenderMessagesStep(recipe)(transition) + data = out[TransitionKey.COMPLEMENTARY_DATA] + + assert data["messages"] == [ + [{"role": "user", "content": "pick the cube"}], + [{"role": "user", "content": "open the drawer"}], + ] + assert data["message_streams"] == [["low_level"], ["low_level"]] + assert data["target_message_indices"] == [[], []] + + +def test_render_messages_step_rejects_mismatched_non_empty_language_batches(): + recipe = TrainingRecipe( + messages=[ + MessageTurn( + role="assistant", + content="${subtask}", + stream="high_level", + target=True, + if_present="subtask", + ) + ] + ) + transition = create_transition( + complementary_data={ + "timestamp": torch.tensor([0.0, 1.0, 2.0]), + "language_persistent": [[], []], + "language_events": [[{"style": "unmatched"}], [], []], + } + ) + + with pytest.raises(ValueError, match="must have equal lengths"): + RenderMessagesStep(recipe)(transition) + + +def test_select_batch_indices_slices_numpy_action(): + action = np.arange(6).reshape(3, 2) + transition = create_transition(action=action) + + selected = _select_batch_indices(transition, [2, 0], batch_size=3) + + np.testing.assert_array_equal(selected[TransitionKey.ACTION], action[[2, 0]]) + + +def test_select_batch_indices_slices_robot_action_dict(): + transition = create_transition( + action={ + "joints": np.arange(6).reshape(3, 2), + "gripper": torch.tensor([[0.0], [1.0], [2.0]]), + } + ) + + selected = _select_batch_indices(transition, [2, 0], batch_size=3) + + np.testing.assert_array_equal(selected[TransitionKey.ACTION]["joints"], np.array([[4, 5], [0, 1]])) + assert torch.equal(selected[TransitionKey.ACTION]["gripper"], torch.tensor([[2.0], [0.0]])) + + +def test_select_batch_indices_rejects_misaligned_list(): + transition = create_transition(complementary_data={"task": ["one", "two"]}) + + with pytest.raises(ValueError, match="expected 3 values, got 2"): + _select_batch_indices(transition, [2, 0], batch_size=3) + + +def test_fallback_low_level_render_rejects_partially_missing_task_batch(): + with pytest.raises(ValueError, match=r"missing task at indices \[1\]"): + _fallback_low_level_render(["pick cube", None, "place cube"]) diff --git a/tests/processor/test_tokenizer_processor.py b/tests/processor/test_tokenizer_processor.py index 52277c78f..10e2b6389 100644 --- a/tests/processor/test_tokenizer_processor.py +++ b/tests/processor/test_tokenizer_processor.py @@ -19,6 +19,7 @@ Tests for the TokenizerProcessorStep class. """ import tempfile +from pathlib import Path from unittest.mock import patch import pytest @@ -26,7 +27,7 @@ import torch from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature from lerobot.lerobot_types import TransitionKey -from lerobot.processor import DataProcessorPipeline, TokenizerProcessorStep +from lerobot.processor import ActionTokenizerProcessorStep, DataProcessorPipeline, TokenizerProcessorStep from lerobot.processor.converters import create_transition, identity_transition from lerobot.utils.constants import ( ACTION, @@ -87,6 +88,51 @@ class MockTokenizer: return result + def save_pretrained(self, save_directory: str | Path) -> None: + save_directory = Path(save_directory) + save_directory.mkdir(parents=True, exist_ok=True) + (save_directory / "tokenizer_config.json").write_text("{}") + + +def test_action_tokenizer_config_preserves_token_mapping(): + processor = object.__new__(ActionTokenizerProcessorStep) + processor.trust_remote_code = True + processor.max_action_tokens = 384 + processor.fast_skip_tokens = 64 + processor.paligemma_tokenizer_name = "custom/paligemma" + processor.allow_truncation = False + processor.action_tokenizer_name = "custom/fast" + processor.action_tokenizer_input_object = None + + assert processor.get_config() == { + "trust_remote_code": True, + "max_action_tokens": 384, + "fast_skip_tokens": 64, + "paligemma_tokenizer_name": "custom/paligemma", + "allow_truncation": False, + "action_tokenizer_name": "custom/fast", + } + + +def test_action_tokenizer_can_reject_truncated_sequences(): + processor = object.__new__(ActionTokenizerProcessorStep) + processor.max_action_tokens = 4 + processor.fast_skip_tokens = 128 + processor.allow_truncation = False + processor.action_tokenizer = lambda _actions: [1, 2, 3] + processor._paligemma_tokenizer = type( + "Tokenizer", + (), + { + "vocab_size": 1000, + "bos_token_id": 2, + "encode": lambda _self, text, **_kwargs: [10, 11] if text == "Action: " else [12, 1], + }, + )() + + with pytest.raises(ValueError, match="max_action_tokens=4"): + processor._tokenize_action(torch.zeros(1, 2, 1)) + @pytest.fixture def mock_tokenizer(): @@ -490,9 +536,11 @@ def test_save_and_load_pretrained_with_tokenizer_name(mock_auto_tokenizer): @skip_if_package_missing("transformers") -def test_save_and_load_pretrained_with_tokenizer_object(): - """Test saving and loading processor with tokenizer object using overrides.""" +@patch("lerobot.processor.tokenizer_processor.AutoTokenizer") +def test_save_and_load_pretrained_with_tokenizer_object(mock_auto_tokenizer): + """Test that a tokenizer object is saved and reloads from its local artifact.""" mock_tokenizer = MockTokenizer(vocab_size=100) + mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer original_processor = TokenizerProcessorStep( tokenizer=mock_tokenizer, max_length=32, task_key="instruction" @@ -506,11 +554,12 @@ def test_save_and_load_pretrained_with_tokenizer_object(): # Save processor robot_processor.save_pretrained(temp_dir) - # Load processor with tokenizer override (since tokenizer object wasn't saved) + assert (Path(temp_dir) / "tokenizer" / "tokenizer_config.json").is_file() + + # Load processor without an object override: the saved artifact is portable. loaded_processor = DataProcessorPipeline.from_pretrained( temp_dir, config_filename="dataprocessorpipeline.json", - overrides={"tokenizer_processor": {"tokenizer": mock_tokenizer}}, to_transition=identity_transition, to_output=identity_transition, ) diff --git a/tests/utils/test_collate.py b/tests/utils/test_collate.py index 2b23b3180..94d87cbcf 100644 --- a/tests/utils/test_collate.py +++ b/tests/utils/test_collate.py @@ -9,7 +9,7 @@ import torch # noqa: E402 from lerobot.utils.collate import lerobot_collate_fn # noqa: E402 -def test_lerobot_collate_preserves_messages_and_drops_raw_language(): +def test_lerobot_collate_preserves_messages_and_raw_language(): batch = [ { "index": torch.tensor(0), @@ -17,14 +17,14 @@ def test_lerobot_collate_preserves_messages_and_drops_raw_language(): "message_streams": ["low_level"], "target_message_indices": [0], "language_persistent": [{"content": "raw"}], - "language_events": [], + "language_events": [{"content": "event a"}], }, { "index": torch.tensor(1), "messages": [{"role": "assistant", "content": "b"}], "message_streams": ["low_level"], "target_message_indices": [0], - "language_persistent": [{"content": "raw"}], + "language_persistent": [{"content": "raw b"}], "language_events": [], }, ] @@ -36,8 +36,8 @@ def test_lerobot_collate_preserves_messages_and_drops_raw_language(): assert out["messages"][1][0]["content"] == "b" assert out["message_streams"] == [["low_level"], ["low_level"]] assert out["target_message_indices"] == [[0], [0]] - assert "language_persistent" not in out - assert "language_events" not in out + assert out["language_persistent"] == [[{"content": "raw"}], [{"content": "raw b"}]] + assert out["language_events"] == [[{"content": "event a"}], []] def test_lerobot_collate_passes_through_standard_batch():