admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-07 10:01:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(pi052): trim PR — remove say tool, debug gates, dead code; move runtime

Browse Source 
Cleanup pass over the language-support PR to cut LOC and scope creep.

Removals:
- SayTool + tools/ package (registry, Tool protocol, [tools] extra) and the
  runtime's tool-dispatch path. Kept <say> training supervision and inference
  stripping so speech-annotated datasets still train.
- WeightedEpisodeAwareSampler + VQA oversampling wiring
  (_build_vqa_oversample_weights, vqa_target_fraction) — training uses plain
  EpisodeAwareSampler again.
- Debug env-gates PI052_DEBUG_TENSORS, PI052_SUBTASK_USE_TASK, EVAL_TASK_OVERRIDE.
- Dead code: broken _tp._DUMP_BUDGET block, unused imports (copy/Tensor,
  RevisionNotFoundError, LeRobotDataset, os), messages_for_vqa, steps.py shim
  (modeling imports pi052_adapter directly), duplicated _emit, builtins.type[T].

Moves:
- Policy-agnostic runtime -> src/lerobot/runtime/ (LanguageConditionedRuntime +
  adapter Protocol + state); pi052 keeps only its adapter + CLI. Tests -> tests/runtime/.

Other:
- Compacted verbose AI-authored comments/docstrings across pi052 (kept the
  hard-won DDP / barrier-timeout / reduce-max / VQA-routing notes).
- Relocated LM-head prediction debug helper to pi052/debug_utils.py.
- Fixed test_render_messages: assert task-fallback render (current behavior)
  instead of the stale no-op expectation.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Pepijn
