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

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>

pyproject.toml

@@ -250,14 +250,6 @@ annotations = [
     # install it locally only if you run your own ``vllm serve``.
 ]
 
-# Tool implementations under src/lerobot/tools/. Each tool's dependencies
-# are isolated so adding a new tool doesn't bloat the base install.
-# Currently only `say` (Kyutai pocket-tts; CPU-only, ~100M params).
-tools = [
-    "pocket-tts>=1.0.0,<3.0.0",
-    "scipy>=1.11.0,<2.0.0",  # SayTool.output_dir uses scipy.io.wavfile
-]
-
 # Development
 dev = ["pre-commit>=3.7.0,<5.0.0", "debugpy>=1.8.1,<1.9.0", "lerobot[grpcio-dep]", "grpcio-tools>=1.73.1,<2.0.0", "mypy>=1.19.1", "ruff>=0.14.1", "lerobot[notebook]"]
 notebook = ["jupyter>=1.0.0,<2.0.0", "ipykernel>=6.0.0,<7.0.0"]

src/lerobot/configs/train.py

@@ -114,14 +114,6 @@ class TrainPipelineConfig(HubMixin):
     ema: EMAConfig = field(default_factory=EMAConfig)
     peft: PeftConfig | None = None
 
-    # VQA oversampling. When set (a fraction in (0, 1)), the training
-    # dataloader uses a WeightedEpisodeAwareSampler that draws frames
-    # carrying a `vqa` language annotation often enough that they make
-    # up roughly this fraction of the training stream. VQA annotations
-    # are typically sparse, so without this they are underrepresented.
-    # `None` (default) keeps uniform episode-aware sampling.
-    vqa_target_fraction: float | None = None
-
     # Sample weighting configuration (e.g., for RA-BC training). Old
     # inline ``use_rabc`` / ``rabc_*`` params are migrated to this
     # field by ``_migrate_legacy_rabc_keys`` above.

src/lerobot/datasets/__init__.py

@@ -49,7 +49,7 @@ from .lerobot_dataset import LeRobotDataset
 from .multi_dataset import MultiLeRobotDataset
 from .pipeline_features import aggregate_pipeline_dataset_features, create_initial_features
 from .pyav_utils import check_video_encoder_parameters_pyav, detect_available_encoders_pyav
-from .sampler import EpisodeAwareSampler, WeightedEpisodeAwareSampler, compute_sampler_state
+from .sampler import EpisodeAwareSampler, compute_sampler_state
 from .streaming_dataset import StreamingLeRobotDataset
 from .utils import DEFAULT_EPISODES_PATH, create_lerobot_dataset_card
 from .video_utils import VideoEncodingManager
