diff --git a/pyproject.toml b/pyproject.toml index ecc5c7208..1d2a05ff8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -250,14 +250,6 @@ annotations = [ # install it locally only if you run your own ``vllm serve``. ] -# Tool implementations under src/lerobot/tools/. Each tool's dependencies -# are isolated so adding a new tool doesn't bloat the base install. -# Currently only `say` (Kyutai pocket-tts; CPU-only, ~100M params). -tools = [ - "pocket-tts>=1.0.0,<3.0.0", - "scipy>=1.11.0,<2.0.0", # SayTool.output_dir uses scipy.io.wavfile -] - # Development dev = ["pre-commit>=3.7.0,<5.0.0", "debugpy>=1.8.1,<1.9.0", "lerobot[grpcio-dep]", "grpcio-tools>=1.73.1,<2.0.0", "mypy>=1.19.1", "ruff>=0.14.1", "lerobot[notebook]"] notebook = ["jupyter>=1.0.0,<2.0.0", "ipykernel>=6.0.0,<7.0.0"] diff --git a/src/lerobot/configs/train.py b/src/lerobot/configs/train.py index 1828ab004..59ec65bbf 100644 --- a/src/lerobot/configs/train.py +++ b/src/lerobot/configs/train.py @@ -114,14 +114,6 @@ class TrainPipelineConfig(HubMixin): ema: EMAConfig = field(default_factory=EMAConfig) peft: PeftConfig | None = None - # VQA oversampling. When set (a fraction in (0, 1)), the training - # dataloader uses a WeightedEpisodeAwareSampler that draws frames - # carrying a `vqa` language annotation often enough that they make - # up roughly this fraction of the training stream. VQA annotations - # are typically sparse, so without this they are underrepresented. - # `None` (default) keeps uniform episode-aware sampling. - vqa_target_fraction: float | None = None - # Sample weighting configuration (e.g., for RA-BC training). Old # inline ``use_rabc`` / ``rabc_*`` params are migrated to this # field by ``_migrate_legacy_rabc_keys`` above. diff --git a/src/lerobot/datasets/__init__.py b/src/lerobot/datasets/__init__.py index 288b85fb9..f207e4e35 100644 --- a/src/lerobot/datasets/__init__.py +++ b/src/lerobot/datasets/__init__.py @@ -49,7 +49,7 @@ from .lerobot_dataset import LeRobotDataset from .multi_dataset import MultiLeRobotDataset from .pipeline_features import aggregate_pipeline_dataset_features, create_initial_features from .pyav_utils import check_video_encoder_parameters_pyav, detect_available_encoders_pyav -from .sampler import EpisodeAwareSampler, WeightedEpisodeAwareSampler, compute_sampler_state +from .sampler import EpisodeAwareSampler, compute_sampler_state from .streaming_dataset import StreamingLeRobotDataset from .utils import DEFAULT_EPISODES_PATH, create_lerobot_dataset_card from .video_utils import VideoEncodingManager @@ -77,7 +77,6 @@ __all__ = [ "DEFAULT_QUANTILES", "EVENT_ONLY_STYLES", "EpisodeAwareSampler", - "WeightedEpisodeAwareSampler", "LANGUAGE_EVENTS", "LANGUAGE_PERSISTENT", "LeRobotDataset", diff --git a/src/lerobot/datasets/sampler.py b/src/lerobot/datasets/sampler.py index 92af81291..af85dff9b 100644 --- a/src/lerobot/datasets/sampler.py +++ b/src/lerobot/datasets/sampler.py @@ -154,81 +154,6 @@ class EpisodeAwareSampler: return self._num_frames -class WeightedEpisodeAwareSampler(EpisodeAwareSampler): - """``EpisodeAwareSampler`` that draws frames *with replacement* in - proportion to per-frame weights. - - Used to oversample frames carrying a sparse annotation (e.g. a VQA - question) so the policy sees them more often than their natural - dataset density. One epoch still yields ``len(self.indices)`` - samples — the weights only change the *composition* of the stream, - not its length. Each epoch re-draws, so the oversampled subset - varies run to run. - """ - - def __init__( - self, - dataset_from_indices: list[int], - dataset_to_indices: list[int], - frame_weights, - *, - episode_indices_to_use: list | None = None, - drop_n_first_frames: int = 0, - drop_n_last_frames: int = 0, - seed: int = 0, - ): - """ - Args: - dataset_from_indices: Episode start indices (see ``EpisodeAwareSampler``). - dataset_to_indices: Episode end indices. - frame_weights: 1-D sequence/tensor of non-negative weights, one per - dataset frame (length == total dataset frames). Higher weight ⇒ - that frame is sampled more often. - episode_indices_to_use / drop_n_first_frames / drop_n_last_frames: - Same meaning as ``EpisodeAwareSampler`` — the episode-boundary - frame filtering is applied first, then weighting is restricted - to the surviving frames. - """ - super().__init__( - dataset_from_indices, - dataset_to_indices, - episode_indices_to_use=episode_indices_to_use, - drop_n_first_frames=drop_n_first_frames, - drop_n_last_frames=drop_n_last_frames, - shuffle=False, - seed=seed, - ) - weights = torch.as_tensor(frame_weights, dtype=torch.double).flatten() - idx = torch.tensor(self.indices, dtype=torch.long) - if weights.numel() <= int(idx.max()): - raise ValueError( - f"frame_weights has {weights.numel()} entries but the sampler " - f"references frame index {int(idx.max())}." - ) - selected = weights[idx] - if not torch.isfinite(selected).all() or bool((selected < 0).any()): - raise ValueError("frame_weights must be finite and non-negative.") - if float(selected.sum()) <= 0.0: - # All surviving frames have zero weight — fall back to uniform. - selected = torch.ones_like(selected) - self._weights = selected - self._indices = idx - - def __iter__(self) -> Iterator[int]: - epoch, start = self._epoch, self._start_index - self._epoch += 1 - self._start_index = 0 - generator = self._epoch_generator(epoch) - picks = torch.multinomial( - self._weights, - num_samples=self._num_frames, - replacement=True, - generator=generator, - ) - for i in picks[start:].tolist(): - yield int(self._indices[i]) - - def compute_sampler_state(step: int, num_frames: int, batch_size: int, num_processes: int) -> dict: """Map an optimization step to an `EpisodeAwareSampler` state for sample-exact resume. diff --git a/src/lerobot/datasets/utils.py b/src/lerobot/datasets/utils.py index d7a54e704..0a4452db9 100644 --- a/src/lerobot/datasets/utils.py +++ b/src/lerobot/datasets/utils.py @@ -26,7 +26,6 @@ import numpy as np import packaging.version import torch from huggingface_hub import DatasetCard, DatasetCardData, HfApi -from huggingface_hub.errors import RevisionNotFoundError from lerobot.utils.utils import flatten_dict, unflatten_dict diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index 36a9c5fb9..0e5fa7fb7 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -96,12 +96,8 @@ def _restore_pi052_pretrained_state( base = Path(pretrained_path) if not base.exists(): - # ``pretrained_path`` may be a HF Hub repo id rather than a local dir. - # ``from_pretrained`` downloads the model weights, but pi052 builds its - # processors fresh (so the generic loader never fetches them), leaving - # the processor JSON + normalizer-stat safetensors un-downloaded. Resolve - # them from the hub here — otherwise the quantile stats are silently left - # at fresh init and the policy runs completely un-normalized. + # Hub repo id, not a local dir: fetch the processor JSON + stats here + # (the generic loader never does for pi052's fresh-built processors). try: from huggingface_hub import snapshot_download # noqa: PLC0415 @@ -398,17 +394,8 @@ def make_pre_post_processors( policy configuration type. """ if pretrained_path and getattr(policy_cfg, "type", None) == "pi052": - # pi052 pipelines don't roundtrip through the saved - # ``policy_preprocessor.json``: ``RenderMessagesStep`` holds a - # Python ``TrainingRecipe`` (not JSON-serializable; saved as - # ``{}``) and ``ActionTokenizerProcessorStep`` saves a host-only - # FAST tokenizer path. Generic ``from_pretrained`` then dies - # with ``RenderMessagesStep.__init__() missing 1 required - # positional argument: 'recipe'`` (job 22164494). - # - # Mirror ``lerobot_pi052_runtime``'s bootstrap: build pipelines - # fresh from ``config.recipe_path`` and transplant the saved - # stateful blobs (normalizer stats) from the checkpoint dir. + # pi052 pipelines don't JSON-roundtrip — rebuild fresh and transplant + # saved state (see ``_restore_pi052_pretrained_state`` for why). from .pi052.processor_pi052 import make_pi052_pre_post_processors preprocessor, postprocessor = make_pi052_pre_post_processors( diff --git a/src/lerobot/policies/pi052/configuration_pi052.py b/src/lerobot/policies/pi052/configuration_pi052.py index ff5ff395d..f79534bf1 100644 --- a/src/lerobot/policies/pi052/configuration_pi052.py +++ b/src/lerobot/policies/pi052/configuration_pi052.py @@ -56,11 +56,9 @@ class PI052Config(PI05Config): # Recipe / language stack --------------------------------------------- recipe_path: str | None = "recipes/subtask_mem.yaml" - """Path (absolute or relative to ``src/lerobot/configs/``) to a - ``TrainingRecipe`` YAML. Defaults to the compact Hi-Robot blend - shipped with this policy. Set to ``None`` to disable recipe - rendering and fall back to π0.5's single-task ``Task: ... Action:`` - prompt path (unannotated datasets keep working that way).""" + """``TrainingRecipe`` YAML path (absolute or relative to + ``src/lerobot/configs/``). ``None`` disables recipe rendering — unannotated + datasets fall back to π0.5's plain ``Task: ... Action:`` prompt.""" apply_chat_template: bool = False """PaliGemma is *not* chat-pretrained — its tokenizer doesn't ship a diff --git a/src/lerobot/policies/pi052/debug_utils.py b/src/lerobot/policies/pi052/debug_utils.py new file mode 100644 index 000000000..c0ae0c59e --- /dev/null +++ b/src/lerobot/policies/pi052/debug_utils.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python + +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Training-time debug helpers for PI052's language head.""" + +import logging +from typing import Any + + +def print_debug_text_predictions(policy: Any, batch: dict[str, Any], step: int, n_samples: int = 5) -> None: + """Forward the current batch and print head-argmax vs label per supervised position. + + Opt-in via ``LEROBOT_DEBUG_PREDS_EVERY=``. Only the + policy types that expose ``debug_text_predictions`` participate + (currently PI052); others are silently skipped. Pretty-prints up to + ``n_samples`` samples from the current batch, showing the prompt, + every supervised position's (label, prediction, ✓/✗), and a + per-sample token-accuracy summary — the cheapest "is text training + actually learning anything" signal. + """ + # Accelerator/DDP wraps the policy in a ``module`` attribute and + # doesn't proxy custom methods through, so a naive + # ``hasattr(policy, "debug_text_predictions")`` returns False on the + # wrapper — and the helper would silently no-op. Walk through any + # ``.module`` indirection (DDP, FSDP, ``accelerator.prepare`` wrappers) + # to reach the raw policy that actually defines the method. + inner = policy + while hasattr(inner, "module") and not hasattr(inner, "debug_text_predictions"): + inner = inner.module + if not hasattr(inner, "debug_text_predictions"): + logging.warning( + "LEROBOT_DEBUG_PREDS_EVERY set but policy %s has no " + "debug_text_predictions method — skipping dump.", + type(inner).__name__, + ) + return + try: + debug = inner.debug_text_predictions(batch, max_samples=n_samples) + except Exception as exc: # noqa: BLE001 + logging.warning("debug_text_predictions failed: %s", exc, exc_info=True) + return + if not debug: + logging.warning( + "debug_text_predictions returned no supervised samples — current batch has no text labels." + ) + return + policy = inner # used below for select_message-style decoding parity + + # Build a tokenizer for decoding — match training side exactly. + try: + from transformers import AutoTokenizer # noqa: PLC0415 + + from lerobot.policies.pi052.text_processor_pi052 import ( # noqa: PLC0415 + register_paligemma_loc_tokens, + ) + + tok_name = getattr(policy.config, "tokenizer_name", None) or "google/paligemma-3b-pt-224" + tokenizer = register_paligemma_loc_tokens(AutoTokenizer.from_pretrained(tok_name)) + except Exception as exc: # noqa: BLE001 + logging.warning("debug preds: tokenizer load failed: %s", exc) + return + + ids = debug["input_ids"] + labels = debug["labels"] + preds = debug["predictions"] + attn = debug["attention_mask"] + + n = ids.shape[0] + print( + f"\n========== STEP {step} DEBUG PREDICTIONS ({n} samples) ==========", + flush=True, + ) + for s in range(n): + a = attn[s].tolist() + real = sum(a) + sid = ids[s].tolist() + sl = labels[s].tolist() + sp = preds[s].tolist() + prompt = tokenizer.decode(sid[:real], skip_special_tokens=False) + print(f"\n --- sample {s + 1}/{n} ---", flush=True) + print(f" prompt: {prompt!r}", flush=True) + + # Ground-truth target (the contiguous supervised label span). + sup_ids = [int(sid[i]) for i in range(real) if sl[i] != -100] + if sup_ids: + print( + f" target (ground truth) : {tokenizer.decode(sup_ids, skip_special_tokens=False)!r}", + flush=True, + ) + + # Training-side teacher-forced argmax on the same prompt+target. + n_sup = n_ok = 0 + teacher_chars: list[int] = [] + for i in range(1, real): + label = sl[i] + if label == -100: + continue + n_sup += 1 + pred = int(sp[i - 1]) + teacher_chars.append(pred) + if label == pred: + n_ok += 1 + teacher_text = tokenizer.decode(teacher_chars, skip_special_tokens=False) if teacher_chars else "" + acc = n_ok / max(n_sup, 1) + print( + f" training argmax (teacher-fed) : {teacher_text!r} acc={n_ok}/{n_sup}={acc:.1%}", + flush=True, + ) + print("=" * 60 + "\n", flush=True) diff --git a/src/lerobot/policies/pi052/fit_fast_tokenizer.py b/src/lerobot/policies/pi052/fit_fast_tokenizer.py index 2f8224c72..298f84909 100644 --- a/src/lerobot/policies/pi052/fit_fast_tokenizer.py +++ b/src/lerobot/policies/pi052/fit_fast_tokenizer.py @@ -122,9 +122,11 @@ def fit_fast_tokenizer( if out_dir.exists() and (out_dir / _CACHE_SENTINEL).exists(): logger.info( - "FAST tokenizer cache hit: %s — re-using fitted tokenizer for " - "dataset=%s base=%s n_samples=%d", - out_dir, dataset_repo_id, base_tokenizer_name, n_samples, + "FAST tokenizer cache hit: %s — re-using fitted tokenizer for dataset=%s base=%s n_samples=%d", + out_dir, + dataset_repo_id, + base_tokenizer_name, + n_samples, ) return str(out_dir) @@ -136,10 +138,7 @@ def fit_fast_tokenizer( # and compiles a ``.pyc`` — concurrent writers occasionally produce # a stale / partial ``.pyc`` and the subsequent ``from .. import # UniversalActionProcessor`` raises ``AttributeError``. - is_leader = ( - int(os.environ.get("RANK", "0")) == 0 - and int(os.environ.get("LOCAL_RANK", "0")) == 0 - ) + is_leader = int(os.environ.get("RANK", "0")) == 0 and int(os.environ.get("LOCAL_RANK", "0")) == 0 if not is_leader: timeout_s = 1800.0 # 30 min — covers ~1024-sample fits on cold caches start = time.monotonic() @@ -155,15 +154,16 @@ def fit_fast_tokenizer( return str(out_dir) logger.info( - "FAST tokenizer cache miss — fitting on dataset=%s " - "base=%s n_samples=%d chunk_size=%d → %s", - dataset_repo_id, base_tokenizer_name, n_samples, chunk_size, out_dir, + "FAST tokenizer cache miss — fitting on dataset=%s base=%s n_samples=%d chunk_size=%d → %s", + dataset_repo_id, + base_tokenizer_name, + n_samples, + chunk_size, + out_dir, ) from transformers import AutoProcessor # noqa: PLC0415 - from lerobot.datasets.lerobot_dataset import LeRobotDataset # noqa: PLC0415 - # Stream a single episode's worth of action chunks at a time so # we don't blow memory on huge datasets. Random episode + # random start offset gives a reasonable spread. @@ -186,16 +186,14 @@ def fit_fast_tokenizer( # for ~2.5 h before NCCL killed it). Reading the ``action`` column # straight from the parquet shards is also faster: each per-episode # ``LeRobotDataset`` instantiation re-parses every meta file. - from huggingface_hub import snapshot_download # noqa: PLC0415 import pyarrow as _pa # noqa: PLC0415 import pyarrow.parquet as _pq # noqa: PLC0415 + from huggingface_hub import snapshot_download # noqa: PLC0415 snap = Path(snapshot_download(repo_id=dataset_repo_id, repo_type="dataset")) data_files = sorted((snap / "data").glob("chunk-*/file-*.parquet")) if not data_files: - raise RuntimeError( - f"FAST fit: no ``data/chunk-*/file-*.parquet`` shards found under {snap!s}." - ) + raise RuntimeError(f"FAST fit: no ``data/chunk-*/file-*.parquet`` shards found under {snap!s}.") # Read just the (episode_index, action) columns once across all # shards. This is the same pattern used elsewhere in the codebase @@ -215,9 +213,7 @@ def fit_fast_tokenizer( # Fallback path for nested-list types: flatten via to_pylist(). acts = np.asarray(acts_col.to_pylist(), dtype=np.float32) if acts.ndim != 2: - raise RuntimeError( - f"FAST fit: expected ``action`` rows to be 1-D vectors; got shape {acts.shape}." - ) + raise RuntimeError(f"FAST fit: expected ``action`` rows to be 1-D vectors; got shape {acts.shape}.") # Episode index → slice (start, stop) into ``acts`` along axis 0. # ``eps`` is monotonically increasing within each parquet shard but @@ -268,7 +264,9 @@ def fit_fast_tokenizer( actions = np.stack(actions_buf, axis=0).astype(np.float32) # (N, H, D) logger.info( "FAST fit: collected %d chunks of shape %s from %d episodes", - actions.shape[0], actions.shape[1:], eps_visited, + actions.shape[0], + actions.shape[1:], + eps_visited, ) # Quantile-normalise per dimension before fitting. diff --git a/src/lerobot/policies/pi052/inference/__init__.py b/src/lerobot/policies/pi052/inference/__init__.py index 012ff440c..6b613c3cd 100644 --- a/src/lerobot/policies/pi052/inference/__init__.py +++ b/src/lerobot/policies/pi052/inference/__init__.py @@ -14,12 +14,11 @@ """PI052 runtime adapter and CLI helpers.""" -from lerobot.policies.language_conditioned import ( +from lerobot.runtime import ( LanguageConditionedRuntime, RuntimeState, Tick, TickClock, - ToolCall, VQAResult, ) @@ -36,7 +35,6 @@ __all__ = [ "StdinReader", "Tick", "TickClock", - "ToolCall", "VQAResult", "make_state_panel", "print_robot_lines", diff --git a/src/lerobot/policies/pi052/inference/pi052_adapter.py b/src/lerobot/policies/pi052/inference/pi052_adapter.py index 8f4820374..2575d210f 100644 --- a/src/lerobot/policies/pi052/inference/pi052_adapter.py +++ b/src/lerobot/policies/pi052/inference/pi052_adapter.py @@ -21,7 +21,7 @@ import re from dataclasses import dataclass from typing import Any -from lerobot.policies.language_conditioned import RuntimeState, ToolCall, VQAResult +from lerobot.runtime import RuntimeState, VQAResult logger = logging.getLogger(__name__) @@ -69,10 +69,6 @@ class PI052PolicyAdapter: suppress_loc_tokens=kind in {"subtask", "memory", "interjection"}, ) - def parse_tool_calls(self, text: str) -> list[ToolCall]: - _plan, speech = split_plan_and_say(text) - return [ToolCall("say", {"text": speech})] if speech else [] - def plan_from_text(self, text: str) -> str: plan, _speech = split_plan_and_say(text) return "" if looks_like_gibberish(plan) else plan @@ -305,7 +301,3 @@ def split_plan_and_say(text: str) -> tuple[str, str]: speech = match.group(1).strip().strip('"').strip("'") plan = (text[: match.start()] + text[match.end() :]).strip() return plan, speech - - -def messages_for_vqa(question: str) -> list[dict[str, Any]]: - return [{"role": "user", "content": question}] diff --git a/src/lerobot/policies/pi052/inference/runtime.py b/src/lerobot/policies/pi052/inference/runtime.py index 57ec3b1d7..03765a138 100644 --- a/src/lerobot/policies/pi052/inference/runtime.py +++ b/src/lerobot/policies/pi052/inference/runtime.py @@ -19,12 +19,11 @@ from __future__ import annotations from collections.abc import Callable from typing import Any -from lerobot.policies.language_conditioned import ( +from lerobot.runtime import ( LanguageConditionedRuntime, RuntimeState, Tick, TickClock, - ToolCall, VQAResult, ) @@ -38,7 +37,6 @@ class PI052Runtime(LanguageConditionedRuntime): self, policy: Any, *, - tools: dict[str, Any] | None = None, observation_provider: Callable[[], dict | None] | None = None, robot_executor: Callable[[Any], None] | None = None, event_collector: Callable[[RuntimeState], None] | None = None, @@ -51,7 +49,6 @@ class PI052Runtime(LanguageConditionedRuntime): policy_adapter=policy if isinstance(policy, PI052PolicyAdapter) else PI052PolicyAdapter(policy), observation_provider=observation_provider, action_executor=robot_executor, - tools=tools or {}, event_collector=event_collector, chunk_hz=chunk_hz, ctrl_hz=ctrl_hz, @@ -67,6 +64,5 @@ __all__ = [ "RuntimeState", "Tick", "TickClock", - "ToolCall", "VQAResult", ] diff --git a/src/lerobot/policies/pi052/inference/runtime_cli.py b/src/lerobot/policies/pi052/inference/runtime_cli.py index 80b03cd9c..afe723fbc 100644 --- a/src/lerobot/policies/pi052/inference/runtime_cli.py +++ b/src/lerobot/policies/pi052/inference/runtime_cli.py @@ -50,9 +50,6 @@ With a real robot:: ``--policy.path`` accepts either a local directory or a Hugging Face Hub repo id. ``--dataset.repo_id`` likewise. - -Tool dispatch (TTS via ``SayTool``) is enabled by default when -``pocket-tts`` is installed; pass ``--no_tts`` to disable. """ from __future__ import annotations @@ -64,16 +61,11 @@ from collections.abc import Callable from contextlib import suppress from typing import Any +from .repl import _emit + logger = logging.getLogger("lerobot.pi052.runtime") -def _emit(state: Any, event_name: str) -> None: - if hasattr(state, "emit"): - state.emit(event_name) - else: - state.setdefault("events_this_tick", []).append(event_name) - - def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: p = argparse.ArgumentParser( description=("Interactive REPL runtime for a trained PI052 hierarchical VLA checkpoint."), @@ -243,18 +235,6 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: "wrong robot, robot not at home pose)." ), ) - p.add_argument( - "--no_tts", - action="store_true", - help="Disable the ``say`` tool dispatch.", - ) - p.add_argument( - "--tts.voice", - dest="tts_voice", - type=str, - default="alba", - help="Pocket-tts voice name (or path to a .wav for cloning).", - ) p.add_argument( "--chunk_hz", type=float, @@ -402,19 +382,11 @@ def _build_observation_provider( device: str, augment: bool = False, ) -> Callable[[], dict | None]: - """Build a closure that feeds dataset frames into the runtime. + """Closure feeding preprocessed dataset frames to the runtime, advancing + ``advance_per_tick`` frames per call and looping at episode end. - Each call returns a preprocessed observation batch (images + - state, batched, on the policy's device, normalized) suitable for - ``policy.select_action`` and ``policy.select_message``. The - closure walks the chosen episode forward by ``advance_per_tick`` - frames per call, looping back to the episode start when it falls - off the end. - - The dataset's ``language_persistent`` / ``language_events`` - columns are stripped before the sample reaches the preprocessor, - so ``RenderMessagesStep`` and ``PI052TextTokenizerStep`` are - no-ops; the runtime supplies its own messages from current state. + Language columns are stripped first — the runtime supplies its own + messages from current state, not the dataset's annotations. """ from lerobot.datasets.lerobot_dataset import LeRobotDataset # noqa: PLC0415 @@ -422,15 +394,9 @@ def _build_observation_provider( if len(ds) == 0: raise ValueError(f"Dataset {dataset_repo_id!r} episode {episode} is empty.") - # Optional: apply the same torchvision-v2 augmentation pipeline - # that training used, so dry-run sees frames from the augmented - # support region (not just the unperturbed dataset frames). When - # the LM head still generates coherent text under this, it has - # learned over the augmentation distribution — the *opposite* of - # the "memorised one specific frame per supervision" failure - # mode. When it collapses to ``\n`` here too, the head is hyper- - # specific to the unperturbed training samples and only the - # retrain can help. + # Optional: replay training's augmentation pipeline so dry-run probes the + # augmented support region — coherent text under jitter means the LM head + # generalized; collapse to "\n" means it memorised unperturbed frames. inference_aug = None if augment: from lerobot.transforms import ( # noqa: PLC0415 @@ -471,15 +437,9 @@ def _bootstrap_state_from_dataset( episode: int, start_frame: int, ) -> dict[str, str]: - """Pull task / active plan / active memory / active subtask at ``start_frame``. - - The model is heavily memorised on the exact training prompts the - recipe rendered from this dataset (canonical task wording, - persistent atoms emitted earlier in the episode). Reconstructing - that state at REPL startup lets the runtime's first prompt line - up with what training looked like — without it the model sees an - out-of-distribution prompt and falls back to its dominant - training mode (VQA JSON spam). + """Pull task / active plan / memory / subtask at ``start_frame``, so the + runtime's first prompt matches the canonical training prompts (an OOD + prompt makes the model fall back to its dominant mode, VQA JSON spam). """ from lerobot.datasets.lerobot_dataset import LeRobotDataset # noqa: PLC0415 @@ -535,20 +495,10 @@ def _select_task_interactively( ds_meta: Any, bootstrap_task: str | None, ) -> str | None: - """Ask the operator which task to run at startup. - - Behaviour: - - * If a dataset is loaded, build a numbered menu of every unique task - string in ``ds_meta.tasks`` (canonical bootstrap task listed first - as the default). Add a ``[c] type a custom task`` option. - * If no dataset is loaded, show a plain ``Enter task:`` prompt. - * Non-TTY runs (scripts, pipes) skip the prompt and return the - bootstrap task so the existing "first stdin line becomes task" - flow in ``_run_repl`` / ``_run_autonomous`` still works. - - Returns the chosen task string, or ``None`` when the operator declines - to pick one (Ctrl-D / empty + no default). + """Interactive task picker: numbered menu of dataset tasks (bootstrap task + as default) plus a custom-input option; plain prompt without a dataset. + Non-TTY runs skip the prompt and return the bootstrap task. Returns + ``None`` when the operator declines (Ctrl-D / empty + no default). """ options: list[str] = [] seen: set[str] = set() @@ -745,18 +695,10 @@ def _build_robot_observation_provider( task: str | None, ds_features: dict[str, Any] | None, ) -> Callable[[], dict | None]: - """Closure that reads from the robot, runs the policy preprocessor. - - Each call: ``robot.get_observation()`` (raw per-joint + per-camera - dict, possibly with scalar floats) → ``build_inference_frame`` - (extract the keys the dataset declared, reshape per-joint floats - into a single ``observation.state`` vector, prefix camera keys - with ``observation.images.``, convert to tensors with batch dim - on device) → wrap in an ``EnvTransition`` (the preprocessor - pipeline is transition-shaped, keyed by ``TransitionKey``) → - preprocessor (rename, normalise) → unwrap and return the flat - observation batch ``policy.select_action`` / ``policy.select_message`` - consume. + """Closure reading from the robot each call: ``robot.get_observation()`` → + ``build_inference_frame`` (state vector + image tensors, batched, on device) + → ``EnvTransition``-wrapped preprocessor (rename, normalise) → flat + observation batch for ``select_action`` / ``select_message``. """ import torch # noqa: PLC0415 @@ -768,19 +710,10 @@ def _build_robot_observation_provider( torch_device = torch.device(device) if isinstance(device, str) else device robot_type = getattr(robot, "robot_type", None) or getattr(getattr(robot, "config", None), "type", None) - # Pre-compute the camera-key → target (H, W) map from - # ``ds_features``. The training distribution sees frames at the - # recorded resolution (e.g. 480×640); a live Mac/USB camera will - # almost always hand us a different native size (720p / 1080p). - # PI052's internal ``resize_with_pad(512, 512)`` does pad the - # input to a fixed canvas, but the *geometry* of that pad differs - # by input aspect ratio — top/left padding varies, so the visual - # tokens at each tile carry different content than what the model - # saw at training. The action expert tolerates this (flow head - # rides broad geometry); the LM head, supervised much more - # tightly on visual features, goes out of distribution and the - # head's distribution at position 0 collapses to its dominant - # mode (a memorised ``\n``-only run in this checkpoint). + # Camera-key → training (H, W) map from ``ds_features``. Live cameras + # rarely match the recorded resolution, and a different aspect ratio + # changes resize_with_pad's padding geometry — the flow head tolerates + # that, but the tightly-supervised LM head goes OOD and collapses. _resize_logged = {"done": False} target_image_shapes: dict[str, tuple[int, int]] = {} if ds_features: @@ -814,12 +747,8 @@ def _build_robot_observation_provider( # columns the robot stream may carry through. _strip_runtime_owned_language_cols(raw) - # Force-match the training-time visual distribution: - # every camera frame the model trained on came from the - # dataset at its recorded (H, W). Resize the live frame to - # that exact shape so the downstream resize_with_pad geometry - # matches training. Without this the LM head is OOD on every - # tick. + # Resize live frames to the training (H, W) so the downstream + # resize_with_pad geometry matches what the model saw in training. if target_image_shapes: try: import cv2 as _cv2 # noqa: PLC0415 @@ -851,13 +780,8 @@ def _build_robot_observation_provider( continue raw[cam_key] = _cv2.resize(img, (target_w, target_h), interpolation=_cv2.INTER_AREA) _resize_logged["done"] = True - # Print the state vector once so the operator can eyeball - # it against the dataset's stats. State OOD is a real - # failure mode for VLAs — the prefix carries state via - # the projection layer, and a neutral home pose can - # easily sit a couple σ off the supervised support - # region. Gated on ``first_call`` so this doesn't spam - # every observation tick. + # One-shot state-vector print so the operator can eyeball it + # against dataset stats (state OOD is a real VLA failure mode). if first_call and "observation.state" in (ds_features or {}): state_names = ds_features["observation.state"].get("names") or [] state_vals = [raw.get(n) for n in state_names] @@ -1374,19 +1298,6 @@ def _make_state_panel_renderer( return _redraw -def _build_tools(no_tts: bool, tts_voice: str) -> dict[str, Any]: - """Instantiate the tools declared on this dataset/policy.""" - if no_tts: - return {} - try: - from lerobot.tools import SayTool # noqa: PLC0415 - - return {"say": SayTool(voice=tts_voice)} - except Exception as exc: # noqa: BLE001 - logger.warning("Could not initialise SayTool (%s) — speech disabled.", exc) - return {} - - def _silence_noisy_loggers() -> None: """Drop chatty third-party loggers down to WARNING. @@ -1529,10 +1440,6 @@ def main(argv: list[str] | None = None) -> int: augment=getattr(args, "dataset_augment_at_inference", False), ) - tools = _build_tools(args.no_tts, args.tts_voice) - if tools: - print(f"[pi052] tools loaded: {list(tools)}", flush=True) - from lerobot.policies.pi052.inference import ( # noqa: PLC0415 LanguageConditionedRuntime, PI052PolicyAdapter, @@ -1540,7 +1447,6 @@ def main(argv: list[str] | None = None) -> int: runtime = LanguageConditionedRuntime( policy_adapter=PI052PolicyAdapter(policy), - tools=tools, observation_provider=observation_provider, action_executor=robot_executor, # No background event collector — the REPL drives ticks diff --git a/src/lerobot/policies/pi052/inference/steps.py b/src/lerobot/policies/pi052/inference/steps.py deleted file mode 100644 index 60ead4ac4..000000000 --- a/src/lerobot/policies/pi052/inference/steps.py +++ /dev/null @@ -1,29 +0,0 @@ -# Copyright 2026 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Compatibility exports for PI052 model helper imports.""" - -from .pi052_adapter import ( - _build_text_batch, - _generate_with_policy, - _get_loc_tokenizer, - looks_like_gibberish as _looks_like_gibberish, -) - -__all__ = [ - "_build_text_batch", - "_generate_with_policy", - "_get_loc_tokenizer", - "_looks_like_gibberish", -] diff --git a/src/lerobot/policies/pi052/inference/ui.py b/src/lerobot/policies/pi052/inference/ui.py index 28b15a8d7..05b127cdc 100644 --- a/src/lerobot/policies/pi052/inference/ui.py +++ b/src/lerobot/policies/pi052/inference/ui.py @@ -25,12 +25,8 @@ Two-zone terminal layout: └───────────────────────────────────────────────────┘ > _ -The state panel re-renders on every state change. Chat lines are -``console.print``'d above the live region so they accumulate naturally -in scrollback. Implemented with :class:`rich.live.Live` plus -:func:`rich.console.Console.input` for the prompt — when an input is -pending, ``rich.Live`` auto-suspends so the input doesn't fight the -panel for cursor position. +Chat lines print above a ``rich.Live`` region (natural scrollback); the +state panel re-renders on change, auto-suspending while input is pending. """ from __future__ import annotations diff --git a/src/lerobot/policies/pi052/inference/vqa.py b/src/lerobot/policies/pi052/inference/vqa.py index f2e6f8a7b..6b690abc8 100644 --- a/src/lerobot/policies/pi052/inference/vqa.py +++ b/src/lerobot/policies/pi052/inference/vqa.py @@ -131,19 +131,10 @@ def _loc_to_norm(idx: int) -> float: def parse_loc_answer(answer: str) -> dict | None: """Parse a PaliGemma ````-format spatial VQA answer. - PI052 trains spatial answers in PaliGemma's native detection - vocabulary, label-first: a point is ``