2026-07-02 14:16:41 +02:00
parent d099ac91b3
commit 4fa9578e3d
32 changed files with 338 additions and 1266 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pyproject.toml
-8
View File
@@ -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"]
src/lerobot/configs/train.py
-8
View File
@@ -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.
src/lerobot/datasets/__init__.py
+1 -2
View File
@@ -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",
src/lerobot/datasets/sampler.py
-75
View File
@@ -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.
src/lerobot/datasets/utils.py
-1
View File
@@ -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
src/lerobot/policies/factory.py
+4 -17
View File
@@ -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(
src/lerobot/policies/pi052/configuration_pi052.py
+3 -5
View File
@@ -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
src/lerobot/policies/pi052/debug_utils.py
+121
View File
@@ -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=<step_interval>``. 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)
src/lerobot/policies/pi052/fit_fast_tokenizer.py
+18 -20
View File
@@ -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.
src/lerobot/policies/pi052/inference/__init__.py
+1 -3
View File
@@ -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",
src/lerobot/policies/pi052/inference/pi052_adapter.py
+1 -9
View File
@@ -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}]
src/lerobot/policies/pi052/inference/runtime.py
+1 -5
View File
@@ -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",
]
src/lerobot/policies/pi052/inference/runtime_cli.py
+28 -122
View File
@@ -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
src/lerobot/policies/pi052/inference/steps.py
-29
View File
@@ -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",
]
src/lerobot/policies/pi052/inference/ui.py
+2 -6
View File
@@ -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
src/lerobot/policies/pi052/inference/vqa.py
+7 -20
View File
@@ -131,19 +131,10 @@ def _loc_to_norm(idx: int) -> float:
def parse_loc_answer(answer: str) -> dict | None:
"""Parse a PaliGemma ``<loc>``-format spatial VQA answer.
PI052 trains spatial answers in PaliGemma's native detection
vocabulary, label-first: a point is ``<label> <locY><locX>``, a box
is ``<label> <locY0><locX0><locY1><locX1>``, and multiple boxes are
joined by `` ; `` (e.g. ``cube <loc..><loc..><loc..><loc..> ; box
<loc..><loc..><loc..><loc..>``). Loc-first formats are also accepted
— this parser strips loc tokens and treats the remainder as the
label, so order is irrelevant. Coordinates come back *normalized*
([0, 1]); the overlay denormalizes them against the chosen camera
frame's pixel size.
Returns ``{"kind", "payload", "normalized": True}`` on success
(``payload`` mirrors the JSON shapes so the overlay code is shared),
or ``None`` when the answer carries no ``<loc>`` tokens.
Point: ``<label> <locY><locX>``; box: ``<label> <locY0><locX0><locY1><locX1>``;
multiple boxes joined by `` ; `` (label/loc order irrelevant). Returns
``{"kind", "payload", "normalized": True}`` with [0, 1] coords mirroring the
JSON shapes (shared overlay code), or ``None`` without ``<loc>`` tokens.
"""
if not answer or "<loc" not in answer:
return None
@@ -178,14 +169,10 @@ def parse_loc_answer(answer: str) -> dict | None:
def parse_vqa_answer(answer: str) -> dict | None:
"""Parse a VQA answer string into ``{"kind", "payload"}``.
"""Parse a VQA answer (``<loc>`` text or JSON) into ``{"kind", "payload"}``.
``kind`` is one of the ``VQA_ANSWER_SHAPES`` names (``bbox``,
``keypoint``, ``count``, ``attribute``, ``spatial``) or ``"unknown"``
when the JSON doesn't match any known shape. PaliGemma ``<loc>``
spatial answers are detected first (PI052 trains them in that native
format). Returns ``None`` when the answer is neither ``<loc>`` text
nor a parseable JSON object.
``kind`` is a ``VQA_ANSWER_SHAPES`` name or ``"unknown"``; ``<loc>`` answers
are tried first. Returns ``None`` when neither format parses.
"""
if not answer or not answer.strip():
return None
src/lerobot/policies/pi052/modeling_pi052.py
+93 -185
View File
@@ -14,30 +14,16 @@
"""π0.5 v2 policy — dual-head training & hierarchical inference.
A thin subclass of :class:`PI05Policy` that:
* keeps the PaliGemma ``lm_head`` unfrozen during fine-tuning
(``PI05Policy`` zeroes / freezes it because it never reads from
the head; ``PI052Config.unfreeze_lm_head`` flips that),
* adds a ``text_loss`` term computed via cross-entropy on
``text_labels`` (built by ``PI052TextTokenizerStep``),
* adds :meth:`select_message` for AR text generation at inference
(the high-level step in the π0.5 paper's two-stage inference loop),
* combines both losses in :meth:`forward` per Eq. (1) of the paper:
L = H(x, f_θ_text) + α * ‖ω - a - f_θ_action(...)‖²
with α controllable via ``config.flow_loss_weight``.
The multi-rate inference runtime in ``lerobot.policies.pi052.inference``
(driven by the ``lerobot-pi052-runtime`` CLI) sits on top of this:
``predict_action_chunk`` for the action expert and ``select_message``
for the LM head.
π0.5 with the PaliGemma LM head re-enabled: adds a text CE loss on
``text_labels`` next to the flow loss (L = H(x, f_θ_text) + α·flow, α via
``config.flow_loss_weight``) and :meth:`select_message` for AR text
generation. The multi-rate runtime in ``lerobot.policies.pi052.inference``
(``lerobot-pi052-runtime`` CLI) drives ``predict_action_chunk`` +
``select_message``. See :class:`PI052Config` for the knobs.
"""
from __future__ import annotations
import builtins
import logging
import math
import types
@@ -72,6 +58,7 @@ logger = logging.getLogger(__name__)
# transformer (PaliGemmaWithExpertModel, sdpa_attention_forward,
# compute_layer_complete, get_gemma_config) lives in lerobot.policies.pi_gemma.
class ActionSelectKwargs(TypedDict, total=False):
inference_delay: int | None
prev_chunk_left_over: Tensor | None
@@ -524,9 +511,7 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
# Precompute the whole timestep schedule on-device once, instead of
# rebuilding a tensor from a Python float every step
# (``torch.tensor(time, device=cuda)`` is a host->device sync ×num_steps).
times = torch.tensor(
[1.0 + s * dt for s in range(num_steps)], dtype=torch.float32, device=device
)
times = torch.tensor([1.0 + s * dt for s in range(num_steps)], dtype=torch.float32, device=device)
x_t = noise
for step in range(num_steps):
@@ -562,20 +547,6 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
if self.rtc_processor is not None and self.rtc_processor.is_debug_enabled():
self.rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
import os as _os # noqa: PLC0415
if _os.environ.get("PI052_DEBUG_TENSORS") == "1" and not getattr(self, "_dbg_act_done", False):
import logging as _lg # noqa: PLC0415
_a = x_t.float()
ad = self.config.max_action_dim
_lg.getLogger(__name__).info(
"PI052_DEBUG predicted norm action chunk shape=%s min=%.3f max=%.3f mean=%.3f std=%.3f (real dims only) (expect ~[-1,1])",
tuple(x_t.shape), _a[..., :12].min().item(), _a[..., :12].max().item(),
_a[..., :12].mean().item(), _a[..., :12].std().item(),
)
self._dbg_act_done = True
return x_t
def denoise_step(
@@ -599,9 +570,7 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None]
position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1
full_att_2d_masks_4d = self._prepare_attention_masks_4d(
full_att_2d_masks, dtype=suffix_embs.dtype
)
full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks, dtype=suffix_embs.dtype)
self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager" # noqa: SLF001
# The expert forward appends the suffix K/V to the prefix cache in-place
@@ -626,7 +595,6 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
return self.action_out_proj(suffix_out)
# FAST action-token vocab size (``lerobot/fast-action-tokenizer``). The
# tokenizer maps a FAST BPE id ``t`` to the PaliGemma vocab id
# ``vocab_size - 1 - fast_skip_tokens - t`` (see ``TokenizerProcessorStep``),
@@ -802,9 +770,7 @@ def _fast_lin_ce(
# Fold the boolean mask into the target via ignore_index. No
# ``.any().item()`` sync — Liger returns 0.0 when every position
# is ignored, preserving graph capture for CUDA graphs.
shift_targets = torch.where(
shift_valid, shift_targets, torch.full_like(shift_targets, -100)
)
shift_targets = torch.where(shift_valid, shift_targets, torch.full_like(shift_targets, -100))
B, T_1, H = shift_hidden.shape
flat_hidden = shift_hidden.reshape(B * T_1, H).to(lm_head_weight.dtype)
@@ -827,6 +793,7 @@ def _fast_lin_ce(
# and V projections. Forward output is bit-equivalent to the standard
# layer; backward differs only on the path action_loss → VLM K/V.
def _compute_layer_ki(
layer_idx,
inputs_embeds,
@@ -912,11 +879,19 @@ def _compute_layer_ki(
att_vlm, _ = sdpa_attention_forward(
paligemma.model.language_model.layers[layer_idx].self_attn,
Q_vlm, K_for_vlm, V_for_vlm, mask_for_vlm, scaling,
Q_vlm,
K_for_vlm,
V_for_vlm,
mask_for_vlm,
scaling,
)
att_action, _ = sdpa_attention_forward(
paligemma.model.language_model.layers[layer_idx].self_attn,
Q_action, K_for_action, V_for_action, mask_for_action, scaling,
Q_action,
K_for_action,
V_for_action,
mask_for_action,
scaling,
)
att = torch.cat([att_vlm, att_action], dim=1)
@@ -983,22 +958,31 @@ def _paligemma_forward_ki(
hasattr(self.gemma_expert.model, "gradient_checkpointing")
and self.gemma_expert.model.gradient_checkpointing
and self.training
) or (
hasattr(self, "gradient_checkpointing") and self.gradient_checkpointing and self.training
)
) or (hasattr(self, "gradient_checkpointing") and self.gradient_checkpointing and self.training)
for layer_idx in range(num_layers):
if use_gc:
inputs_embeds = torch.utils.checkpoint.checkpoint(
_compute_layer_ki,
layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond,
use_reentrant=False, preserve_rng_state=False,
paligemma=self.paligemma, gemma_expert=self.gemma_expert,
layer_idx,
inputs_embeds,
attention_mask,
position_ids,
adarms_cond,
use_reentrant=False,
preserve_rng_state=False,
paligemma=self.paligemma,
gemma_expert=self.gemma_expert,
)
else:
inputs_embeds = _compute_layer_ki(
layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond,
paligemma=self.paligemma, gemma_expert=self.gemma_expert,
layer_idx,
inputs_embeds,
attention_mask,
position_ids,
adarms_cond,
paligemma=self.paligemma,
gemma_expert=self.gemma_expert,
)
outputs_embeds = []
@@ -1057,8 +1041,7 @@ class PI052Policy(PreTrainedPolicy):
backbone._pi052_orig_forward = backbone.forward
backbone.forward = types.MethodType(_paligemma_forward_ki, backbone)
logger.info(
"PI052: knowledge insulation enabled — action→VLM K/V "
"gradients are blocked in attention."
"PI052: knowledge insulation enabled — action→VLM K/V gradients are blocked in attention."
)
# Per-env hierarchical-inference state. Sized lazily on the first
@@ -1167,9 +1150,8 @@ class PI052Policy(PreTrainedPolicy):
):
return self._pi05_flow_forward(batch, reduction=reduction)
run_flow = (
self.config.flow_loss_weight > 0
and (predict_actions_t is None or bool(predict_actions_t.any().item()))
run_flow = self.config.flow_loss_weight > 0 and (
predict_actions_t is None or bool(predict_actions_t.any().item())
)
run_text = self.config.text_loss_weight > 0 and text_labels is not None
@@ -1330,13 +1312,24 @@ class PI052Policy(PreTrainedPolicy):
num_repeats = int(getattr(self.config, "flow_num_repeats", 1))
if num_repeats > 1:
prefix_out, flow_loss = self._amortized_prefix_and_flow(
actions, prefix_embs, prefix_pad, prefix_att,
non_fast_prefix_len, fast_len, predict_actions_t, num_repeats,
actions,
prefix_embs,
prefix_pad,
prefix_att,
non_fast_prefix_len,
fast_len,
predict_actions_t,
num_repeats,
)
else:
prefix_out, flow_loss = self._combined_prefix_and_flow(
actions, prefix_embs, prefix_pad, prefix_att,
non_fast_prefix_len, fast_len, predict_actions_t,
actions,
prefix_embs,
prefix_pad,
prefix_att,
non_fast_prefix_len,
fast_len,
predict_actions_t,
)
text_loss, fast_loss = self._prefix_ce_losses(
@@ -1371,9 +1364,7 @@ class PI052Policy(PreTrainedPolicy):
suffix_embs, suffix_pad, suffix_att, adarms_cond = self.model.embed_suffix(x_t, time)
# ---- bf16 alignment (mirrors PI05Pytorch.forward) -----------
first_layer = (
self.model.paligemma_with_expert.paligemma.model.language_model.layers[0]
)
first_layer = self.model.paligemma_with_expert.paligemma.model.language_model.layers[0]
if first_layer.self_attn.q_proj.weight.dtype == torch.bfloat16:
suffix_embs = suffix_embs.to(dtype=torch.bfloat16)
prefix_embs = prefix_embs.to(dtype=torch.bfloat16)
@@ -1407,9 +1398,7 @@ class PI052Policy(PreTrainedPolicy):
non_fast_valid = prefix_pad[:, :non_fast_prefix_len].sum(dim=1, keepdim=True)
suffix_pos = non_fast_valid + torch.cumsum(suffix_pad, dim=1) - 1
position_ids = torch.cat([position_ids[:, : prefix_pad.shape[1]], suffix_pos], dim=1)
att_2d_masks_4d = self.model._prepare_attention_masks_4d(
att_2d_masks, dtype=prefix_embs.dtype
)
att_2d_masks_4d = self.model._prepare_attention_masks_4d(att_2d_masks, dtype=prefix_embs.dtype)
# ---- forward (capture BOTH expert outputs) ------------------
(prefix_out, suffix_out), _ = self.model.paligemma_with_expert.forward(
@@ -1422,9 +1411,7 @@ class PI052Policy(PreTrainedPolicy):
)
# ---- flow loss (mirrors PI05Pytorch.forward) ----------------
suffix_out_slice = suffix_out[:, -self.model.config.chunk_size :].to(
dtype=torch.float32
)
suffix_out_slice = suffix_out[:, -self.model.config.chunk_size :].to(dtype=torch.float32)
v_t = self.model.action_out_proj(suffix_out_slice)
flow_per_dim = F.mse_loss(u_t, v_t, reduction="none")
# Truncate to the actual action dimensionality (PI05 pads
@@ -1670,9 +1657,7 @@ class PI052Policy(PreTrainedPolicy):
use_cache=False,
)
if vlm_out is None:
raise RuntimeError(
"PI052 text+fast loss: VLM forward returned no hidden states."
)
raise RuntimeError("PI052 text+fast loss: VLM forward returned no hidden states.")
lm_head = self.model.paligemma_with_expert.paligemma.lm_head
@@ -1682,7 +1667,7 @@ class PI052Policy(PreTrainedPolicy):
# embed_prefix lays out as [images, language]; with FAST
# appended the full sequence is [images, language, FAST].
if fast_len > 0:
text_hidden = vlm_out[:, -(fast_len + lang_len):-fast_len, :]
text_hidden = vlm_out[:, -(fast_len + lang_len) : -fast_len, :]
else:
text_hidden = vlm_out[:, -lang_len:, :]
text_loss = _shifted_lin_ce(
@@ -1693,11 +1678,7 @@ class PI052Policy(PreTrainedPolicy):
)
fast_loss: Tensor | None = None
if (
action_tokens is not None
and action_code_mask is not None
and fast_len > 0
):
if action_tokens is not None and action_code_mask is not None and fast_len > 0:
fast_hidden = vlm_out[:, -fast_len:, :]
fast_loss = _fast_lin_ce(
fast_hidden,
@@ -1714,9 +1695,7 @@ class PI052Policy(PreTrainedPolicy):
# ------------------------------------------------------------------
@torch.no_grad()
def debug_text_predictions(
self, batch: dict[str, Tensor], max_samples: int = 5
) -> dict[str, Tensor]:
def debug_text_predictions(self, batch: dict[str, Tensor], max_samples: int = 5) -> dict[str, Tensor]:
"""Run the text-loss forward but return argmax predictions instead of CE.
Lets a periodic training-loop hook compare what the LM head emits
@@ -1764,10 +1743,7 @@ class PI052Policy(PreTrainedPolicy):
position_ids = torch.cumsum(prefix_pad, dim=1) - 1
att_2d_4d = self.model._prepare_attention_masks_4d(att_2d)
backbone = self.model.paligemma_with_expert
backbone_dtype = (
backbone.paligemma.model.language_model.layers[0]
.self_attn.q_proj.weight.dtype
)
backbone_dtype = backbone.paligemma.model.language_model.layers[0].self_attn.q_proj.weight.dtype
if att_2d_4d.dtype != backbone_dtype:
att_2d_4d = att_2d_4d.to(dtype=backbone_dtype)
@@ -1778,7 +1754,7 @@ class PI052Policy(PreTrainedPolicy):
inputs_embeds=[prefix_embs, None],
use_cache=False,
)
text_hidden = vlm_out[:, -sub_labels.shape[1]:, :]
text_hidden = vlm_out[:, -sub_labels.shape[1] :, :]
lm_head = backbone.paligemma.lm_head
text_logits = lm_head(text_hidden.to(lm_head.weight.dtype))
preds = text_logits.argmax(dim=-1)
@@ -1810,27 +1786,19 @@ class PI052Policy(PreTrainedPolicy):
suppress_loc_tokens: bool = False,
use_kv_cache: bool = True,
) -> str:
"""Generate text continuation from a multimodal prefix.
"""Generate text continuation from a multimodal prefix (used by PI052Runtime).
Consumed by :class:`lerobot.policies.pi052.inference.PI052Runtime`
for the high-level / VQA / memory-update text streams.
``suppress_loc_tokens`` masks PaliGemma's reserved ``<locDDDD>``
ids ([256000, 257024)) to ``-inf`` before sampling. PaliGemma's
pretraining puts heavy first-token mass on these ids for any
``Assistant:`` continuation; with a small fine-tuning text-CE
budget (or aggressive LR decay) the LM head can drift back
toward that prior even when teacher-forced argmax stays at
100%. Callsites that legitimately emit ``<loc>`` (VQA spatial
answers) must keep this ``False``; subtask / memory / plan
generation should pass ``True``.
``suppress_loc_tokens=True`` masks PaliGemma's reserved ``<locDDDD>`` ids
([256000, 257024)) before sampling — the pretraining prior drifts back to
them on small text-CE budgets. Pass ``True`` for subtask/memory/plan,
``False`` for VQA (spatial answers legitimately emit ``<loc>``).
"""
self.eval()
if tokenizer is None:
from transformers import AutoTokenizer # noqa: PLC0415
from .inference.steps import _get_loc_tokenizer # noqa: PLC0415
from .inference.pi052_adapter import _get_loc_tokenizer # noqa: PLC0415
from .text_processor_pi052 import register_paligemma_loc_tokens # noqa: PLC0415
tok_name = getattr(self.config, "tokenizer_name", None) or "google/paligemma-3b-pt-224"
@@ -1840,7 +1808,7 @@ class PI052Policy(PreTrainedPolicy):
special_ids: set[int] = set()
try:
for sid in (tokenizer.all_special_ids or []):
for sid in tokenizer.all_special_ids or []:
if sid is not None:
special_ids.add(int(sid))
except Exception: # noqa: BLE001
@@ -1886,10 +1854,7 @@ class PI052Policy(PreTrainedPolicy):
# fp32 even when the rest is bf16, so ``next(parameters())``
# would land on one of those and we'd skip the cast. q_proj is
# always cast with the rest, so its dtype is the one SDPA sees.
backbone_dtype = (
backbone.paligemma.model.language_model.layers[0]
.self_attn.q_proj.weight.dtype
)
backbone_dtype = backbone.paligemma.model.language_model.layers[0].self_attn.q_proj.weight.dtype
for _ in range(max_new_tokens):
if cache is None:
@@ -1989,7 +1954,7 @@ class PI052Policy(PreTrainedPolicy):
return self._action_queue.popleft()
def _with_low_level_subtask_prompt(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
from .inference.steps import _build_text_batch # noqa: PLC0415
from .inference.pi052_adapter import _build_text_batch # noqa: PLC0415
from .text_processor_pi052 import discretize_state_str # noqa: PLC0415
n = self._batch_size_from_observation(batch)
@@ -2015,20 +1980,10 @@ class PI052Policy(PreTrainedPolicy):
# own task + observation, then stack the per-env prompts into a single
# (n, L) batch for the action expert. This keeps batch_size > 1 correct
# (env i is conditioned on env i's subtask, not a broadcast of env 0).
# Diagnostic toggle (PI052_SUBTASK_USE_TASK=1): skip the learned subtask
# generator and condition the action expert on the raw task text. Isolates
# whether the generator is the eval bottleneck — eval-only, off by default.
import os # noqa: PLC0415
use_task_directly = os.environ.get("PI052_SUBTASK_USE_TASK") == "1"
rows: list[tuple[Tensor, Tensor | None]] = []
tokenizer = None
for i in range(n):
if use_task_directly:
subtask = tasks[i]
self.last_subtasks[i] = subtask
elif regenerate or not self.last_subtasks[i]:
if regenerate or not self.last_subtasks[i]:
obs_i = self._slice_observation(batch, i)
subtask = self._generate_low_level_subtask(obs_i, tasks[i], i)
else:
@@ -2043,27 +1998,6 @@ class PI052Policy(PreTrainedPolicy):
[{"role": "user", "content": content}],
add_generation_prompt=False,
)
if (
os.environ.get("PI052_DEBUG_TENSORS") == "1"
and i == 0
and not getattr(self, "_dbg_prompt_done", False)
):
import logging as _lg # noqa: PLC0415
_tok = text_batch["tokenizer"]
_ids = text_batch["lang_tokens"][0]
_decoded = _tok.decode(_ids.tolist())
_log = _lg.getLogger(__name__)
_log.info("PI052_DEBUG eval low-level content[0]: %r", content)
_log.info("PI052_DEBUG eval decoded prompt[0]: %r", _decoded)
if torch.is_tensor(state_all):
_s = state_all[i].float()
_log.info(
"PI052_DEBUG eval norm state[0]: min=%.3f max=%.3f mean=%.3f | digits=%s",
_s.min().item(), _s.max().item(), _s.mean().item(),
discretize_state_str(state_all[i]),
)
self._dbg_prompt_done = True
rows.append((text_batch["lang_tokens"], text_batch["lang_masks"]))
tokenizer = text_batch["tokenizer"]
@@ -2072,9 +2006,7 @@ class PI052Policy(PreTrainedPolicy):
# Scalar aliases mirror env 0 for back-compat / single-env overlays.
self.last_subtask = self.last_subtasks[0] if self.last_subtasks else None
self.last_subtask_raw = self.last_subtasks_raw[0] if self.last_subtasks_raw else None
self.last_subtask_source = (
self.last_subtasks_source[0] if self.last_subtasks_source else "unset"
)
self.last_subtask_source = self.last_subtasks_source[0] if self.last_subtasks_source else "unset"
out = dict(batch)
out[OBS_LANGUAGE_TOKENS] = tokens
@@ -2082,7 +2014,10 @@ class PI052Policy(PreTrainedPolicy):
return out
def _generate_low_level_subtask(self, obs_i: dict[str, Tensor], task: str, i: int) -> str:
from .inference.steps import _generate_with_policy, _looks_like_gibberish # noqa: PLC0415
from .inference.pi052_adapter import ( # noqa: PLC0415
_generate_with_policy,
looks_like_gibberish as _looks_like_gibberish,
)
msg = ""
if task:
@@ -2163,9 +2098,7 @@ class PI052Policy(PreTrainedPolicy):
return out
@staticmethod
def _stack_token_rows(
rows: list[tuple[Tensor, Tensor | None]], tokenizer: Any
) -> tuple[Tensor, Tensor]:
def _stack_token_rows(rows: list[tuple[Tensor, Tensor | None]], tokenizer: Any) -> tuple[Tensor, Tensor]:
"""Right-pad per-env ``(1, L_i)`` token/mask rows and stack to ``(n, L)``.
Right-padding with a False attention mask matches the training-time
@@ -2264,7 +2197,7 @@ class PI052Policy(PreTrainedPolicy):
# ------------------------------------------------------------------
@classmethod
def from_pretrained(
cls: builtins.type[T],
cls: type[T],
pretrained_name_or_path: str | Path,
*,
config: PreTrainedConfig | None = None,
@@ -2493,13 +2426,9 @@ class PI052Policy(PreTrainedPolicy):
base_lr = float(self.config.optimizer_lr)
groups: list[dict[str, object]] = []
if backbone_params:
groups.append(
{"params": backbone_params, "lr": base_lr * backbone_scale, "name": "backbone"}
)
groups.append({"params": backbone_params, "lr": base_lr * backbone_scale, "name": "backbone"})
if expert_params:
groups.append(
{"params": expert_params, "lr": base_lr * expert_scale, "name": "action_expert"}
)
groups.append({"params": expert_params, "lr": base_lr * expert_scale, "name": "action_expert"})
if head_params:
groups.append({"params": head_params, "lr": base_lr * head_scale, "name": "lm_head"})
# Sanity: a non-trivial head scale that matches no params would silently
@@ -2514,13 +2443,18 @@ class PI052Policy(PreTrainedPolicy):
"PI052Policy LR groups (base=%.3g): backbone=%.3g (×%.3g, n=%d), "
"action_expert=%.3g (×%.3g, n=%d), lm_head=%.3g (×%.3g, n=%d)",
base_lr,
base_lr * backbone_scale, backbone_scale, len(backbone_params),
base_lr * expert_scale, expert_scale, len(expert_params),
base_lr * head_scale, head_scale, len(head_params),
base_lr * backbone_scale,
backbone_scale,
len(backbone_params),
base_lr * expert_scale,
expert_scale,
len(expert_params),
base_lr * head_scale,
head_scale,
len(head_params),
)
return groups
def init_rtc_processor(self):
"""Initialize RTC processor if RTC is enabled in config."""
self.rtc_processor = None
@@ -2558,25 +2492,10 @@ class PI052Policy(PreTrainedPolicy):
f"(batch: {batch.keys()}) (image_features: {self.config.image_features})"
)
# Diagnostic (PI052_DEBUG_TENSORS=1): dump raw + processed image stats
# once, to compare the eval env's image pipeline against training.
import os as _os # noqa: PLC0415
_dbg = _os.environ.get("PI052_DEBUG_TENSORS") == "1" and not getattr(self, "_dbg_img_done", False)
# Preprocess image features present in the batch
for key in present_img_keys:
img = batch[key]
if _dbg and key == present_img_keys[0]:
import logging as _lg # noqa: PLC0415
_r = img.float()
_lg.getLogger(__name__).info(
"PI052_DEBUG raw img[%s] shape=%s dtype=%s min=%.3f max=%.3f mean=%.3f",
key, tuple(img.shape), str(img.dtype), _r.min().item(), _r.max().item(), _r.mean().item(),
)
# Ensure tensor is on the same device as the model
if img.device != device:
img = img.to(device)
@@ -2599,16 +2518,6 @@ class PI052Policy(PreTrainedPolicy):
# Normalize from [0,1] to [-1,1] as expected by siglip
img = img * 2.0 - 1.0
if _dbg and key == present_img_keys[0]:
import logging as _lg # noqa: PLC0415
_p = img.float()
_lg.getLogger(__name__).info(
"PI052_DEBUG processed img[%s] shape=%s min=%.3f max=%.3f mean=%.3f (expect ~[-1,1])",
key, tuple(img.shape), _p.min().item(), _p.max().item(), _p.mean().item(),
)
self._dbg_img_done = True
# from openpi preprocess_observation_pytorch: Convert back to [B, C, H, W] format if it was originally channels-first
if is_channels_first:
img = img.permute(0, 3, 1, 2) # [B, H, W, C] -> [B, C, H, W]
@@ -2633,7 +2542,6 @@ class PI052Policy(PreTrainedPolicy):
actions = pad_vector(batch[ACTION], self.config.max_action_dim)
return actions
@torch.no_grad()
def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor:
"""Predict a chunk of actions given environment observations."""
src/lerobot/policies/pi052/text_processor_pi052.py
+9 -28
View File
@@ -182,22 +182,12 @@ def _sample_indices(value: Any, batch_size: int) -> list[int | None]:
return [int(value)] * batch_size
# ---------------------------------------------------------------------------
# VQA spatial answers → PaliGemma <loc> format (PI052 only)
#
# PaliGemma is pre-trained on detection / pointing with a ``<locNNNN>``
# vocabulary (normalized [0, 1023]). The recipe's bbox / keypoint VQA
# answers are stored as JSON in Qwen2.5-VL's grounding convention:
# **0–1000 normalized coordinates**, NOT pixels. (Verified empirically
# on the published datasets: x and y both span 0..1000 with ~30% of
# values exceeding the camera's pixel dimensions — they're not pixels.)
# Converting to ``<loc>`` is therefore camera-resolution-independent:
# ``loc_idx = round(coord / 1000 * 1023)``. We do the conversion here —
# not in the dataset — so the dataset keeps the raw JSON and stays
# backbone-agnostic.
# ---------------------------------------------------------------------------
# VQA spatial answers → PaliGemma <loc> format (PI052 only).
# Dataset JSON uses Qwen2.5-VL's 0–1000 *normalized* grounding coords (not
# pixels — verified empirically); PaliGemma's <locNNNN> vocab is [0, 1023], so
# ``loc_idx = round(coord / 1000 * 1023)`` is resolution-independent. Converted
# here, not in the dataset, so the raw JSON stays backbone-agnostic.
# The 0–1000 scale Qwen2.5-VL emits for grounding coordinates.
_VQA_COORD_SCALE = 1000.0
@@ -276,10 +266,7 @@ def _vqa_answer_to_loc(answer: dict[str, Any]) -> str | None:
label = str(det.get("label", "")).strip()
if not label:
continue
toks = (
f"{_loc_token(y1)}{_loc_token(x1)}"
f"{_loc_token(y2)}{_loc_token(x2)}"
)
toks = f"{_loc_token(y1)}{_loc_token(x1)}{_loc_token(y2)}{_loc_token(x2)}"
parts.append(f"{label} {toks}")
return " ; ".join(parts) if parts else None
return None
@@ -393,9 +380,7 @@ class PI052TextTokenizerStep(ProcessorStep):
return self._tokenizer
from transformers import AutoTokenizer # noqa: PLC0415
self._tokenizer = register_paligemma_loc_tokens(
AutoTokenizer.from_pretrained(self.tokenizer_name)
)
self._tokenizer = register_paligemma_loc_tokens(AutoTokenizer.from_pretrained(self.tokenizer_name))
return self._tokenizer
# ------------------------------------------------------------------
@@ -517,9 +502,7 @@ class PI052TextTokenizerStep(ProcessorStep):
break
# Append EOS to supervised target turns so the LM head learns to
# stop (the span covers it → it becomes a supervised label).
prompt, spans = _format_messages(
messages, target_indices, getattr(tokenizer, "eos_token", None)
)
prompt, spans = _format_messages(messages, target_indices, getattr(tokenizer, "eos_token", None))
encoded = tokenizer(
prompt,
@@ -627,9 +610,7 @@ def _classify_for_dropout(message: dict[str, Any]) -> str | None:
if isinstance(content, list):
text_parts = [b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text"]
content = " ".join(text_parts)
elif content is None:
return None
elif not isinstance(content, str):
elif content is None or not isinstance(content, str):
return None
s = content.strip()
if s.startswith("Plan:") or s.startswith("Previous plan"):
src/lerobot/policies/pi_gemma.py
+2 -3
View File
@@ -14,11 +14,10 @@
from __future__ import annotations
import copy
from typing import TYPE_CHECKING, Literal
import torch
from torch import Tensor, nn
from torch import nn
from torch.nn import functional as F # noqa: N812
from lerobot.utils.import_utils import _transformers_available
@@ -391,6 +390,7 @@ __all__ = [
# width/depth variant config (renamed from GemmaConfig to avoid clashing with
# transformers' GemmaConfig).
def sdpa_attention_forward(
module,
query: torch.Tensor,
@@ -754,4 +754,3 @@ class PaliGemmaWithExpertModel(
prefix_past_key_values = None
return [prefix_output, suffix_output], prefix_past_key_values
src/lerobot/policies/language_conditioned/__init__.py → src/lerobot/runtime/__init__.py
+2 -4
View File
@@ -12,15 +12,14 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""Generic runtime primitives for language-conditioned policies."""
"""Policy-agnostic high/low-level runtime for language-conditioned policies."""
from .runtime import (
from .language_runtime import (
LanguageConditionedPolicyAdapter,
LanguageConditionedRuntime,
RuntimeState,
Tick,
TickClock,
ToolCall,
VQAResult,
)
@@ -30,6 +29,5 @@ __all__ = [
"RuntimeState",
"Tick",
"TickClock",
"ToolCall",
"VQAResult",
]
src/lerobot/policies/language_conditioned/runtime.py → src/lerobot/runtime/language_runtime.py
-43
View File
@@ -26,14 +26,6 @@ from typing import Any, Protocol
logger = logging.getLogger(__name__)
@dataclass
class ToolCall:
"""A pending runtime tool invocation."""
name: str
arguments: dict[str, Any] = field(default_factory=dict)
@dataclass
class VQAResult:
"""Text answer plus optional parsed spatial payload."""
@@ -51,7 +43,6 @@ class RuntimeState:
language_context: dict[str, str] = field(default_factory=dict)
action_queue: deque[Any] = field(default_factory=deque)
events: set[str] = field(default_factory=set)
pending_tools: list[ToolCall] = field(default_factory=list)
log_lines: list[str] = field(default_factory=list)
mode: str = "action"
stop: bool = False
@@ -64,7 +55,6 @@ class RuntimeState:
"current_plan": ("language_context", "plan"),
"current_subtask": ("language_context", "subtask"),
"current_memory": ("language_context", "memory"),
"tool_calls_pending": ("pending_tools", None),
"events_this_tick": ("events", None),
"_tick": ("tick", None),
}
@@ -148,8 +138,6 @@ class LanguageConditionedPolicyAdapter(Protocol):
user_text: str | None = None,
) -> str: ...
def parse_tool_calls(self, text: str) -> list[ToolCall]: ...
def answer_vqa(
self,
question: str,
@@ -210,7 +198,6 @@ class LanguageConditionedRuntime:
policy_adapter: LanguageConditionedPolicyAdapter
observation_provider: Callable[[], dict[str, Any] | None] | None = None
action_executor: Callable[[Any], None] | None = None
tools: dict[str, Any] = field(default_factory=dict)
event_collector: Callable[[RuntimeState], None] | None = None
chunk_hz: float = 4.0
ctrl_hz: float = 50.0
@@ -271,7 +258,6 @@ class LanguageConditionedRuntime:
self.maybe_handle_user_events()
self.maybe_enqueue_action_chunk(force=force_rates)
self.dispatch_action(force=force_rates)
self.dispatch_tools()
self.state.events.clear()
def _current_observation(self) -> dict[str, Any] | None:
@@ -315,14 +301,6 @@ class LanguageConditionedRuntime:
out = self.policy_adapter.select_text("interjection", observation, self.state, user_text=text)
if not out:
return
calls = self.policy_adapter.parse_tool_calls(out)
for call in calls:
self.state.pending_tools.append(call)
if calls:
self.state.emit("tool_call_pending")
for call in calls:
if call.name == "say" and call.arguments.get("text"):
self.state.log(f" speech: {call.arguments['text']}")
plan = getattr(self.policy_adapter, "plan_from_text", lambda value: value)(out)
if plan:
self.state.set_context("plan", plan, label="plan")
@@ -394,27 +372,6 @@ class LanguageConditionedRuntime:
if latest is not None and self.action_executor is not None:
self.action_executor(latest)
def dispatch_tools(self) -> None:
if not (self.state.take_event("tool_call_pending") or self.state.pending_tools):
return
pending = list(self.state.pending_tools)
self.state.pending_tools = []
for call in pending:
name = call.name if isinstance(call, ToolCall) else (call.get("function") or {}).get("name")
args = (
call.arguments
if isinstance(call, ToolCall)
else (call.get("function") or {}).get("arguments", {})
)
tool = self.tools.get(name)
if tool is None:
self.state.log(f" [warn] tool {name!r} not registered — skipping call")
continue
try:
tool.call(args)
except Exception as exc: # noqa: BLE001
self.state.log(f" [error] tool dispatch failed: {exc}")
def _handle_action_deadline(self) -> None:
deadline = self.state.action_deadline
if self.state.mode == "action" and deadline is not None and time.monotonic() >= deadline:
src/lerobot/scripts/lerobot_eval.py
+1 -15
View File
@@ -52,7 +52,6 @@ You can learn about the CLI options for this script in the `EvalPipelineConfig`
import concurrent.futures as cf
import json
import logging
import os
import threading
import time
from collections import defaultdict
@@ -114,9 +113,7 @@ def _wrap_text_to_width(text: str, cv2, font, scale: int, thickness: int, max_wi
return lines or [""]
def _annotate_eval_frames(
frames: np.ndarray, task: str | None, subtask: str | None
) -> np.ndarray:
def _annotate_eval_frames(frames: np.ndarray, task: str | None, subtask: str | None) -> np.ndarray:
"""Overlay the high-level task and predicted subtask onto rendered frames.
``frames`` is ``(n_envs, H, W, C)`` uint8. Best-effort: if OpenCV isn't
@@ -240,17 +237,6 @@ def rollout(
except (AttributeError, NotImplementedError):
observation["task"] = [""] * env.num_envs
# Diagnostic (EVAL_TASK_OVERRIDE): replace the env task string with a
# fixed hand-written instruction for every env. Isolates whether the
# action head can execute a given phrasing, independent of the env's
# own description. Logs the original once for comparison.
_task_override = os.environ.get("EVAL_TASK_OVERRIDE")
if _task_override:
if step == 0:
logging.info("EVAL_TASK_OVERRIDE active: env task[0]=%r -> %r",
observation["task"][0], _task_override)
observation["task"] = [_task_override] * env.num_envs
# Apply environment-specific preprocessing (e.g., LiberoProcessorStep for LIBERO)
observation = env_preprocessor(observation)
src/lerobot/scripts/lerobot_train.py
+14 -205
View File
@@ -50,7 +50,6 @@ from lerobot.configs import parser
from lerobot.configs.train import TrainPipelineConfig
from lerobot.datasets import (
EpisodeAwareSampler,
WeightedEpisodeAwareSampler,
compute_sampler_state,
make_dataset,
)
@@ -118,14 +117,6 @@ def update_policy(
if sample_weighter is not None:
sample_weights, weight_stats = sample_weighter.compute_batch_weights(batch)
# Diagnostic-only: skip DDP gradient all-reduce to isolate compute vs comms
# in the per-step time. Training is incorrect under this flag; use for probes.
sync_ctx = (
accelerator.no_sync(policy)
if os.environ.get("LEROBOT_DEBUG_NO_GRAD_SYNC") == "1"
else nullcontext()
)
# Let accelerator handle mixed precision
with accelerator.autocast():
if sample_weights is not None:
@@ -151,8 +142,7 @@ def update_policy(
# TODO(rcadene): policy.unnormalize_outputs(out_dict)
# Use accelerator's backward method
with sync_ctx:
accelerator.backward(loss)
accelerator.backward(loss)
# Clip gradients if specified
if grad_clip_norm > 0:
@@ -185,161 +175,6 @@ def update_policy(
return train_metrics, output_dict
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=<step_interval>``. 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)
def _build_vqa_oversample_weights(dataset: Any, target_fraction: float) -> "torch.Tensor | None":
"""Build per-frame sampling weights that oversample VQA-annotated frames.
Scans the dataset's ``language_events`` column for frames carrying a
``vqa``-style annotation and returns a weight tensor (length == total
dataset frames) such that, under multinomial sampling, VQA frames make up
roughly ``target_fraction`` of the training stream.
Returns ``None`` (⇒ fall back to uniform episode-aware sampling) when VQA
frames cannot be detected or there are none.
"""
if not 0.0 < target_fraction < 1.0:
logging.warning(
"vqa_target_fraction must be in (0, 1); got %s — VQA oversampling disabled.",
target_fraction,
)
return None
hf = getattr(dataset, "hf_dataset", None)
if hf is None or "language_events" not in getattr(hf, "column_names", []):
logging.warning("Dataset has no `language_events` column — VQA oversampling disabled.")
return None
events_col = hf["language_events"]
n_frames = len(events_col)
is_vqa = torch.zeros(n_frames, dtype=torch.bool)
for i, rows in enumerate(events_col):
if rows and any((row or {}).get("style") == "vqa" for row in rows):
is_vqa[i] = True
n_vqa = int(is_vqa.sum())
if n_vqa == 0:
logging.warning("No `vqa` annotations found in the dataset — VQA oversampling disabled.")
return None
n_other = n_frames - n_vqa
# Solve target = (n_vqa·w) / (n_vqa·w + n_other) for the VQA weight w.
# Clamp to ≥ 1 so VQA frames are never *down*-weighted below uniform.
weight = (target_fraction * n_other) / ((1.0 - target_fraction) * max(n_vqa, 1))
weight = max(weight, 1.0)
weights = torch.ones(n_frames, dtype=torch.double)
weights[is_vqa] = weight
logging.info(
"VQA oversampling: %d/%d frames carry a `vqa` annotation (%.2f%%); "
"weighting them x%.2f to target ~%.0f%% of the training stream.",
n_vqa,
n_frames,
100.0 * n_vqa / n_frames,
weight,
100.0 * target_fraction,
)
return weights
@parser.wrap()
def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
"""
@@ -632,29 +467,14 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
to_indices = dataset.meta.episodes["dataset_to_index"]
seed = cfg.seed if cfg.seed is not None else 0
# When `vqa_target_fraction` is set, oversample VQA-annotated
# frames via a weighted sampler; otherwise plain episode-aware.
vqa_weights = None
if cfg.vqa_target_fraction is not None:
vqa_weights = _build_vqa_oversample_weights(dataset, cfg.vqa_target_fraction)
if vqa_weights is not None:
sampler = WeightedEpisodeAwareSampler(
from_indices,
to_indices,
vqa_weights,
episode_indices_to_use=dataset.episodes,
drop_n_last_frames=getattr(active_cfg, "drop_n_last_frames", 0),
seed=seed,
)
else:
sampler = EpisodeAwareSampler(
from_indices,
to_indices,
episode_indices_to_use=dataset.episodes,
drop_n_last_frames=getattr(active_cfg, "drop_n_last_frames", 0),
shuffle=True,
seed=seed,
)
sampler = EpisodeAwareSampler(
from_indices,
to_indices,
episode_indices_to_use=dataset.episodes,
drop_n_last_frames=getattr(active_cfg, "drop_n_last_frames", 0),
shuffle=True,
seed=seed,
)
if cfg.resume and step > 0:
# The resume offset depends on the (num_processes, batch_size) that produced `step`, so
# use the values recorded in the checkpoint (falling back to the current ones for older
@@ -851,22 +671,13 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
is_saving_step = step % cfg.save_freq == 0 or step == cfg.steps
is_eval_step = cfg.eval_freq > 0 and step % cfg.eval_freq == 0
# Optional periodic head-prediction dump for the LM head:
# ``LEROBOT_DEBUG_PREDS_EVERY=1000`` prints 5 samples + per-token
# (label, argmax, ✓/✗) every 1000 steps. Cheap diagnostic to see
# whether the text head is actually learning what we expect, vs
# collapsing to a fixed token. Refilling the recipe-sample dump
# budget at the same cadence also redumps the raw input shapes.
# Optional LM-head diagnostic (``LEROBOT_DEBUG_PREDS_EVERY=<steps>``): prints
# per-token (label, argmax) for a few samples to check the text head is learning.
_debug_preds_every = int(os.environ.get("LEROBOT_DEBUG_PREDS_EVERY", "0"))
if _debug_preds_every > 0 and step % _debug_preds_every == 0 and is_main_process:
try:
from lerobot.policies.pi052 import text_processor_pi052 as _tp # noqa: PLC0415
from lerobot.policies.pi052.debug_utils import print_debug_text_predictions # noqa: PLC0415
_tp._DUMPED_SO_FAR = 0
_tp._DUMP_BUDGET = max(_tp._DUMP_BUDGET, 5)
except Exception as exc: # noqa: BLE001
logging.debug("Could not reset PI052 debug dump budget: %s", exc, exc_info=True)
_print_debug_text_predictions(policy, batch, step, n_samples=5)
print_debug_text_predictions(policy, batch, step, n_samples=5)
if is_log_step:
# Collective reduce must run on every rank, before the main-process gate below.
@@ -1043,9 +854,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
# weights too, to a sibling ``<repo_id>-ema`` repo, so both are
# fully loadable and you can benchmark/deploy whichever is better.
# Non-fatal: the live model is already up if this fails.
if ema is not None and not (
not cfg.is_reward_model_training and cfg.policy.use_peft
):
if ema is not None and not (not cfg.is_reward_model_training and cfg.policy.use_peft):
ema_model = ema.ema_model
ema_repo_id = f"{active_cfg.repo_id}-ema"
orig_repo_id = ema_model.config.repo_id
src/lerobot/tools/__init__.py
-29
View File
@@ -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.
"""LeRobot tool implementations.
Storage of the tool catalog (``meta/info.json["tools"]``) and the
``SAY_TOOL_SCHEMA`` constant live in PR 1
(``lerobot.datasets.language``). This package holds the *runnable*
implementations one file per tool, plus the registry that maps tool
names to classes.
See ``docs/source/tools.mdx`` for the authoring guide.
"""
from .base import Tool
from .registry import TOOL_REGISTRY, get_tools
from .say import SayTool
__all__ = ["Tool", "TOOL_REGISTRY", "get_tools", "SayTool"]
src/lerobot/tools/base.py
-58
View File
@@ -1,58 +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.
"""Tool protocol — the contract every runnable tool implementation honors.
Tools are the executable side of the OpenAI-style function-calling
abstraction the v3.1 language schema (PR 1) carries on assistant
messages: the schema describes *what can be called*, the tool
implementation describes *how to call it*.
Implementations live one-per-file under :mod:`lerobot.tools` (e.g.
``say.py`` for ``SayTool``) and are registered in
:mod:`lerobot.tools.registry`. The runtime instantiates them lazily so
heavy dependencies (torch models, audio backends, network clients,
hardware drivers) only load when the dataset actually declares the tool.
"""
from __future__ import annotations
from typing import Any, Protocol, runtime_checkable
@runtime_checkable
class Tool(Protocol):
"""Minimum surface every tool must expose."""
#: Name matching ``schema["function"]["name"]``. The runtime dispatcher
#: routes incoming ``tool_calls`` to the implementation by this key.
name: str
#: OpenAI-style function-call schema. Same dict the dataset stores in
#: ``meta/info.json["tools"]`` and the chat template renders into the
#: prompt.
schema: dict[str, Any]
def call(self, arguments: dict[str, Any]) -> Any:
"""Execute the tool with the model-provided arguments.
``arguments`` is the parsed dict from
``tool_calls[i]["function"]["arguments"]`` (already JSON-decoded
when the model emits a JSON-string by the chat-template
convention). Implementations validate the dict against their own
schema; the runtime only routes by name.
Return value is implementation-defined — typically a tensor
(TTS audio), a Path (saved file), a dict (structured result), or
``None`` (side-effect-only call).
"""
src/lerobot/tools/registry.py
-70
View File
@@ -1,70 +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.
"""Tool registry — name → implementation class.
Adding a new tool:
1. Drop a file under ``src/lerobot/tools/`` that defines a class
conforming to :class:`lerobot.tools.base.Tool` (must expose ``name``,
``schema``, ``call(arguments)``).
2. Register the class here under :data:`TOOL_REGISTRY`.
3. (Optional) Pre-populate ``meta/info.json["tools"]`` on your dataset
to advertise the schema to the chat-template + policy. The PR 2
annotation pipeline preserves anything you put there.
See ``docs/source/tools.mdx`` for the full authoring guide.
"""
from __future__ import annotations
from typing import Any
from .base import Tool
from .say import SayTool
#: Map from ``function.name`` to a class implementing :class:`Tool`.
#: The runtime instantiates entries lazily — registering a tool here is
#: essentially free (no model load happens until ``call`` runs).
TOOL_REGISTRY: dict[str, type] = {
"say": SayTool,
}
def get_tools(meta: Any, **kwargs: Any) -> dict[str, Tool]:
"""Build name → tool-instance dict from a dataset's declared catalog.
``meta`` is anything with a ``.tools`` attribute returning the
OpenAI-style schema list — typically a
:class:`lerobot.datasets.dataset_metadata.LeRobotDatasetMetadata`.
Each entry whose ``function.name`` is registered here is
instantiated with the schema dict; tools whose name is unknown to
the registry are skipped (the schema still rides through the chat
template, the model just can't actually invoke that tool at
inference).
Extra keyword arguments are forwarded to every constructor — useful
for runtime defaults like ``output_dir=Path("./tts_log")``.
"""
declared = list(meta.tools)
instances: dict[str, Tool] = {}
for schema in declared:
try:
name = schema["function"]["name"]
except (KeyError, TypeError):
continue
cls = TOOL_REGISTRY.get(name)
if cls is None:
continue
instances[name] = cls(schema=schema, **kwargs)
return instances
src/lerobot/tools/say.py
-169
View File
@@ -1,169 +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.
"""``SayTool`` — text-to-speech tool wrapping Kyutai's pocket-tts.
The first concrete tool implementation. PI052 and downstream runtime
dispatchers consume this when the model emits an assistant message
with ``tool_calls=[{function: {name: "say", arguments: {text: ...}}}]``.
Why pocket-tts:
- runs on CPU (no GPU dependency); ~6× real-time on a MacBook Air M4
- ~100M parameters, ~200ms first-chunk latency
- streamable, voice-cloneable
- pip-installable, MIT-style permissive license
The pocket-tts model is loaded **lazily** the first time ``call(...)``
runs (or eagerly via ``preload()``). Loading takes a few seconds and
several hundred MB of RAM, so we don't pay the cost when the tool is
merely *registered* — only when it's *invoked*.
Optional dependency. Install with::
pip install lerobot[tools]
# or directly:
pip install pocket-tts
"""
from __future__ import annotations
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
from lerobot.datasets.language import SAY_TOOL_SCHEMA
logger = logging.getLogger(__name__)
@dataclass
class SayTool:
"""Speak a short utterance via Kyutai's pocket-tts.
Parameters
----------
schema:
Optional schema override; defaults to the canonical
``SAY_TOOL_SCHEMA`` from PR 1. Custom voices or extended
argument shapes can pass in a modified schema, but the
implementation only reads ``arguments["text"]``.
voice:
One of the pocket-tts catalog voices (``alba``, ``marius``,
``javert``, ``jean``, ``fantine``, ``cosette``, ``eponine``,
``azelma``) or a path to a ``.wav`` / ``.safetensors`` voice
file for cloning. See the pocket-tts model card for licensing.
output_dir:
If set, every ``call(...)`` writes a ``<timestamp>.wav`` audio
file there in addition to returning the PCM tensor.
``None`` (default) skips disk writes — useful for live
playback paths that hand the tensor directly to a sounddevice
/ WebAudio sink.
"""
schema: dict[str, Any] = field(default_factory=lambda: dict(SAY_TOOL_SCHEMA))
voice: str = "alba"
output_dir: Path | None = None
name: str = field(init=False, default="say")
_model: Any = field(init=False, default=None, repr=False)
_voice_state: Any = field(init=False, default=None, repr=False)
_sample_rate: int = field(init=False, default=24000, repr=False)
# ------------------------------------------------------------------
# Lazy model load
# ------------------------------------------------------------------
def preload(self) -> None:
"""Load the pocket-tts model + voice state into memory.
Optional — ``call(...)`` triggers this automatically on first
invocation. Useful when you want the multi-second load to
happen at startup rather than on the first ``say`` the policy
emits.
"""
if self._model is not None and self._voice_state is not None:
return
try:
from pocket_tts import TTSModel # noqa: PLC0415 (optional dep)
except ImportError as exc: # pragma: no cover (env-dependent)
raise ImportError(
"SayTool requires pocket-tts. Install with `pip install "
"lerobot[tools]` or `pip install pocket-tts`."
) from exc
logger.info("SayTool: loading pocket-tts model + voice=%r", self.voice)
self._model = TTSModel.load_model()
self._voice_state = self._model.get_state_for_audio_prompt(self.voice)
self._sample_rate = int(getattr(self._model, "sample_rate", 24000))
# ------------------------------------------------------------------
# Tool protocol
# ------------------------------------------------------------------
def call(self, arguments: dict[str, Any]) -> Any:
"""Speak ``arguments["text"]`` and return the PCM tensor.
Optionally also writes ``<output_dir>/<timestamp>.wav`` when
``self.output_dir`` is set. The returned tensor is a 1-D
``torch.Tensor`` of float32 PCM samples at
``self.sample_rate`` Hz — directly playable by
``sounddevice.play(audio.numpy(), self.sample_rate)`` or
encodable by ``scipy.io.wavfile.write``.
"""
text = arguments.get("text")
if not isinstance(text, str) or not text.strip():
raise ValueError(
f"SayTool.call expects arguments={{'text': str}}, got {arguments!r}"
)
self.preload()
audio = self._model.generate_audio(self._voice_state, text)
if self.output_dir is not None:
self._write_wav(audio, text)
return audio
@property
def sample_rate(self) -> int:
"""PCM sample rate of the returned tensor (Hz)."""
return self._sample_rate
# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------
def _write_wav(self, audio: Any, text: str) -> Path:
"""Write a ``.wav`` next to ``output_dir`` for offline inspection."""
import time as _time # noqa: PLC0415
try:
import scipy.io.wavfile # noqa: PLC0415
except ImportError as exc: # pragma: no cover
raise ImportError(
"SayTool.output_dir requires scipy. `pip install scipy`."
) from exc
out_dir = Path(self.output_dir)
out_dir.mkdir(parents=True, exist_ok=True)
# One file per call; suffix with a millisecond timestamp + a
# short text snippet so a directory listing is informative.
snippet = "".join(c if c.isalnum() else "_" for c in text[:32]).strip("_")
ts_ms = int(_time.time() * 1000)
path = out_dir / f"say_{ts_ms}_{snippet}.wav"
# ``audio`` is a torch tensor; pocket-tts uses CPU, so a plain
# ``.numpy()`` is safe.
scipy.io.wavfile.write(path, self.sample_rate, audio.numpy())
return path
tests/datasets/test_sampler.py
+1 -47
View File
@@ -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, WeightedEpisodeAwareSampler, compute_sampler_state
from lerobot.datasets.sampler import EpisodeAwareSampler, compute_sampler_state
def calculate_episode_data_index(hf_dataset: Dataset) -> dict[str, torch.Tensor]:
@@ -152,52 +152,6 @@ def test_partial_episode_drop_warns(caplog):
assert "Episode 0" in caplog.text
# --- WeightedEpisodeAwareSampler --------------------------------------------
def test_weighted_sampler_respects_episode_drop_and_length():
"""The episode-boundary frame filtering is applied before weighting,
and one epoch still yields ``len(indices)`` samples."""
# One episode, 10 frames; drop the last 2.
sampler = WeightedEpisodeAwareSampler([0], [10], frame_weights=torch.ones(10), drop_n_last_frames=2)
assert sampler.indices == list(range(8))
assert len(sampler) == 8
draws = list(sampler)
assert len(draws) == 8
# Dropped frames 8 and 9 must never be sampled.
assert all(d in set(range(8)) for d in draws)
def test_weighted_sampler_oversamples_high_weight_frames():
"""A heavily-weighted frame dominates the draws."""
torch.manual_seed(0)
# 100 frames, frame 7 is weighted 1000x.
weights = torch.ones(100)
weights[7] = 1000.0
sampler = WeightedEpisodeAwareSampler([0], [100], frame_weights=weights)
counts = {}
for _ in range(20): # 20 epochs
for d in sampler:
counts[d] = counts.get(d, 0) + 1
total = sum(counts.values())
# Frame 7 should be the overwhelming majority of the 2000 draws.
assert counts.get(7, 0) / total > 0.9
def test_weighted_sampler_zero_weights_fall_back_to_uniform():
"""If every surviving frame has zero weight, sampling is uniform
rather than crashing."""
sampler = WeightedEpisodeAwareSampler([0], [6], frame_weights=torch.zeros(6))
draws = set(sampler)
assert draws.issubset(set(range(6)))
assert len(list(sampler)) == 6
def test_weighted_sampler_rejects_short_weight_vector():
with pytest.raises(ValueError, match="frame_weights"):
WeightedEpisodeAwareSampler([0], [10], frame_weights=torch.ones(5))
# --- seeded (seed, epoch) shuffling, resume, and state ---
EPISODE_BOUNDS = ([0, 2, 3], [2, 3, 6]) # episodes of 2, 1 and 3 frames
tests/policies/pi052/test_pi052_runtime_adapter.py
+2 -5
View File
@@ -1,7 +1,7 @@
from types import SimpleNamespace
from lerobot.policies.language_conditioned import RuntimeState
from lerobot.policies.pi052.inference.pi052_adapter import PI052PolicyAdapter, split_plan_and_say
from lerobot.runtime import RuntimeState
def test_pi052_adapter_builds_recipe_prompts_from_runtime_state():
@@ -28,13 +28,11 @@ def test_pi052_adapter_builds_recipe_prompts_from_runtime_state():
]
def test_pi052_adapter_parses_say_tool_calls_and_plan_text():
def test_pi052_adapter_strips_say_markers_from_plan_text():
adapter = PI052PolicyAdapter(policy=object())
text = "Move to the sink. <say>heading to the sink</say>"
assert split_plan_and_say(text) == ("Move to the sink.", "heading to the sink")
assert adapter.parse_tool_calls(text)[0].name == "say"
assert adapter.parse_tool_calls(text)[0].arguments == {"text": "heading to the sink"}
assert adapter.plan_from_text(text) == "Move to the sink."
@@ -48,7 +46,6 @@ def test_pi052_runtime_cli_smoke_does_not_load_model(monkeypatch):
"_load_policy_and_preprocessor",
lambda policy_path, dataset_repo_id: (fake_policy, None, None, None),
)
monkeypatch.setattr(runtime_cli, "_build_tools", lambda no_tts, tts_voice: {})
monkeypatch.setattr(runtime_cli, "_run_repl", lambda runtime, initial_task, max_ticks: 0)
assert runtime_cli.main(["--policy.path=fake", "--no_robot", "--task=clean", "--max_ticks=0"]) == 0
tests/processor/test_render_messages_processor.py
+21 -1
View File
@@ -12,7 +12,9 @@ from lerobot.processor.render_messages_processor import RenderMessagesStep # no
from lerobot.types import TransitionKey # noqa: E402
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 +23,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
tests/policies/language_conditioned/test_runtime.py → tests/runtime/test_language_runtime.py
+5 -22
View File
@@ -1,7 +1,6 @@
from lerobot.policies.language_conditioned import (
from lerobot.runtime import (
LanguageConditionedRuntime,
RuntimeState,
ToolCall,
VQAResult,
)
@@ -18,11 +17,7 @@ class FakeAdapter:
def select_text(self, kind, observation, state, user_text=None):
self.text_calls.append((kind, user_text))
return "new plan <say>ok</say>"
def parse_tool_calls(self, text):
assert text == "new plan <say>ok</say>"
return [ToolCall("say", {"text": "ok"})]
return "new plan"
def answer_vqa(self, question, camera, observation, state):
return VQAResult(answer=f"answer: {question}")
@@ -32,14 +27,6 @@ class FakeAdapter:
state.set_context("subtask", "pick cup", label="subtask")
class FakeTool:
def __init__(self):
self.calls = []
def call(self, args):
self.calls.append(args)
def test_runtime_tick_updates_language_enqueues_and_dispatches_action():
adapter = FakeAdapter()
executed = []
@@ -59,24 +46,20 @@ def test_runtime_tick_updates_language_enqueues_and_dispatches_action():
assert " subtask: pick cup" in logs
def test_runtime_handles_user_interjection_and_dispatches_tools():
def test_runtime_handles_user_interjection():
adapter = FakeAdapter()
tool = FakeTool()
runtime = LanguageConditionedRuntime(
policy_adapter=adapter,
observation_provider=lambda: {"observation.state": 1},
tools={"say": tool},
)
runtime.set_task("clean")
runtime.state.extra["recent_interjection"] = "please say ok"
runtime.state.emit("user_interjection")
logs = runtime.step_once()
runtime.step_once()
assert ("interjection", "please say ok") in adapter.text_calls
assert runtime.state.language_context["plan"] == "new plan <say>ok</say>"
assert tool.calls == [{"text": "ok"}]
assert " speech: ok" in logs
assert runtime.state.language_context["plan"] == "new plan"
def test_runtime_state_aliases_legacy_keys_to_language_context():
uv.lock
Generated
+1 -52
View File
@@ -424,15 +424,6 @@ dependencies = [
]
sdist = { url = "https://files.pythonhosted.org/packages/5c/37/0211f82891a9f14efcfd2b2096f8d9e4351398ad637fdd1ee59cfc580b0e/bddl-1.0.1.tar.gz", hash = "sha256:1fa4e6e5050b93888ff6fd8455c39bfb29d3864ce06b4c37c0f781f513a2ae26", size = 164809, upload-time = "2022-03-08T01:48:23.564Z" }
[[package]]
name = "beartype"
version = "0.22.9"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/c7/94/1009e248bbfbab11397abca7193bea6626806be9a327d399810d523a07cb/beartype-0.22.9.tar.gz", hash = "sha256:8f82b54aa723a2848a56008d18875f91c1db02c32ef6a62319a002e3e25a975f", size = 1608866, upload-time = "2025-12-13T06:50:30.72Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/71/cc/18245721fa7747065ab478316c7fea7c74777d07f37ae60db2e84f8172e8/beartype-0.22.9-py3-none-any.whl", hash = "sha256:d16c9bbc61ea14637596c5f6fbff2ee99cbe3573e46a716401734ef50c3060c2", size = 1333658, upload-time = "2025-12-13T06:50:28.266Z" },
]
[[package]]
name = "beautifulsoup4"
version = "4.15.0"
@@ -3160,10 +3151,6 @@ test = [
{ name = "pytest-cov" },
{ name = "pytest-timeout" },
]
tools = [
{ name = "pocket-tts" },
{ name = "scipy" },
]
topreward = [
{ name = "transformers" },
]
@@ -3388,7 +3375,6 @@ requires-dist = [
{ name = "peft", marker = "extra == 'peft-dep'", specifier = ">=0.18.0,<1.0.0" },
{ name = "pillow", specifier = ">=10.0.0,<13.0.0" },
{ name = "placo", marker = "extra == 'placo-dep'", specifier = ">=0.9.6,<0.9.16" },
{ name = "pocket-tts", marker = "extra == 'tools'", specifier = ">=1.0.0,<3.0.0" },
{ name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.7.0,<5.0.0" },
{ name = "protobuf", marker = "extra == 'grpcio-dep'", specifier = ">=6.31.1,<8.0.0" },
{ name = "protobuf", marker = "extra == 'reachy2'", specifier = "<=6.32.0" },
@@ -3414,7 +3400,6 @@ requires-dist = [
{ name = "scikit-image", marker = "extra == 'video-benchmark'", specifier = ">=0.23.2,<0.26.0" },
{ name = "scipy", marker = "extra == 'all'", specifier = ">=1.14.0,<2.0.0" },
{ name = "scipy", marker = "extra == 'scipy-dep'", specifier = ">=1.14.0,<2.0.0" },
{ name = "scipy", marker = "extra == 'tools'", specifier = ">=1.11.0,<2.0.0" },
{ name = "sentencepiece", marker = "extra == 'sentencepiece-dep'", specifier = ">=0.2.0,<0.3.0" },
{ name = "setuptools", specifier = ">=71.0.0,<81.0.0" },
{ name = "teleop", marker = "extra == 'phone'", specifier = ">=0.1.0,<0.2.0" },
@@ -3432,7 +3417,7 @@ requires-dist = [
{ name = "transformers", marker = "extra == 'transformers-dep'", specifier = ">=5.4.0,<5.6.0" },
{ name = "wandb", marker = "extra == 'training'", specifier = ">=0.24.0,<0.28.0" },
]
provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "sentencepiece-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "hilserl", "vla-jepa", "async", "peft", "annotations", "tools", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"]
provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "sentencepiece-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "hilserl", "vla-jepa", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"]
[[package]]
name = "librt"
@@ -4839,33 +4824,6 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
]
[[package]]
name = "pocket-tts"
version = "2.1.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "beartype" },
{ name = "einops" },
{ name = "fastapi" },
{ name = "huggingface-hub" },
{ name = "numpy" },
{ name = "pydantic" },
{ name = "python-multipart" },
{ name = "requests" },
{ name = "safetensors" },
{ name = "scipy" },
{ name = "sentencepiece" },
{ name = "torch", version = "2.11.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux'" },
{ name = "torch", version = "2.11.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux'" },
{ name = "typer" },
{ name = "typing-extensions" },
{ name = "uvicorn" },
]
sdist = { url = "https://files.pythonhosted.org/packages/f9/2c/7445f57163bb40e2b2fab4df70d18a4216c4965cdf74196344d95859fc07/pocket_tts-2.1.0.tar.gz", hash = "sha256:6f244f445413400f686506f5ccfb75048547caab7b455b927f4a854c551c60a8", size = 642108, upload-time = "2026-05-04T14:00:29.207Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/cf/63/d16958d388efee3f0fc7287e1418ed652ddbc2b61ff4f581f0ad0abce625/pocket_tts-2.1.0-py3-none-any.whl", hash = "sha256:7b8f01d3e52aa7df84887b711994586bdc875e024a8b40a15f757feeeb29f752", size = 68096, upload-time = "2026-05-04T14:00:27.547Z" },
]
[[package]]
name = "pre-commit"
version = "4.6.0"
@@ -5539,15 +5497,6 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/27/be/0631a861af4d1c875f096c07d34e9a63639560a717130e7a87cbc82b7e3f/python_json_logger-4.1.0-py3-none-any.whl", hash = "sha256:132994765cf75bf44554be9aa49b06ef2345d23661a96720262716438141b6b2", size = 15021, upload-time = "2026-03-29T04:39:55.266Z" },
]
[[package]]
name = "python-multipart"
version = "0.0.32"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/5b/42/55c32bb9b12693c092ad250a0e82edb5b31ddeda6eb772de5f308b3804ad/python_multipart-0.0.32.tar.gz", hash = "sha256:be54b7f3fa167bb83e4fcd936b887b708f4e57fe75911c02aebf53efaf8d938e", size = 46881, upload-time = "2026-06-04T16:18:58.647Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/e1/04/e8135ebd1ad02c56ec633277529b2602ff99ff634be76cdba5744cf554fd/python_multipart-0.0.32-py3-none-any.whl", hash = "sha256:ff6d3f776f16878c894e52e107296ffc890e913c611b1a4ec6c44e2821fe2e23", size = 30042, upload-time = "2026-06-04T16:18:57.319Z" },
]
[[package]]
name = "python-xlib"
version = "0.33"