@@ -77,7 +77,6 @@ __all__ = [
     "DEFAULT_QUANTILES",
     "EVENT_ONLY_STYLES",
     "EpisodeAwareSampler",
-    "WeightedEpisodeAwareSampler",
     "LANGUAGE_EVENTS",
     "LANGUAGE_PERSISTENT",
     "LeRobotDataset",

src/lerobot/datasets/sampler.py

@@ -154,81 +154,6 @@ class EpisodeAwareSampler:
         return self._num_frames
 
 
-class WeightedEpisodeAwareSampler(EpisodeAwareSampler):
-    """``EpisodeAwareSampler`` that draws frames *with replacement* in
-    proportion to per-frame weights.
-
-    Used to oversample frames carrying a sparse annotation (e.g. a VQA
-    question) so the policy sees them more often than their natural
-    dataset density. One epoch still yields ``len(self.indices)``
-    samples — the weights only change the *composition* of the stream,
-    not its length. Each epoch re-draws, so the oversampled subset
-    varies run to run.
-    """
-
-    def __init__(
-        self,
-        dataset_from_indices: list[int],
-        dataset_to_indices: list[int],
-        frame_weights,
-        *,
-        episode_indices_to_use: list | None = None,
-        drop_n_first_frames: int = 0,
-        drop_n_last_frames: int = 0,
-        seed: int = 0,
-    ):
-        """
-        Args:
-            dataset_from_indices: Episode start indices (see ``EpisodeAwareSampler``).
-            dataset_to_indices: Episode end indices.
-            frame_weights: 1-D sequence/tensor of non-negative weights, one per
-                dataset frame (length == total dataset frames). Higher weight ⇒
-                that frame is sampled more often.
-            episode_indices_to_use / drop_n_first_frames / drop_n_last_frames:
-                Same meaning as ``EpisodeAwareSampler`` — the episode-boundary
-                frame filtering is applied first, then weighting is restricted
-                to the surviving frames.
-        """
-        super().__init__(
-            dataset_from_indices,
-            dataset_to_indices,
-            episode_indices_to_use=episode_indices_to_use,
-            drop_n_first_frames=drop_n_first_frames,
-            drop_n_last_frames=drop_n_last_frames,
-            shuffle=False,
-            seed=seed,
-        )
-        weights = torch.as_tensor(frame_weights, dtype=torch.double).flatten()
-        idx = torch.tensor(self.indices, dtype=torch.long)
-        if weights.numel() <= int(idx.max()):
-            raise ValueError(
-                f"frame_weights has {weights.numel()} entries but the sampler "
-                f"references frame index {int(idx.max())}."
-            )
-        selected = weights[idx]
-        if not torch.isfinite(selected).all() or bool((selected < 0).any()):
-            raise ValueError("frame_weights must be finite and non-negative.")
-        if float(selected.sum()) <= 0.0:
-            # All surviving frames have zero weight — fall back to uniform.
-            selected = torch.ones_like(selected)
-        self._weights = selected
-        self._indices = idx
-
-    def __iter__(self) -> Iterator[int]:
-        epoch, start = self._epoch, self._start_index
-        self._epoch += 1
-        self._start_index = 0
-        generator = self._epoch_generator(epoch)
-        picks = torch.multinomial(
-            self._weights,
-            num_samples=self._num_frames,
-            replacement=True,
-            generator=generator,
-        )
-        for i in picks[start:].tolist():
-            yield int(self._indices[i])
-
-
 def compute_sampler_state(step: int, num_frames: int, batch_size: int, num_processes: int) -> dict:
     """Map an optimization step to an `EpisodeAwareSampler` state for sample-exact resume.
 

src/lerobot/datasets/utils.py

@@ -26,7 +26,6 @@ import numpy as np
 import packaging.version
 import torch
 from huggingface_hub import DatasetCard, DatasetCardData, HfApi
-from huggingface_hub.errors import RevisionNotFoundError
 
 from lerobot.utils.utils import flatten_dict, unflatten_dict
 

src/lerobot/policies/factory.py

@@ -96,12 +96,8 @@ def _restore_pi052_pretrained_state(
 
     base = Path(pretrained_path)
     if not base.exists():
-        # ``pretrained_path`` may be a HF Hub repo id rather than a local dir.
+        # Hub repo id, not a local dir: fetch the processor JSON + stats here
-        # ``from_pretrained`` downloads the model weights, but pi052 builds its
+        # (the generic loader never does for pi052's fresh-built processors).
-        # 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.
         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
+        # pi052 pipelines don't JSON-roundtrip — rebuild fresh and transplant
-        # ``policy_preprocessor.json``: ``RenderMessagesStep`` holds a
+        # saved state (see ``_restore_pi052_pretrained_state`` for why).
-        # 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.
         from .pi052.processor_pi052 import make_pi052_pre_post_processors
 
         preprocessor, postprocessor = make_pi052_pre_post_processors(

src/lerobot/policies/pi052/configuration_pi052.py

@@ -56,11 +56,9 @@ class PI052Config(PI05Config):
 
     # Recipe / language stack ---------------------------------------------
     recipe_path: str | None = "recipes/subtask_mem.yaml"
-    """Path (absolute or relative to ``src/lerobot/configs/``) to a
+    """``TrainingRecipe`` YAML path (absolute or relative to
-    ``TrainingRecipe`` YAML. Defaults to the compact Hi-Robot blend
+    ``src/lerobot/configs/``). ``None`` disables recipe rendering — unannotated
-    shipped with this policy. Set to ``None`` to disable recipe
+    datasets fall back to π0.5's plain ``Task: ... Action:`` prompt."""
-    rendering and fall back to π0.5's single-task ``Task: ... Action:``
-    prompt path (unannotated datasets keep working that way)."""
 
     apply_chat_template: bool = False
     """PaliGemma is *not* chat-pretrained — its tokenizer doesn't ship a

src/lerobot/policies/pi052/debug_utils.py

@@ -0,0 +1,121 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Training-time debug helpers for PI052's language head."""
+
+import logging
+from typing import Any
+
+
+def print_debug_text_predictions(policy: Any, batch: dict[str, Any], step: int, n_samples: int = 5) -> None:
+    """Forward the current batch and print head-argmax vs label per supervised position.
+
+    Opt-in via ``LEROBOT_DEBUG_PREDS_EVERY=<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

@@ -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 "
+            "FAST tokenizer cache hit: %s — re-using fitted tokenizer for dataset=%s base=%s n_samples=%d",
-            "dataset=%s base=%s n_samples=%d",
+            out_dir,
-            out_dir, dataset_repo_id, base_tokenizer_name, n_samples,
+            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 = (
+    is_leader = int(os.environ.get("RANK", "0")) == 0 and int(os.environ.get("LOCAL_RANK", "0")) == 0
-        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 "
+        "FAST tokenizer cache miss — fitting on dataset=%s base=%s n_samples=%d chunk_size=%d → %s",
-        "base=%s n_samples=%d chunk_size=%d → %s",
+        dataset_repo_id,
-        dataset_repo_id, base_tokenizer_name, n_samples, chunk_size, out_dir,
+        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(
+        raise RuntimeError(f"FAST fit: no ``data/chunk-*/file-*.parquet`` shards found under {snap!s}.")
-            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(
+        raise RuntimeError(f"FAST fit: expected ``action`` rows to be 1-D vectors; got shape {acts.shape}.")
-            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

@@ -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

@@ -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

@@ -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

@@ -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 +
+    Language columns are stripped first — the runtime supplies its own
-    state, batched, on the policy's device, normalized) suitable for
+    messages from current state, not the dataset's annotations.
-    ``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.
     """
     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
+    # Optional: replay training's augmentation pipeline so dry-run probes the
-    # that training used, so dry-run sees frames from the augmented
+    # augmented support region — coherent text under jitter means the LM head
-    # support region (not just the unperturbed dataset frames). When
+    # generalized; collapse to "\n" means it memorised unperturbed frames.
-    # 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.
     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``.
+    """Pull task / active plan / memory / subtask at ``start_frame``, so the
-
+    runtime's first prompt matches the canonical training prompts (an OOD
-    The model is heavily memorised on the exact training prompts the
+    prompt makes the model fall back to its dominant mode, VQA JSON spam).
-    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).
     """
     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.
+    """Interactive task picker: numbered menu of dataset tasks (bootstrap task
-
+    as default) plus a custom-input option; plain prompt without a dataset.
-    Behaviour:
+    Non-TTY runs skip the prompt and return the bootstrap task. Returns
-
+    ``None`` when the operator declines (Ctrl-D / empty + no default).
-    * 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).
     """
     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.
+    """Closure reading from the robot each call: ``robot.get_observation()`` →
-
+    ``build_inference_frame`` (state vector + image tensors, batched, on device)
-    Each call: ``robot.get_observation()`` (raw per-joint + per-camera
+    → ``EnvTransition``-wrapped preprocessor (rename, normalise) → flat
-    dict, possibly with scalar floats) → ``build_inference_frame``
+    observation batch for ``select_action`` / ``select_message``.
-    (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.
     """
     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
+    # Camera-key → training (H, W) map from ``ds_features``. Live cameras
-    # ``ds_features``. The training distribution sees frames at the
+    # rarely match the recorded resolution, and a different aspect ratio
-    # recorded resolution (e.g. 480×640); a live Mac/USB camera will
+    # changes resize_with_pad's padding geometry — the flow head tolerates
-    # almost always hand us a different native size (720p / 1080p).
+    # that, but the tightly-supervised LM head goes OOD and collapses.
-    # 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).
     _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:
+        # Resize live frames to the training (H, W) so the downstream
-        # every camera frame the model trained on came from the
+        # resize_with_pad geometry matches what the model saw in training.
-        # 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.
         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
+                # One-shot state-vector print so the operator can eyeball it
-                # it against the dataset's stats. State OOD is a real
+                # against dataset stats (state OOD is a real VLA failure mode).
-                # 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.
                 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

@@ -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

@@ -25,12 +25,8 @@ Two-zone terminal layout:
     └───────────────────────────────────────────────────┘
     > _
 
-The state panel re-renders on every state change. Chat lines are
+Chat lines print above a ``rich.Live`` region (natural scrollback); the
-``console.print``'d above the live region so they accumulate naturally
+state panel re-renders on change, auto-suspending while input is pending.
-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.
 """
 
 from __future__ import annotations

src/lerobot/policies/pi052/inference/vqa.py

@@ -131,19 +131,10 @@ def _loc_to_norm(idx: int) -> float:
 def parse_loc_answer(answer: str) -> dict | None:
     """Parse a PaliGemma ``<loc>``-format spatial VQA answer.
 
-    PI052 trains spatial answers in PaliGemma's native detection
+    Point: ``<label> <locY><locX>``; box: ``<label> <locY0><locX0><locY1><locX1>``;
-    vocabulary, label-first: a point is ``<label> <locY><locX>``, a box
+    multiple boxes joined by `` ; `` (label/loc order irrelevant). Returns
-    is ``<label> <locY0><locX0><locY1><locX1>``, and multiple boxes are
+    ``{"kind", "payload", "normalized": True}`` with [0, 1] coords mirroring the
-    joined by `` ; `` (e.g. ``cube <loc..><loc..><loc..><loc..> ; box
+    JSON shapes (shared overlay code), or ``None`` without ``<loc>`` tokens.
-    <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.
     """
     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``,
+    ``kind`` is a ``VQA_ANSWER_SHAPES`` name or ``"unknown"``; ``<loc>`` answers
-    ``keypoint``, ``count``, ``attribute``, ``spatial``) or ``"unknown"``
+    are tried first. Returns ``None`` when neither format parses.
-    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.
     """
     if not answer or not answer.strip():
         return None

src/lerobot/policies/pi052/modeling_pi052.py

@@ -14,30 +14,16 @@
 
 """π0.5 v2 policy — dual-head training & hierarchical inference.
 
-A thin subclass of :class:`PI05Policy` that:
+π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
-* keeps the PaliGemma ``lm_head`` unfrozen during fine-tuning
+``config.flow_loss_weight``) and :meth:`select_message` for AR text
-  (``PI05Policy`` zeroes / freezes it because it never reads from
+generation. The multi-rate runtime in ``lerobot.policies.pi052.inference``
-  the head; ``PI052Config.unfreeze_lm_head`` flips that),
+(``lerobot-pi052-runtime`` CLI) drives ``predict_action_chunk`` +
-* adds a ``text_loss`` term computed via cross-entropy on
+``select_message``. See :class:`PI052Config` for the knobs.
-  ``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.
 """
 
 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(
+        times = torch.tensor([1.0 + s * dt for s in range(num_steps)], dtype=torch.float32, device=device)
-            [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_4d = self._prepare_attention_masks_4d(full_att_2d_masks, dtype=suffix_embs.dtype)
-            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_targets = torch.where(shift_valid, shift_targets, torch.full_like(shift_targets, -100))
-        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 (
+    ) or (hasattr(self, "gradient_checkpointing") and self.gradient_checkpointing and self.training)
-        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,
+                layer_idx,
-                use_reentrant=False, preserve_rng_state=False,
+                inputs_embeds,
-                paligemma=self.paligemma, gemma_expert=self.gemma_expert,
+                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,
+                layer_idx,
-                paligemma=self.paligemma, gemma_expert=self.gemma_expert,
+                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 "
+                "PI052: knowledge insulation enabled — action→VLM K/V gradients are blocked in attention."
-                "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 = (
+        run_flow = self.config.flow_loss_weight > 0 and (
-            self.config.flow_loss_weight > 0
+            predict_actions_t is None or bool(predict_actions_t.any().item())
-            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,
+                actions,
-                non_fast_prefix_len, fast_len, predict_actions_t, num_repeats,
+                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,
+                actions,
-                non_fast_prefix_len, fast_len, predict_actions_t,
+                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 = (
+        first_layer = self.model.paligemma_with_expert.paligemma.model.language_model.layers[0]
-            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_4d = self.model._prepare_attention_masks_4d(att_2d_masks, dtype=prefix_embs.dtype)
-            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(
+        suffix_out_slice = suffix_out[:, -self.model.config.chunk_size :].to(dtype=torch.float32)
-            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(
+            raise RuntimeError("PI052 text+fast loss: VLM forward returned no hidden states.")
-                "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 (
+        if action_tokens is not None and action_code_mask is not None and fast_len > 0:
-            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(
+    def debug_text_predictions(self, batch: dict[str, Tensor], max_samples: int = 5) -> dict[str, Tensor]:
-        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_dtype = backbone.paligemma.model.language_model.layers[0].self_attn.q_proj.weight.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`
+        ``suppress_loc_tokens=True`` masks PaliGemma's reserved ``<locDDDD>`` ids
-        for the high-level / VQA / memory-update text streams.
+        ([256000, 257024)) before sampling — the pretraining prior drifts back to
-
+        them on small text-CE budgets. Pass ``True`` for subtask/memory/plan,
-        ``suppress_loc_tokens`` masks PaliGemma's reserved ``<locDDDD>``
+        ``False`` for VQA (spatial answers legitimately emit ``<loc>``).
-        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``.
         """
         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_dtype = backbone.paligemma.model.language_model.layers[0].self_attn.q_proj.weight.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:
+            if regenerate or not self.last_subtasks[i]:
-                subtask = tasks[i]
-                self.last_subtasks[i] = subtask
-            elif 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_subtask_source = self.last_subtasks_source[0] if self.last_subtasks_source else "unset"
-            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(
+    def _stack_token_rows(rows: list[tuple[Tensor, Tensor | None]], tokenizer: Any) -> tuple[Tensor, Tensor]:
-        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(
+            groups.append({"params": backbone_params, "lr": base_lr * backbone_scale, "name": "backbone"})
-                {"params": backbone_params, "lr": base_lr * backbone_scale, "name": "backbone"}
-            )
         if expert_params:
-            groups.append(
+            groups.append({"params": expert_params, "lr": base_lr * expert_scale, "name": "action_expert"})
-                {"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 * backbone_scale,
-            base_lr * expert_scale, expert_scale, len(expert_params),
+            backbone_scale,
-            base_lr * head_scale, head_scale, len(head_params),
+            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

@@ -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).
-# 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
-# PaliGemma is pre-trained on detection / pointing with a ``<locNNNN>``
+# ``loc_idx = round(coord / 1000 * 1023)`` is resolution-independent. Converted
-# vocabulary (normalized [0, 1023]). The recipe's bbox / keypoint VQA
+# here, not in the dataset, so the raw JSON stays backbone-agnostic.
-# 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.
-# ---------------------------------------------------------------------------
 
-# 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 = (
+            toks = f"{_loc_token(y1)}{_loc_token(x1)}{_loc_token(y2)}{_loc_token(x2)}"
-                f"{_loc_token(y1)}{_loc_token(x1)}"
-                f"{_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(
+        self._tokenizer = register_paligemma_loc_tokens(AutoTokenizer.from_pretrained(self.tokenizer_name))
-            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(
+        prompt, spans = _format_messages(messages, target_indices, getattr(tokenizer, "eos_token", None))
-            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:
+    elif content is None or not isinstance(content, str):
-        return None
-    elif not isinstance(content, str):
         return None
     s = content.strip()
     if s.startswith("Plan:") or s.startswith("Previous plan"):

src/lerobot/policies/pi_gemma.py

@@ -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/runtime/__init__.py

@@ -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/runtime/language_runtime.py

@@ -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

@@ -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(
+def _annotate_eval_frames(frames: np.ndarray, task: str | None, subtask: str | None) -> np.ndarray:
-    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

@@ -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
+        sampler = EpisodeAwareSampler(
-        # frames via a weighted sampler; otherwise plain episode-aware.
+            from_indices,
-        vqa_weights = None
+            to_indices,
-        if cfg.vqa_target_fraction is not None:
+            episode_indices_to_use=dataset.episodes,
-            vqa_weights = _build_vqa_oversample_weights(dataset, cfg.vqa_target_fraction)
+            drop_n_last_frames=getattr(active_cfg, "drop_n_last_frames", 0),
-        if vqa_weights is not None:
+            shuffle=True,
-            sampler = WeightedEpisodeAwareSampler(
+            seed=seed,
-                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,
-            )
         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:
+        # Optional LM-head diagnostic (``LEROBOT_DEBUG_PREDS_EVERY=<steps>``): prints
-        # ``LEROBOT_DEBUG_PREDS_EVERY=1000`` prints 5 samples + per-token
+        # per-token (label, argmax) for a few samples to check the text head is learning.
-        # (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.
         _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.debug_utils import print_debug_text_predictions  # noqa: PLC0415
-                from lerobot.policies.pi052 import text_processor_pi052 as _tp  # noqa: PLC0415
 
-                _tp._DUMPED_SO_FAR = 0
+            print_debug_text_predictions(policy, batch, step, n_samples=5)
-                _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)
 
         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 (
+            if ema is not None and not (not cfg.is_reward_model_training and cfg.policy.use_peft):
-                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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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/runtime/test_language_runtime.py

@@ -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>"
+        return "new plan"
-
-    def parse_tool_calls(self, text):
-        assert text == "new plan <say>ok</say>"
-        return [ToolCall("say", {"text": "ok"})]
 
     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 runtime.state.language_context["plan"] == "new plan"
-    assert tool.calls == [{"text": "ok"}]
-    assert "  speech: ok" in logs
 
 
 def test_runtime_state_aliases_legacy_keys_to_language_context():

uv.lock

@@ -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"