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

style: compact comments in language runtime

Browse Source
This commit is contained in:
Pepijn
2026-07-15 13:52:52 +02:00
parent 1eed8df1c4
commit 7c125c0028
20 changed files with 205 additions and 917 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/policies/pi052/__init__.py
+2 -14
View File
@@ -12,21 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""π0.5 v2 — full reproduction of the π0.5 paper's hierarchical
inference recipe on lerobot.
"""π0.5 with recipe-driven language supervision and hierarchical inference.
Extends :class:`lerobot.policies.pi05.PI05Policy` with:
* recipe-driven training (PR 1's :class:`RenderMessagesStep`),
* PaliGemma ``lm_head`` cross-entropy on supervised subtask spans
(the "high-level subtask prediction" of the paper, §IV.D),
* AR text generation at inference (:meth:`PI052Policy.select_message`),
* per-component prompt dropout (Pi 0.7 §V.E) for regularising the
text head against missing context at inference.
See ``src/lerobot/configs/recipes/subtask_mem.yaml`` for the compact
training recipe and
``examples/training/pi052_hirobot.slurm`` for the launcher.
PI052 adds supervised PaliGemma text generation, prompt dropout, and autoregressive inference to PI0.5.
"""
from .configuration_pi052 import PI052Config
src/lerobot/policies/pi052/configuration_pi052.py
+14 -111
View File
@@ -12,28 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""π0.5 v2 (with text head) — reproduction of the π0.5 paper's
hierarchical inference recipe.
"""PI0.5 with its PaliGemma text head enabled for hierarchical language/action training.
Same architecture as the existing ``PI05Policy`` (PaliGemma 2B VLM +
~300M Gemma action expert, joint training with FAST tokens during
pre-train and flow matching during post-train), but with the
PaliGemma ``lm_head`` re-enabled so the same model can be supervised
to predict both:
* **subtask strings** at the high level (cross-entropy on the LM
head), and
* **action chunks** at the low level (flow matching on the
action-expert tokens).
This is the dual-head co-training pattern from the paper:
L = H(x, f_θ_text) + α * ‖ω - a - f_θ_action(a_τ, o, ℓ)‖²
with α = 10.0 per § IV.D of arxiv:2504.16054. The π0.5 model splits
inference into a text-prediction step followed by an action-prediction
step, which the multi-rate runtime (``lerobot.runtime``, via the
``lerobot-language-runtime`` CLI) drives at separate rates.
The runtime generates high-level text and low-level flow-matched actions at separate rates.
"""
from dataclasses import dataclass
@@ -47,12 +28,7 @@ from ..pi05.configuration_pi05 import PI05Config
@PreTrainedConfig.register_subclass("pi052")
@dataclass
class PI052Config(PI05Config):
"""π0.5 with the PaliGemma LM head re-enabled for subtask prediction.
Recipe-driven dual-head training: the flow head supervises actions,
the LM head supervises subtask / plan / memory / VQA text. The
flow:text loss split is the milder 5:1 (see ``flow_loss_weight``).
"""
"""PI0.5 configuration for recipe-driven text and action supervision."""
# Recipe / language stack ---------------------------------------------
recipe_path: str | None = "recipes/subtask_mem.yaml"
@@ -67,12 +43,7 @@ class PI052Config(PI05Config):
mirroring how the π0.5 paper's high-level inference samples text
auto-regressively after the prefix."""
# Loss weights --------------------------------------------------------
# Paper §IV.D uses α=10 between the flow and text terms, assuming
# text is a rare auxiliary task. With the recipe stack the flow-only
# `low_level` branch fires on a large share of samples, so α=10
# swamps the LM head and collapses generation into degenerate
# repetition. We use the milder 5:1 split here.
# Balance frequent recipe text supervision against the paper's α=10 flow weight.
text_loss_weight: float = 1.0
"""Weight on the LM-head cross-entropy term. Set to ``0`` to disable
text training entirely (reverts to flow-only / π0.5 behaviour)."""
@@ -90,23 +61,12 @@ class PI052Config(PI05Config):
because it never reads from it. Must be ``True`` for π0.5-style
hierarchical inference."""
# Per-component prompt dropout (Pi0.7 §V.E) ---------------------------
# Randomly drop non-target context messages so the LM head learns
# to handle missing /
# stale plan / memory at inference. Defaults to 0.0 so behaviour
# is identical until explicitly enabled.
# Optional context dropout improves tolerance to missing or stale language state.
plan_dropout_prob: float = 0.0
memory_dropout_prob: float = 0.0
subtask_dropout_prob: float = 0.0
# FAST discrete-action supervision — paper §III.B-C ------------------
# When enabled, actions are *also* tokenised via the FAST tokenizer
# ("physical-intelligence/fast") and supervised with cross-entropy
# on the PaliGemma LM head — exactly as in the paper's pre-training
# objective (Eq. 1 mixes FAST CE + flow MSE + subtask CE). The
# ActionTokenizerProcessorStep is wired into the preprocessor
# pipeline when this flag is set; the loss is computed in
# PI052Policy.forward.
# FAST adds discrete-action CE to the text and flow objectives from paper §III.B-C.
enable_fast_action_loss: bool = True
"""If True, tokenise actions with the FAST tokenizer and add a
cross-entropy loss on the LM head. On by default to match the
@@ -158,76 +118,25 @@ class PI052Config(PI05Config):
"""Number of action chunks to sample for the fit. The FAST paper uses
a few thousand; 1024 is a reasonable default for medium datasets."""
# Knowledge insulation — paper §III.B --------------------------------
# When enabled, gradients from the action expert's flow loss are
# blocked from flowing back into the VLM's K/V projections. This
# prevents the action loss from over-fitting the language backbone
# to robot-specific features. Implemented in ``modeling_pi052`` as
# a per-instance monkey-patch on ``paligemma_with_expert.forward``
# that splits queries into VLM and action halves and ``.detach()``-s
# the VLM K/V tensors used in the action-half's attention.
# Knowledge insulation detaches VLM K/V from action-loss gradients (paper §III.B).
knowledge_insulation: bool = True
"""If True, route every transformer layer through the KI
attention path that blocks action→VLM gradient flow on K/V."""
# Learning-rate defaults --------------------------------------------
# pi052 inherits π0.5's openpi-validated optimizer config (peak LR
# 2.5e-5, cosine→2.5e-6, 1k warmup, AdamW (0.9, 0.95), wd=0.01,
# grad_clip=1.0). The only place pi052 needs to diverge from pi05
# is the LM-head LR multiplier: pi05 has no text supervision so the
# head doesn't get gradients; pi052 always has text supervision
# (subtask / memory / VQA) via the recipe, and under KI the LM head
# only sees gradients on ~30–45% of the batch (the text-CE mask
# share of the recipe). Under aggressive cosine decay this is too
# weak to keep the head pinned, so it drifts back toward PaliGemma's
# pretrained ``<loc>`` first-token bias. 5x is the documented fix
# (see ``PI05Config.lm_head_lr_scale`` docstring); the wiring is
# already in ``PI05Policy.get_optim_params`` — it splits the LM head
# + tied ``embed_tokens`` into their own param group while sharing
# the same cosine lambda, so the 5x ratio is preserved across decay.
# Boost sparse text-head updates while retaining PI0.5's optimizer schedule.
lm_head_lr_scale: float = 5.0
# Separate LRs for the VLM backbone vs the action expert (paper §III.B).
# The backbone is a pretrained PaliGemma; the action expert is trained
# from scratch, so their initialisation scales differ and a single global
# LR under-trains one of them. These multipliers scale the base
# ``optimizer_lr`` for each group; the cosine scheduler applies the same
# lambda to every group so the ratios hold across decay. ``backbone_lr_scale``
# covers the PaliGemma tower (except the LM head / tied embeddings, which keep
# their own ``lm_head_lr_scale``); ``action_expert_lr_scale`` covers the Gemma
# expert plus the action/time projection heads. Defaults of 1.0 reproduce the
# single-LR behaviour (back-compat with existing checkpoints/configs).
# Scale pretrained backbone and new action-expert groups independently; 1.0 preserves legacy behavior.
backbone_lr_scale: float = 1.0
action_expert_lr_scale: float = 1.0
# Amortized flow training (paper §III.B, K_repeat). The VLM/backbone forward
# dominates step cost; to extract more learning signal per VLM pass the action
# expert runs ``flow_num_repeats`` denoising targets per sample, each with an
# independent noise + timestep draw, all attending to the single shared VLM
# prefix. The per-repeat flow losses are averaged, so the backbone gradient
# stays well-scaled. Pairs naturally with ``knowledge_insulation`` (which
# additionally detaches the prefix K/V on the action path), the paper's
# setting — but the amortized path is also correct without it. Set to 1 to
# recover the original single-draw combined forward.
# Reuse each VLM prefix across independent denoising draws; 1 restores single-draw flow.
flow_num_repeats: int = 5
# PaLM-style z-loss on text CE. Penalises the log-partition function
# ``z = log Σ exp(logits)`` drifting away from zero — without it, large-
# vocab models (PaliGemma is 257k) can let ``logsumexp`` grow unbounded
# while CE stays low, because a uniform additive logit bias cancels in
# softmax. PaLM appendix B / Chinchilla report z-loss is essential for
# stable large-vocab CE; it especially helps under ``lm_head_lr_scale=
# 5.0`` which amplifies drift risk on the LM head. ``1e-4`` is the
# commonly cited weight; set 0 to disable entirely.
# PaLM-style z-loss stabilizes large-vocabulary CE; 0 disables it.
text_ce_z_loss_weight: float = 1e-4
# Liger Triton kernels (rope + geglu + layer_norm) are now patched
# unconditionally at model build time — see ``_enable_hf_kernels``
# in ``modeling_pi052``. The patch is process-global, idempotent
# and degrades gracefully if ``liger-kernel`` is missing. Measured
# at -4.5% step time on H100 (bench job 22161421); peak memory
# unchanged. ``fused_linear_cross_entropy`` ships separately via
# ``_shifted_lin_ce`` / ``_fast_lin_ce``.
# Liger patches are optional, process-global, and idempotent.
use_flashrt_fp8_mlp: bool = False
"""Opt-in: swap every Gemma GeGLU MLP (action expert + prefix LM) and the
SigLIP vision MLP to FlashRT fused FP8 kernels (Hugging Face Kernel Hub
@@ -255,10 +164,7 @@ class PI052Config(PI05Config):
checkpoints load instead of raising ``DecodingError: The fields
use_flex_attention are not valid for PI052Config``."""
# Optimizer foreach/fused. pi052 carries these locally because the shared
# PI05Config (kept identical to upstream main) does not define them; the
# checkpoints we train serialize both keys into config.json, so they must
# be valid PI052Config fields and flow into the AdamW preset below.
# Keep serialized PI052 AdamW options local because PI05Config lacks them.
optimizer_foreach: bool | None = False
optimizer_fused: bool | None = True
@@ -275,10 +181,7 @@ class PI052Config(PI05Config):
def __post_init__(self) -> None:
super().__post_init__()
# Backbone needs gradients flowing through the text head when
# we're training it. Override the π0.5 default
# (``train_expert_only=True``) unless the user explicitly opts
# out of text training via ``text_loss_weight=0``.
# Override PI0.5's expert-only default when training text.
if self.text_loss_weight > 0 and self.unfreeze_lm_head:
self.train_expert_only = False
if self.flow_num_repeats < 1:
src/lerobot/policies/pi052/debug_utils.py
+2 -16
View File
@@ -20,22 +20,8 @@ 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.
"""Print supervised text predictions and token accuracy for up to ``n_samples`` rows."""
# Unwrap distributed wrappers that do not proxy custom policy methods.
inner = policy
while hasattr(inner, "module") and not hasattr(inner, "debug_text_predictions"):
inner = inner.module
src/lerobot/policies/pi052/fit_fast_tokenizer.py
+10 -75
View File
@@ -12,27 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""Dataset-specific FAST action tokenizer fitting.
"""Fit and cache a FAST tokenizer for a dataset's action distribution.
The published ``physical-intelligence/fast`` tokenizer is a *universal*
codebook fitted on a heterogeneous mix of robot datasets. Per Pertsch
et al. 2025 (the FAST paper, [64] in the π0.5 paper) and §III.C of
π0.5 itself, the recommended practice is to **finetune the tokenizer on
your specific dataset's action distribution** before training the
policy — same way one would adapt a language tokenizer to a domain
corpus. Without this finetune step, action sequences from your robot
may require more tokens per chunk than necessary, lowering effective
compression and slowing convergence of the action-CE loss.
This module provides a single utility, :func:`fit_fast_tokenizer`,
that does the finetune. The training entry point invokes it
automatically when the policy's ``enable_fast_action_loss`` and
``auto_fit_fast_tokenizer`` flags are both ``True`` and no cached
fitted tokenizer is found at ``fast_tokenizer_cache_dir``.
The fitted tokenizer is saved to
``{cache_dir}/{dataset_hash}_{base_hash}/`` so successive training
runs over the same dataset re-use it.
Training invokes this automatically when FAST loss and automatic fitting are enabled.
"""
from __future__ import annotations
@@ -47,11 +29,7 @@ import numpy as np
logger = logging.getLogger(__name__)
# Marker file the cache-hit check looks for. ``ProcessorMixin.save_pretrained``
# writes ``processor_config.json`` (NOT ``preprocessor_config.json`` —
# that's the image / feature-extractor convention). Centralised here so
# the cache-hit check and the rank-N readiness wait agree on the same
# sentinel.
# ``ProcessorMixin.save_pretrained`` writes this shared cache sentinel.
_CACHE_SENTINEL = "processor_config.json"
@@ -130,14 +108,7 @@ def fit_fast_tokenizer(
)
return str(out_dir)
# DDP-safe fit: only the (local) main process actually fits + saves;
# other ranks poll the cache sentinel until the leader is done.
# Without this guard, all N ranks fit concurrently and race on
# ``save_pretrained`` + ``AutoProcessor.from_pretrained`` (the latter
# copies ``processing_action_tokenizer.py`` into ``HF_MODULES_CACHE``
# and compiles a ``.pyc`` — concurrent writers occasionally produce
# a stale / partial ``.pyc`` and the subsequent ``from .. import
# UniversalActionProcessor`` raises ``AttributeError``.
# Only the local main process writes the tokenizer; other ranks wait on the cache sentinel.
is_leader = int(os.environ.get("RANK", "0")) == 0 and int(os.environ.get("LOCAL_RANK", "0")) == 0
if not is_leader:
timeout_s = 1800.0 # 30 min — covers ~1024-sample fits on cold caches
@@ -164,28 +135,11 @@ def fit_fast_tokenizer(
from transformers import AutoProcessor # 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.
#
# Actions are read straight from the underlying HF dataset's
# ``action`` *column* — never via ``ds[i]``. ``ds[i]`` builds a full
# training item (delta-timestamp expansion + video decode + image
# transforms); a single bad video frame would then throw and, since
# the failure was swallowed at debug level, silently starve the fit
# of every chunk. The action column carries no video, so reading it
# directly is both faster and immune to decode errors.
# Read action columns directly to avoid video decoding and bound memory to sampled episodes.
rng = np.random.default_rng(seed)
actions_buf: list[np.ndarray] = []
# Resolve the dataset's data parquet shards directly, sidestepping
# ``LeRobotDataset(repo_id, episodes=[N])`` which on v3-format
# datasets routes through HF datasets'' split lookup and raises
# ``ValueError: Instruction "train" corresponds to no data!`` for
# every episode (job 22182985 looped through 13,293 skipped episodes
# 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.
# Read v3 parquet shards directly to avoid split lookup failures and repeated metadata parsing.
import pyarrow as _pa # noqa: PLC0415
import pyarrow.parquet as _pq # noqa: PLC0415
from huggingface_hub import snapshot_download # noqa: PLC0415
@@ -195,18 +149,12 @@ def fit_fast_tokenizer(
if not data_files:
raise RuntimeError(f"FAST fit: no ``data/chunk-*/file-*.parquet`` shards found under {snap!s}.")
# Read just the (episode_index, action) columns once across all
# shards. This is the same pattern used elsewhere in the codebase
# for whole-dataset audits and stays under ~2 GB even on 32 k-episode
# / 29 M-frame datasets because the action column is a fixed-length
# float vector.
# Load only episode indices and fixed-width actions across all shards.
tables = [_pq.read_table(f, columns=["episode_index", "action"]) for f in data_files]
table = _pa.concat_tables(tables)
eps = table["episode_index"].to_numpy()
acts_col = table["action"]
# ``action`` may be a fixed-shape ListArray or a 2-D NumericArray;
# ``to_numpy(zero_copy_only=False)`` produces an object array of
# 1-D NumPy actions either way, which we stack into (N, D).
# Normalize Arrow action representations into an (N, D) array.
try:
acts = np.stack(acts_col.to_numpy(zero_copy_only=False)).astype(np.float32)
except Exception: # noqa: BLE001
@@ -215,9 +163,7 @@ def fit_fast_tokenizer(
if acts.ndim != 2:
raise RuntimeError(f"FAST fit: expected ``action`` rows to be 1-D vectors; got shape {acts.shape}.")
# Episode index → slice (start, stop) into ``acts`` along axis 0.
# ``eps`` is monotonically increasing within each parquet shard but
# we make no assumption across shards — sort once and group.
# Sort once because episode order is only guaranteed within each shard.
order = np.argsort(eps, kind="stable")
eps_sorted = eps[order]
boundaries = np.searchsorted(eps_sorted, np.arange(int(eps_sorted.max()) + 2))
@@ -269,18 +215,7 @@ def fit_fast_tokenizer(
eps_visited,
)
# Quantile-normalise per dimension before fitting.
#
# The FAST tokenizer DCT-transforms actions, scales by ``scale`` and
# rounds to integer tokens; the integer *range* must fit the
# codebook (vocab_size, default 1024). Raw motor units (e.g. encoder
# ticks) blow that range up — hence "Vocab size 1024 is too small".
# More importantly, at training time ``ActionTokenizerProcessorStep``
# runs *after* the QUANTILES ``NormalizerProcessorStep``, so it
# encodes normalised actions. Fitting on raw actions would mismatch
# that space. We replicate QUANTILES normalisation here (per-dim
# [q01, q99] → [-1, 1], clipped) so the fit and the training-time
# encode see the same distribution.
# Match training-time quantile normalization so FAST sees the same bounded action space.
flat = actions.reshape(-1, actions.shape[-1])
q01 = np.quantile(flat, 0.01, axis=0)
q99 = np.quantile(flat, 0.99, axis=0)
src/lerobot/policies/pi052/flashrt_fp8.py
+40 -33
View File
@@ -12,18 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""Optional FP8 MLP swap for PI052 using the FlashRT Hugging Face Kernel Hub.
"""Optional FlashRT FP8 MLP kernels for PI052.
Replaces every Gemma GeGLU MLP (action expert + prefix language model) with the
fused ``fp8_geglu_mlp_bf16`` kernel and the SigLIP vision-tower MLP with
``fp8_gelu_mlp_bf16``. Static activation scales are calibrated once on a real
observation; weights are quantized once. This is opt-in and degrades gracefully
to the BF16 path if ``kernels`` or the FlashRT packages are unavailable.
Use:
policy = PI052Policy.from_pretrained(...)
batch = preprocessor(observation) # one representative observation
policy.apply_flashrt_fp8_mlp(batch) # calibrate + swap in place
The opt-in swap calibrates once on a real observation and falls back to BF16 when unavailable.
"""
from __future__ import annotations
@@ -32,7 +23,7 @@ import logging
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.functional as F # noqa: N812
logger = logging.getLogger(__name__)
@@ -43,6 +34,8 @@ def _roundtrip_fp8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
"""Quantize->dequantize an activation through FP8 E4M3 at ``scale`` (f32)."""
q = torch.clamp(x.float() / scale.float(), -_FP8_MAX, _FP8_MAX).to(torch.float8_e4m3fn)
return q.float() * scale.float()
_SWIGLU_REPO = "flashrt/flashrt-fp8-swiglu-ffn"
_GELU_REPO = "flashrt/flashrt-fp8-ffn"
_GEMM_REPO = "flashrt/flashrt-gemm-epilogues"
@@ -75,16 +68,10 @@ class _FlashRTGeGLU(nn.Module):
self.in_features = mlp.gate_proj.weight.shape[1]
device = mlp.gate_proj.weight.device
gate_up = torch.cat([mlp.gate_proj.weight, mlp.up_proj.weight], dim=0).float()
# Fold the preceding RMSNorm weight (1 + w) into the gate/up GEMM and feed
# the kernel channel_scale = 1/(1+w). This is exact (it just moves the
# per-channel (1+w) from the activation to the weight) and is what keeps
# FP8 accurate: the normed activation rms(x) is uniform, while
# rms(x)*(1+w) has per-channel outliers that per-tensor FP8 quantizes
# poorly. Mirrors the FlashRT runtime (norm runs with ones, weight folds
# 1+w). Only the fixed-weight (non-adaptive) norms fold; adaptive-RMSNorm
# layers pass fuse_weight=None (channel_scale = ones).
# Fold fixed RMSNorm weights into the GEMM to avoid FP8 activation outliers.
# Adaptive RMSNorm instead uses an identity channel scale.
if fuse_weight is not None:
f = (1.0 + fuse_weight.detach().float())
f = 1.0 + fuse_weight.detach().float()
gate_up = gate_up * f[None, :]
channel_scale = (1.0 / f).to(torch.bfloat16)
else:
@@ -122,10 +109,17 @@ class _FlashRTGeGLU(nn.Module):
self._calibrate_step(x)
shape = x.shape
flat = x.reshape(-1, self.in_features).to(torch.bfloat16)
x_fp8 = self.quant_ops.channel_scale_quantize_fp8_static_bf16(flat, self.channel_scale, self.input_scale)
x_fp8 = self.quant_ops.channel_scale_quantize_fp8_static_bf16(
flat, self.channel_scale, self.input_scale
)
out = self.ffn_ops.fp8_geglu_mlp_bf16(
x_fp8, self.gate_up_fp8, self.down_fp8,
self.input_scale, self.gate_up_scale, self.hidden_scale, self.down_scale,
x_fp8,
self.gate_up_fp8,
self.down_fp8,
self.input_scale,
self.gate_up_scale,
self.hidden_scale,
self.down_scale,
)
return out.reshape(shape)
@@ -150,7 +144,9 @@ class _FlashRTGeluMLP(nn.Module):
self.register_buffer("down_bias", mlp.fc2.bias.detach().to(torch.bfloat16))
self.register_buffer("input_scale", _static_scale(in_amax, safety).to(device))
self.register_buffer("hidden_scale", _static_scale(hid_amax, safety).to(device))
self.register_buffer("channel_scale", torch.ones(self.in_features, device=device, dtype=torch.bfloat16))
self.register_buffer(
"channel_scale", torch.ones(self.in_features, device=device, dtype=torch.bfloat16)
)
self.safety = safety
self.calibrating = False
self._ia = 0.0
@@ -172,10 +168,19 @@ class _FlashRTGeluMLP(nn.Module):
shape = x.shape
dtype = x.dtype
flat = x.reshape(-1, self.in_features).to(torch.bfloat16)
x_fp8 = self.quant_ops.channel_scale_quantize_fp8_static_bf16(flat, self.channel_scale, self.input_scale)
x_fp8 = self.quant_ops.channel_scale_quantize_fp8_static_bf16(
flat, self.channel_scale, self.input_scale
)
out = self.ffn_ops.fp8_gelu_mlp_bf16(
x_fp8, self.up_fp8, self.up_bias, self.down_fp8, self.down_bias,
self.input_scale, self.up_scale, self.hidden_scale, self.down_scale,
x_fp8,
self.up_fp8,
self.up_bias,
self.down_fp8,
self.down_bias,
self.input_scale,
self.up_scale,
self.hidden_scale,
self.down_scale,
)
return out.reshape(*shape[:-1], self.out_features).to(dtype)
@@ -192,7 +197,9 @@ def _run_forward(policy, batches) -> None:
saved = {name: vars(model).pop(name) for name in ("sample_actions", "forward") if name in vars(model)}
with torch.inference_mode():
for batch in batches:
policy.predict_action_chunk({k: (v.clone() if torch.is_tensor(v) else v) for k, v in batch.items()})
policy.predict_action_chunk(
{k: (v.clone() if torch.is_tensor(v) else v) for k, v in batch.items()}
)
torch.cuda.synchronize()
vars(model).update(saved)
@@ -244,9 +251,8 @@ def apply_fp8_mlp(policy, batch, *, safety: float = 1.05) -> bool:
model = policy.model
calibrating = []
gemma_layers = (
list(model.paligemma_with_expert.gemma_expert.model.layers)
+ list(model.paligemma_with_expert.paligemma.model.language_model.layers)
gemma_layers = list(model.paligemma_with_expert.gemma_expert.model.layers) + list(
model.paligemma_with_expert.paligemma.model.language_model.layers
)
for layer in gemma_layers:
fw = _fixed_norm_weight(layer.post_attention_layernorm)
@@ -267,6 +273,7 @@ def apply_fp8_mlp(policy, batch, *, safety: float = 1.05) -> bool:
logger.info(
"PI052: FlashRT FP8 enabled (%d Gemma + %d SigLIP MLPs).",
len(gemma_layers), len(siglip),
len(gemma_layers),
len(siglip),
)
return True
src/lerobot/policies/pi052/inference/__init__.py
+1 -7
View File
@@ -12,13 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""PI052 bridge to the generic language-conditioned runtime.
The runtime, REPL, and CLI are policy-agnostic and live in
:mod:`lerobot.runtime`. PI052 supplies only :class:`PI052PolicyAdapter`;
the ``lerobot-rollout --language`` entry point wires it into
:func:`lerobot.runtime.cli.run`.
"""
"""PI052 adapter for the policy-agnostic language runtime."""
from .pi052_adapter import PI052PolicyAdapter
src/lerobot/policies/pi052/inference/pi052_adapter.py
+2 -10
View File
@@ -12,13 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""PI052 adapter for the generic language-conditioned runtime.
Supplies only the PI052-specific primitives — acting, text generation,
and prompt templates. The high-level control loop (throttling, output
rejection, the subtask -> memory cascade) is inherited from
:class:`lerobot.runtime.adapter.BaseLanguageAdapter`.
"""
"""PI052 actions and text generation for the generic language runtime."""
from __future__ import annotations
@@ -46,9 +40,7 @@ class PI052PolicyAdapter(BaseLanguageAdapter):
)
subtask = state.language_context.get("subtask") or state.task or ""
# Condition the action expert on subtask + discretized state, matching
# training and lerobot-eval's low-level prompt ("{subtask}, State: {..};").
# Without the state the action expert is off-distribution.
# Match the training prompt by conditioning on both subtask and discretized state.
content = subtask
obs_state = observation.get(OBS_STATE)
if isinstance(obs_state, torch.Tensor) and obs_state.numel() > 0:
src/lerobot/policies/pi052/modeling_pi052.py
+48 -241
View File
@@ -12,15 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""π0.5 v2 policy — dual-head training & hierarchical inference.
"""PI0.5 with joint flow/text training and hierarchical language inference."""
π0.5 with the PaliGemma LM head re-enabled: adds a text CE loss on
``text_labels`` next to the flow loss (L = H(x, f_θ_text) + α·flow, α via
``config.flow_loss_weight``) and :meth:`select_message` for AR text
generation. The multi-rate runtime in ``lerobot.policies.pi052.inference``
(``lerobot-language-runtime`` CLI) drives ``predict_action_chunk`` +
``select_message``. See :class:`PI052Config` for the knobs.
"""
# ruff: noqa: N806, N812
from __future__ import annotations
@@ -54,9 +48,7 @@ from .configuration_pi052 import PI052Config
logger = logging.getLogger(__name__)
# PI0.5 flow-matching model + helpers (pi052-specific). The generic dual-expert
# transformer (PaliGemmaWithExpertModel, sdpa_attention_forward,
# compute_layer_complete, get_gemma_config) lives in lerobot.policies.pi_gemma.
# Generic dual-expert transformer helpers live in ``lerobot.policies.pi_gemma``.
class ActionSelectKwargs(TypedDict, total=False):
@@ -339,9 +331,7 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
# Process language tokens
def lang_embed_func(tokens):
# embed_language_tokens -> Gemma embed_tokens, which is GemmaTextScaledWordEmbedding
# (transformers >=5.4.0): it already multiplies by sqrt(hidden_size) internally. Do NOT
# scale again here or text tokens get double-scaled (~45x) and break alignment.
# GemmaTextScaledWordEmbedding already applies sqrt(hidden_size); do not scale twice.
return self.paligemma_with_expert.embed_language_tokens(tokens)
lang_emb = self._apply_checkpoint(lang_embed_func, tokens)
@@ -402,9 +392,7 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
embs = torch.cat(embs, dim=1)
pad_masks = torch.cat(pad_masks, dim=1)
# The suffix mask is the constant [1, 0, ..., 0]; build it on-device
# rather than via torch.tensor(python_list, device=cuda), which is a
# host->device sync on every denoise step.
# Build the constant suffix mask on-device to avoid a per-step host sync.
n = len(att_masks)
att_masks = torch.zeros(n, dtype=embs.dtype, device=embs.device)
att_masks[0] = 1
@@ -436,13 +424,7 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
att_2d_masks_4d = self._prepare_attention_masks_4d(att_2d_masks, dtype=prefix_embs.dtype)
# Selective AC: rely on the per-layer checkpoint inside
# ``PaliGemmaWithExpertModel.forward`` (which wraps each
# transformer block individually). The previous outer
# ``_apply_checkpoint(forward_func, ...)`` doubled up — it
# re-ran the full backbone forward during backward *and* each
# block's own checkpoint re-ran during that recompute. Pure
# waste with SDPA, which already streams attention activations.
# The model already checkpoints each layer; an outer checkpoint would duplicate recomputation.
(_, suffix_out), _ = self.paligemma_with_expert.forward(
attention_mask=att_2d_masks_4d,
position_ids=position_ids,
@@ -508,9 +490,7 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
dt = -1.0 / num_steps
# 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).
# Precompute timesteps on-device to avoid a host sync per denoising step.
times = torch.tensor([1.0 + s * dt for s in range(num_steps)], dtype=torch.float32, device=device)
x_t = noise
@@ -573,12 +553,7 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks, dtype=suffix_embs.dtype)
self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager" # noqa: SLF001
# The expert forward appends the suffix K/V to the prefix cache in-place
# (GemmaAttention.update runs even with use_cache=False), so each step
# must start from a prefix-only cache. Instead of deep-copying the whole
# cache every step, let it append and crop back to the prefix length
# afterwards (the prefix K/V are read-only, so this is exact and keeps
# the loop a single graph).
# Crop appended suffix K/V after each step instead of copying the read-only prefix cache.
outputs_embeds, _ = self.paligemma_with_expert.forward(
attention_mask=full_att_2d_masks_4d,
position_ids=position_ids,
@@ -595,14 +570,7 @@ 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``),
# so action tokens occupy the top ``_FAST_ACTION_VOCAB_SIZE`` ids below the
# ``fast_skip_tokens`` margin. The upper part collides with the reserved
# ``<loc>`` block; the lower part sits just under it and otherwise leaks into
# generated text as high-codepoint gibberish (the action-trained LM head puts
# heavy mass on these ids), so ``select_message`` masks it.
# FAST tokens occupy the high vocabulary range and must be masked during text generation.
_FAST_ACTION_VOCAB_SIZE = 2048
@@ -645,11 +613,6 @@ def _enable_hf_kernels() -> None:
logger.info("PI052: HF kernels (Liger) enabled — rope, geglu, layer_norm fused.")
# ----------------------------------------------------------------------
# Loss helpers (shared between fused and prefix-only paths)
# ----------------------------------------------------------------------
def _mask_per_sample(per_sample: Tensor, predict_actions_t: Tensor | None) -> Tensor:
"""Mean over samples where ``predict_actions_t`` is True, else over all."""
if predict_actions_t is None:
@@ -683,9 +646,7 @@ def _shifted_lin_ce(
(same ``z²·w`` formula on per-position logsumexp). Setting it
to 0 disables the z-loss term at zero cost.
"""
# Liger is imported lazily so the module still imports on machines
# without liger-kernel — the call site only fires from the training
# forward, which always pulls in the kernel.
# Keep Liger optional until the training path needs it.
from liger_kernel.transformers.fused_linear_cross_entropy import ( # noqa: PLC0415
LigerFusedLinearCrossEntropyLoss,
)
@@ -767,9 +728,7 @@ def _fast_lin_ce(
if predict_actions_t is not None:
sample_mask = predict_actions_t[:, None].expand_as(shift_valid)
shift_valid = shift_valid & sample_mask
# 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.
# Encode the mask with ignore_index to avoid a host sync and preserve graph capture.
shift_targets = torch.where(shift_valid, shift_targets, torch.full_like(shift_targets, -100))
B, T_1, H = shift_hidden.shape
@@ -783,15 +742,7 @@ def _fast_lin_ce(
return loss_fn(lm_head_weight, flat_hidden, flat_labels)
# ----------------------------------------------------------------------
# Knowledge insulation — ported from pi05_full (branch ``feat/add-pi05``)
# ----------------------------------------------------------------------
#
# Per-layer attention that splits the queries into VLM and action
# parts, computing attention for action queries with .detach()'d VLM
# K/V so the action loss's gradient cannot flow back into the VLM's K
# and V projections. Forward output is bit-equivalent to the standard
# layer; backward differs only on the path action_loss → VLM K/V.
# Knowledge insulation keeps the forward equivalent while detaching VLM K/V for action-query gradients.
def _compute_layer_ki(
@@ -805,9 +756,7 @@ def _compute_layer_ki(
):
from transformers.models.gemma import modeling_gemma # noqa: PLC0415
# ``_gated_residual`` is a lerobot helper (adaRMSNorm gated residual),
# not part of HF's ``modeling_gemma``. pi05's own layer code imports
# it from ``pi_gemma`` — mirror that here.
# ``_gated_residual`` is LeRobot's adaRMSNorm helper, not a Transformers symbol.
from ..pi_gemma import _gated_residual # noqa: PLC0415
models = [paligemma.model.language_model, gemma_expert.model]
@@ -865,11 +814,7 @@ def _compute_layer_ki(
mask_for_vlm = attention_mask[:, :, :vlm_len, :]
mask_for_action = attention_mask[:, :, vlm_len:, :]
# ``_prepare_attention_masks_4d`` always returns fp32 (0.0 / -inf
# literals), but PaliGemma weights are bf16 when ``dtype=bfloat16``,
# making q bf16. SDPA's ``scaled_dot_product_attention`` then raises
# "invalid dtype for bias - should match query's dtype". Cast each
# mask slice to the corresponding query dtype right before use.
# SDPA requires each fp32-generated mask slice to match its query dtype.
if mask_for_vlm.dtype != Q_vlm.dtype:
mask_for_vlm = mask_for_vlm.to(dtype=Q_vlm.dtype)
if mask_for_action.dtype != Q_action.dtype:
@@ -1005,10 +950,7 @@ class PI052Policy(PreTrainedPolicy):
name = "pi052"
def __init__(self, config: PI052Config, **kwargs: Any) -> None:
# Patch ops BEFORE the backbone is built (the backbone constructed
# below instantiates the Gemma/Siglip layers we want to swap).
# Always-on — the patch is process-global / idempotent and degrades
# gracefully if liger-kernel is missing.
# Patch before constructing Gemma/SigLIP layers; the operation is optional and idempotent.
_enable_hf_kernels()
# ---- inlined PI05Policy.__init__ ----------------------------------
@@ -1024,18 +966,11 @@ class PI052Policy(PreTrainedPolicy):
self.reset()
# ---- end inlined PI05Policy.__init__ ------------------------------
# ``PI05Policy.__init__`` zeroes the PaliGemma ``lm_head`` and
# freezes a few terminal layers when ``train_expert_only`` is
# the (default) True. We re-enable the head if the user
# wants text supervision.
# Re-enable layers PI0.5 freezes when text supervision is requested.
if config.text_loss_weight > 0 and config.unfreeze_lm_head:
self._unfreeze_lm_head()
# Knowledge insulation: bind a custom ``forward`` on the
# PaliGemmaWithExpertModel instance that uses
# :func:`_compute_layer_ki` for the dual-expert layer pass.
# The bind is per-instance, so this doesn't leak into stock
# ``pi05`` policies that share the same class.
# Bind knowledge insulation per instance so stock PI0.5 policies remain unchanged.
if getattr(config, "knowledge_insulation", False):
backbone = self.model.paligemma_with_expert
backbone._pi052_orig_forward = backbone.forward
@@ -1044,11 +979,7 @@ class PI052Policy(PreTrainedPolicy):
"PI052: knowledge insulation enabled — action→VLM K/V gradients are blocked in attention."
)
# Per-env hierarchical-inference state. Sized lazily on the first
# select_action() call once the batch size (number of parallel envs)
# is known. ``last_subtasks[i]`` is the subtask currently conditioning
# env ``i``'s action expert; scalar ``last_subtask`` mirrors env 0 for
# back-compat (e.g. the eval video overlay).
# Size per-environment inference state lazily; the scalar mirrors env 0 for compatibility.
self.last_subtasks: list[str] | None = None
self.last_subtasks_raw: list[str] | None = None
self.last_subtasks_source: list[str] | None = None
@@ -1089,10 +1020,6 @@ class PI052Policy(PreTrainedPolicy):
return apply_fp8_mlp(self, batch, safety=safety)
# ------------------------------------------------------------------
# Head unfreeze helper
# ------------------------------------------------------------------
def _unfreeze_lm_head(self) -> None:
"""Walk the PaliGemma submodules and re-enable gradients on
``lm_head`` + the immediately preceding norm / last text-model
@@ -1101,10 +1028,7 @@ class PI052Policy(PreTrainedPolicy):
if hasattr(backbone, "lm_head"):
for p in backbone.lm_head.parameters():
p.requires_grad_(True)
# The text model's final norm and last transformer block —
# find them dynamically by walking up from the LM head so we
# don't hard-code module names that may drift across transformers
# versions.
# Discover terminal text layers dynamically across Transformers versions.
text_model = getattr(backbone, "model", None)
text_model = getattr(text_model, "language_model", text_model)
if text_model is None:
@@ -1118,10 +1042,6 @@ class PI052Policy(PreTrainedPolicy):
for p in layers[-1].parameters():
p.requires_grad_(True)
# ------------------------------------------------------------------
# Forward (dual loss: flow + text)
# ------------------------------------------------------------------
def forward(
self,
batch: dict[str, Tensor],
@@ -1137,12 +1057,7 @@ class PI052Policy(PreTrainedPolicy):
text_labels = batch.get("text_labels")
predict_actions_t = batch.get("predict_actions")
# Fall through to PI05Policy only on fully unannotated batches
# (no recipe applied → no routing fields). For recipe-applied
# batches we keep control of the loss dispatch even if all
# samples are text-only — delegating would silently train flow
# on text-only frames (PI05Policy.forward ignores
# ``predict_actions``).
# Delegate only unannotated batches; PI0.5 ignores recipe action-routing masks.
if (
text_labels is None
and predict_actions_t is None
@@ -1150,9 +1065,7 @@ class PI052Policy(PreTrainedPolicy):
):
return self._pi05_flow_forward(batch, reduction=reduction)
# Whether any sample in the batch wants actions predicted. This is a data-dependent branch, so
# it needs a host-side bool (one CUDA sync); compute it once and reuse for both flow and FAST
# instead of syncing twice.
# Compute the host-side action-routing decision once for both flow and FAST.
predict_any = predict_actions_t is None or bool(predict_actions_t.any().item())
run_flow = self.config.flow_loss_weight > 0 and predict_any
run_text = self.config.text_loss_weight > 0 and text_labels is not None
@@ -1180,19 +1093,7 @@ class PI052Policy(PreTrainedPolicy):
if action_tokens is None or action_mask is None or action_code_mask is None:
run_fast = False
# ------------------------------------------------------------
# Dispatch: full fusion when flow is active, otherwise the
# prefix-only text+FAST helper (no suffix forward needed).
#
# Full fusion (flow ON):
# ONE backbone forward with prefix=[images, lang, FAST] +
# suffix=[noisy_actions], suffix→FAST attention masked out.
# All three losses computed from slices of the single output.
#
# Prefix-only fusion (flow OFF, e.g. text-only recipes):
# ONE prefix-only forward, both text + FAST losses computed
# from slices. No suffix forward → cheaper.
# ------------------------------------------------------------
# Flow uses one fused prefix/suffix pass; text-only batches skip the suffix.
if run_flow:
flow_loss, text_loss, fast_loss = self._compute_all_losses_fused(
batch,
@@ -1237,17 +1138,12 @@ class PI052Policy(PreTrainedPolicy):
"nothing to train."
)
# Keep loss components as detached tensors (no CUDA sync here); the training loop converts
# them to python floats only on logging steps (see update_policy's log_metrics gate).
# Keep metrics detached on-device until logging to avoid extra CUDA synchronization.
loss_dict["loss"] = total.detach() if total.dim() == 0 else float("nan")
if reduction == "none":
return total.expand(batch[OBS_LANGUAGE_TOKENS].shape[0]), loss_dict
return total, loss_dict
# ------------------------------------------------------------------
# Text loss
# ------------------------------------------------------------------
def _compute_all_losses_fused(
self,
batch: dict[str, Tensor],
@@ -1282,9 +1178,7 @@ class PI052Policy(PreTrainedPolicy):
)
non_fast_prefix_len = prefix_embs.shape[1] # images + language only
# Causal-mask the supervised text-target span so the text-CE is
# genuine next-token prediction, not a bidirectional copy task
# (see ``_mark_target_span_causal``).
# Make supervised text causal rather than a bidirectional copy task.
if text_labels is not None:
lang_start = non_fast_prefix_len - text_labels.shape[1]
if lang_start >= 0:
@@ -1294,8 +1188,7 @@ class PI052Policy(PreTrainedPolicy):
fast_len = 0
if action_tokens is not None and action_mask is not None:
# embed_language_tokens already applies the Gemma sqrt(hidden) scale (tf>=5.4.0);
# do not scale FAST action tokens again (would double-scale).
# Gemma embedding already applies its hidden-size scale.
fast_emb = self.model.paligemma_with_expert.embed_language_tokens(action_tokens)
fast_len = action_tokens.shape[1]
ones_att = torch.ones(
@@ -1307,12 +1200,7 @@ class PI052Policy(PreTrainedPolicy):
prefix_pad = torch.cat([prefix_pad, action_mask.to(prefix_pad.dtype)], dim=1)
prefix_att = torch.cat([prefix_att, ones_att], dim=1)
# ---- flow: one combined forward, or amortized over K repeats ----
# ``flow_num_repeats == 1`` keeps the single combined [prefix; suffix]
# forward. ``> 1`` runs the VLM prefix once and replays the action
# expert K times against fresh noise/timestep draws, reusing the
# cached prefix KV (paper §III.B). Both return ``prefix_out`` for the
# shared text/FAST CE tail.
# Amortized flow reuses one VLM prefix across fresh denoising targets.
num_repeats = int(getattr(self.config, "flow_num_repeats", 1))
if num_repeats > 1:
prefix_out, flow_loss = self._amortized_prefix_and_flow(
@@ -1373,18 +1261,11 @@ class PI052Policy(PreTrainedPolicy):
suffix_embs = suffix_embs.to(dtype=torch.bfloat16)
prefix_embs = prefix_embs.to(dtype=torch.bfloat16)
# ---- combined attention -------------------------------------
pad_masks = torch.cat([prefix_pad, suffix_pad], dim=1)
att_masks = torch.cat([prefix_att, suffix_att], dim=1)
att_2d_masks = make_att_2d_masks(pad_masks, att_masks)
# Critical: zero out suffix → FAST attention. Without this the
# action expert reads the FAST tokens and trivially decodes
# them back to the same continuous actions it's supposed to
# predict from noise. Cumulative-block attention from
# ``make_att_2d_masks`` doesn't enforce this on its own
# because suffix tokens have a strictly higher cumsum than
# FAST tokens and therefore attend to them by default.
# Block suffix-to-FAST attention to prevent trivial action leakage.
if fast_len > 0:
fast_start = non_fast_prefix_len
fast_end = non_fast_prefix_len + fast_len # = prefix_pad.shape[1]
@@ -1392,19 +1273,12 @@ class PI052Policy(PreTrainedPolicy):
position_ids = torch.cumsum(pad_masks, dim=1) - 1
if fast_len > 0:
# The flow suffix is a PARALLEL action representation to the FAST
# block, not a continuation of it (the two never attend to each
# other). At inference there is no FAST block, so the suffix RoPE
# positions start at the valid image+language count. Match that here
# so flow->prefix relative positions are train==inference; otherwise
# the suffix is offset by n_fast (per-sample, 33-111) and the trained
# head reads the wrong RoPE conditioning at deploy time.
# Position flow parallel to FAST so its RoPE offsets match inference without FAST.
non_fast_valid = prefix_pad[:, :non_fast_prefix_len].sum(dim=1, keepdim=True)
suffix_pos = non_fast_valid + torch.cumsum(suffix_pad, dim=1) - 1
position_ids = torch.cat([position_ids[:, : prefix_pad.shape[1]], suffix_pos], dim=1)
att_2d_masks_4d = self.model._prepare_attention_masks_4d(att_2d_masks, dtype=prefix_embs.dtype)
# ---- forward (capture BOTH expert outputs) ------------------
(prefix_out, suffix_out), _ = self.model.paligemma_with_expert.forward(
attention_mask=att_2d_masks_4d,
position_ids=position_ids,
@@ -1477,17 +1351,13 @@ class PI052Policy(PreTrainedPolicy):
if use_bf16:
s_embs = s_embs.to(dtype=torch.bfloat16)
suffix_blocks.append(s_embs)
# adaRMS time conditioning is per-sample; broadcast it across this
# block's chunk tokens so each block carries its own timestep.
# Broadcast each sample's timestep conditioning across its action chunk.
adarms_blocks.append(adarms[:, None, :].expand(batch_size, chunk, adarms.shape[-1]))
suffix_embs = torch.cat(suffix_blocks, dim=1) # (B, k*chunk, D)
adarms_cond = torch.cat(adarms_blocks, dim=1) # (B, k*chunk, cond_dim)
# ---- block-diagonal attention over [prefix | block_1..k] ----
# Prefix rows keep their own (causal/text) attention and never see the
# action blocks. Each action block attends to the valid prefix (minus
# FAST) and only to itself.
# Each action block attends to the non-FAST prefix and itself, never other blocks.
prefix_att_2d = make_att_2d_masks(prefix_pad, prefix_att) # (B, P, P)
device = prefix_pad.device
prefix_rows = torch.cat(
@@ -1507,11 +1377,7 @@ class PI052Policy(PreTrainedPolicy):
att_2d = torch.cat([prefix_rows, action_rows], dim=1) # (B, P + k*chunk, P + k*chunk)
att_2d_4d = model._prepare_attention_masks_4d(att_2d, dtype=prefix_embs.dtype)
# Positions: prefix as usual; every block restarts at the prefix offset
# (each block is an independent denoising of the same chunk). The flow
# blocks are PARALLEL to the FAST block, not a continuation, so offset by
# the valid image+language count (excluding FAST) — matching inference
# (no FAST block) so flow->prefix RoPE positions are train==inference.
# Restart every independent flow block after the non-FAST prefix to match inference RoPE.
if fast_len > 0:
prefix_offsets = prefix_pad[:, :non_fast_prefix_len].sum(dim=-1)[:, None]
else:
@@ -1565,9 +1431,7 @@ class PI052Policy(PreTrainedPolicy):
text_hidden = prefix_out[:, -(fast_len + lang_len) : -fast_len, :]
else:
text_hidden = prefix_out[:, -lang_len:, :]
# Liger fused linear-CE: skip the explicit ``lm_head(...)``
# materialisation; the kernel multiplies on-the-fly and
# never holds the full (B, T, 257k) logits tensor.
# Liger avoids materializing the full vocabulary logits tensor.
text_loss = _shifted_lin_ce(
text_hidden,
lm_head.weight,
@@ -1619,9 +1483,7 @@ class PI052Policy(PreTrainedPolicy):
images, img_masks, lang_tokens, lang_masks
)
# Causal-mask the supervised text-target span (see
# ``_mark_target_span_causal``) before the FAST tokens are
# appended — same fix as ``_compute_all_losses_fused``.
# Make supervised text causal before appending FAST tokens.
if text_labels is not None:
lang_start = prefix_embs.shape[1] - text_labels.shape[1]
if lang_start >= 0:
@@ -1694,10 +1556,6 @@ class PI052Policy(PreTrainedPolicy):
return text_loss, fast_loss
# ------------------------------------------------------------------
# Diagnostic: forward + argmax for supervised text positions
# ------------------------------------------------------------------
@torch.no_grad()
def debug_text_predictions(self, batch: dict[str, Tensor], max_samples: int = 5) -> dict[str, Tensor]:
"""Run the text-loss forward but return argmax predictions instead of CE.
@@ -1773,10 +1631,6 @@ class PI052Policy(PreTrainedPolicy):
if was_training:
self.train()
# ------------------------------------------------------------------
# select_message — AR text generation at inference
# ------------------------------------------------------------------
def select_message(
self,
batch: dict[str, Tensor],
@@ -1815,7 +1669,7 @@ class PI052Policy(PreTrainedPolicy):
for sid in tokenizer.all_special_ids or []:
if sid is not None:
special_ids.add(int(sid))
except Exception: # noqa: BLE001
except Exception: # noqa: BLE001 # nosec B110
pass
if eos_token_id is not None:
special_ids.add(int(eos_token_id))
@@ -1838,42 +1692,24 @@ class PI052Policy(PreTrainedPolicy):
generated: list[int] = []
new_emb = None
# KV-cache decode: encode the (image-heavy) prefix once, then feed only
# the newly sampled token each step, attending to the cached keys. This
# turns an O(n_tokens * prefix_len) recompute into O(prefix_len + n_tokens)
# and is the dominant cost here (the prefix carries ~3*256 image tokens).
# With ``use_kv_cache=False`` the loop reduces exactly to the original
# recompute path (cache stays ``None`` so every step re-runs the full
# prefix), which we keep as a fallback / parity reference.
# Cache the image-heavy prefix; disabling the cache retains the full-recompute parity path.
cache = None
backbone = self.model.paligemma_with_expert
lm_head = backbone.paligemma.lm_head
# ``_prepare_attention_masks_4d`` always returns fp32 (0.0 / -inf
# literals). When weights are bf16, HF's PaliGemma SDPA raises
# "invalid dtype for bias - should match query's dtype". Pull the
# dtype from an attention *projection* weight specifically:
# ``to_bfloat16_for_selected_params`` keeps norms / embeddings in
# 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.
# Use q_proj's dtype because norms and embeddings may remain fp32 while SDPA queries are bf16.
backbone_dtype = backbone.paligemma.model.language_model.layers[0].self_attn.q_proj.weight.dtype
for _ in range(max_new_tokens):
if cache is None:
# First step (and every step when caching is disabled): run the
# full bidirectional-prefix forward. ``current_*`` already grow
# in the no-cache fallback below.
# Run the full bidirectional prefix initially or whenever caching is disabled.
step_embs = current_embs
att_2d = make_att_2d_masks(current_pad, current_att)
position_ids = torch.cumsum(current_pad, dim=1) - 1
att_2d_4d = self.model._prepare_attention_masks_4d(att_2d, dtype=backbone_dtype)
else:
# Incremental step: only the last token. It attends to every
# valid cached key (``current_pad`` already includes this token),
# so pad positions in the prefix stay masked just like the
# recompute path.
# Incremental decoding feeds only the last token while retaining prefix padding masks.
step_embs = new_emb
att_2d = current_pad[:, None, :]
att_2d_4d = self.model._prepare_attention_masks_4d(att_2d, dtype=backbone_dtype)
@@ -1894,10 +1730,7 @@ class PI052Policy(PreTrainedPolicy):
if special_ids and len(generated) < min_new_tokens:
for sid in special_ids:
logits_step[..., sid] = float("-inf")
# Mask FAST action tokens that fall *below* the ``<loc>`` block.
# They are never valid text, but the action-trained head leaks
# them as gibberish; unlike the loc/seg block this region is never
# legitimately emitted (even by VQA), so suppress it on every call.
# Suppress FAST-only vocabulary that otherwise leaks into generated text.
vocab_size = logits_step.shape[-1]
fast_skip = int(getattr(self.config, "fast_skip_tokens", 128))
fast_lo = vocab_size - 1 - fast_skip - (_FAST_ACTION_VOCAB_SIZE - 1)
@@ -1913,9 +1746,7 @@ class PI052Policy(PreTrainedPolicy):
# embed_language_tokens already applies the Gemma sqrt(hidden) scale (tf>=5.4.0).
new_emb = backbone.embed_language_tokens(next_ids.unsqueeze(0))
# ``current_pad`` tracks valid keys for both paths (cache mask +
# position ids). Only the recompute path needs the full embedding /
# block-attention history re-fed each step.
# Both paths track valid keys, but only recompute retains full embedding history.
current_pad = torch.cat([current_pad, ones_step], dim=1)
if not use_kv_cache:
current_embs = torch.cat([current_embs, new_emb], dim=1)
@@ -1964,26 +1795,16 @@ class PI052Policy(PreTrainedPolicy):
n = self._batch_size_from_observation(batch)
self._ensure_subtask_state(n)
tasks = self._tasks_from_batch(batch, n)
# Normalized state for the low-level action prompt (mirrors training:
# "User: {subtask}, State: {256-bin};"). batch state is already
# normalized by the eval preprocessor's NormalizerProcessorStep.
# Mirror training by appending the already normalized state to low-level prompts.
state_all = batch.get(OBS_STATE)
# Decide whether to (re)generate subtasks this chunk or hold the last
# ones. Training conditions the action expert on the subtask active over
# an interval (seconds), not a fresh subtask every 0.25s; regenerating
# every chunk also makes the subtask thrash. With subtask_replan_steps>0
# we regenerate only every ~that many env steps and reuse the held
# subtask in between (state is still refreshed each chunk).
# Hold subtasks for the configured interval to match training and avoid rapid replanning.
replan = int(getattr(self.config, "subtask_replan_steps", 0) or 0)
hold_chunks = max(1, round(replan / self.config.n_action_steps)) if replan > 0 else 1
regenerate = self._subtask_chunk_counter % hold_chunks == 0 or not any(self.last_subtasks or [])
self._subtask_chunk_counter += 1
# Generate one subtask per parallel env, each conditioned on that env's
# 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).
# Generate and batch one independently conditioned subtask per environment.
rows: list[tuple[Tensor, Tensor | None]] = []
tokenizer = None
for i in range(n):
@@ -2033,9 +1854,7 @@ class PI052Policy(PreTrainedPolicy):
)
self.last_subtasks_raw[i] = msg or ""
# Faithful hierarchical inference: condition the action expert on the
# model's own generated subtask verbatim (this is exactly what the
# ``low_level_execution`` recipe did at training — ``user: ${subtask}``).
# Feed the generated subtask verbatim, matching low-level training.
if msg and not _looks_like_gibberish(msg):
subtask = " ".join(msg.strip().split())
self._last_good_subtasks[i] = subtask
@@ -2044,11 +1863,7 @@ class PI052Policy(PreTrainedPolicy):
logger.info("PI052 eval subtask[%d]: %r (task=%r)", i, subtask, task)
return subtask
# Generation unusable (empty / gibberish). Training never fed such a
# prompt to the action expert, so the least-OOD choice is to reuse this
# env's last accepted subtask; on the first chunk (none yet) derive one
# from the task so the action expert still gets an imperative command
# rather than the raw high-level instruction.
# Reuse the last valid subtask, or derive an initial imperative, when generation fails.
debug = getattr(self, "_last_select_message_debug", "") or ""
if not task:
reason = "No task string was available in the batch."
@@ -2193,11 +2008,7 @@ class PI052Policy(PreTrainedPolicy):
return sorted_ix.gather(-1, choice).squeeze(-1)
return torch.multinomial(probs, num_samples=1).squeeze(-1)
# ------------------------------------------------------------------
# Inlined from PI05Policy (vendored; pi052 does not inherit pi05).
# Kept verbatim except PI05Policy.forward -> _pi05_flow_forward (the
# flow-only fallback used by PI052Policy.forward on unannotated batches).
# ------------------------------------------------------------------
# PI0.5 flow-only fallback for unannotated batches.
@classmethod
def from_pretrained(
cls: type[T],
@@ -2406,9 +2217,7 @@ class PI052Policy(PreTrainedPolicy):
if head_scale == 1.0 and backbone_scale == 1.0 and expert_scale == 1.0:
return self.parameters()
# Both ``lm_head.weight`` and the tied ``embed_tokens.weight`` go in the
# head group — boosting only the projection without the embedding pulls
# them apart and breaks the tie PaliGemma was pre-trained with.
# Keep the tied LM projection and embeddings in the same optimizer group.
head_substrings = (
"paligemma_with_expert.paligemma.lm_head.",
"paligemma_with_expert.paligemma.model.language_model.embed_tokens.",
@@ -2550,9 +2359,7 @@ class PI052Policy(PreTrainedPolicy):
"""Predict a chunk of actions given environment observations."""
self.eval()
# Opt-in FlashRT FP8: calibrate static scales on the first real observation
# and swap the MLPs in place. Guard set before the call so the calibration
# forward (which re-enters predict_action_chunk) does not recurse.
# Guard before first-observation FP8 calibration to prevent recursive prediction.
if self.config.use_flashrt_fp8_mlp and not getattr(self, "_fp8_applied", False):
self._fp8_applied = True
self.apply_flashrt_fp8_mlp(batch)
src/lerobot/policies/pi052/processor_pi052.py
+9 -35
View File
@@ -12,25 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""π0.5 v2 pre/post-processor factory.
"""PI052 processor factory with optional recipe rendering and text tokenization.
When ``config.recipe_path`` is set, the pre-processor pipeline becomes:
rename observations
add batch dim
relative-action prep (inherited from π0.5)
NormalizerProcessorStep
RenderMessagesStep — recipe → messages, target_message_indices,
message_streams (PR 1 of the steerable
stack)
PI052TextTokenizerStep — messages → input_ids + label mask +
predict_actions
DeviceProcessorStep
When ``recipe_path`` is ``None`` we delegate to the plain π0.5 pipeline
so unannotated datasets keep working.
Post-processor is unchanged from π0.5.
Without a recipe it delegates to the standard PI0.5 pipeline.
"""
from __future__ import annotations
@@ -55,9 +39,8 @@ from lerobot.processor import (
policy_action_to_transition,
transition_to_policy_action,
)
# RenderMessagesStep is intentionally not re-exported from
# ``lerobot.processor`` because it pulls in optional language-stack deps;
# import it directly.
# Import directly to keep optional language dependencies out of ``lerobot.processor``.
from lerobot.processor.render_messages_processor import RenderMessagesStep
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
@@ -108,21 +91,11 @@ def make_pi052_pre_post_processors(
),
]
# FAST tokenizer for discrete-action CE supervision (paper §III.C).
# Only inserted when explicitly enabled — keeps the post-training-
# style recipe (flow + text) as the default. When on, the step
# writes ACTION_TOKENS / ACTION_TOKEN_MASK into
# ``COMPLEMENTARY_DATA`` and the modeling forward picks them up.
# Add FAST action-token supervision only when explicitly enabled.
if getattr(config, "enable_fast_action_loss", False):
# Per Pertsch et al. 2025 (FAST [64], π0.5 §III.C): fit the
# tokenizer on this dataset's action distribution rather than
# using the universal codebook off the shelf. We do this once
# and cache to disk, keyed on (dataset, base, n_samples).
# Fit once on this dataset and cache by dataset, base tokenizer, and sample count.
action_tokenizer_path = config.action_tokenizer_name
if (
getattr(config, "auto_fit_fast_tokenizer", False)
and dataset_repo_id is not None
):
if getattr(config, "auto_fit_fast_tokenizer", False) and dataset_repo_id is not None:
from .fit_fast_tokenizer import fit_fast_tokenizer # noqa: PLC0415
cache_dir = Path(config.fast_tokenizer_cache_dir).expanduser()
@@ -141,7 +114,8 @@ def make_pi052_pre_post_processors(
"FAST tokenizer fit failed (%s) — falling back to "
"the universal base tokenizer %r. Train will still "
"work but compression will be suboptimal.",
exc, config.action_tokenizer_name,
exc,
config.action_tokenizer_name,
)
input_steps.append(
src/lerobot/policies/pi052/text_processor_pi052.py
+14 -68
View File
@@ -12,25 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""π0.5 v2 text-tokenisation step.
"""Tokenize PI052's plain-text rendered messages and build text/action supervision masks.
PaliGemma is *not* chat-pretrained, so we can't lean on
``tokenizer.apply_chat_template``. Instead we concatenate the rendered
messages as plain text with simple ``User: ... Assistant: ...`` role
delimiters — matching the prompt format π0.5 uses in the paper
(``Task: ... State: ... Action: ...``).
Outputs:
* ``OBS_LANGUAGE_TOKENS`` / ``OBS_LANGUAGE_ATTENTION_MASK`` — the
concatenated prompt tokenised by the PaliGemma tokenizer (the same
one ``processor_pi05`` already uses).
* ``text_labels`` — same shape as token ids, ``-100`` everywhere except
positions belonging to messages whose index is in
``target_message_indices``. ``modeling_pi052`` runs cross-entropy on
those positions via the PaliGemma ``lm_head``.
* ``predict_actions`` — bool tensor, ``True`` iff any of the rendered
target messages has ``message_streams[i] == "low_level"``.
PaliGemma is not chat-trained, so messages use explicit role delimiters instead of a chat template.
"""
from __future__ import annotations
@@ -182,11 +166,7 @@ def _sample_indices(value: Any, batch_size: int) -> list[int | None]:
return [int(value)] * batch_size
# VQA spatial answers → PaliGemma <loc> format (PI052 only).
# Dataset JSON uses Qwen2.5-VL's 0–1000 *normalized* grounding coords (not
# pixels — verified empirically); PaliGemma's <locNNNN> vocab is [0, 1023], so
# ``loc_idx = round(coord / 1000 * 1023)`` is resolution-independent. Converted
# here, not in the dataset, so the raw JSON stays backbone-agnostic.
# Convert normalized Qwen2.5-VL coordinates to PaliGemma's resolution-independent <loc> range.
_VQA_COORD_SCALE = 1000.0
@@ -333,14 +313,9 @@ def _format_messages(
for i, m in enumerate(messages):
role = m.get("role", "user")
content = m.get("content", "") or ""
# Role tag + newline. The model has to learn to emit the same
# role tokens at generation time, which is fine for greedy
# decoding because the chat template is implicit in the
# supervised target span.
# Supervise the explicit role format used again during generation.
header = f"{role.capitalize()}: "
# A supervised target turn ends with EOS so the model learns to
# terminate; the span below covers content + EOS. Non-target
# turns (and inference) carry no EOS.
# Include EOS only in supervised target spans so generation learns to stop.
body = content + eos_token if (eos_token and i in targets) else content
# span covers the content (+ EOS) portion only — never the role
# tag — so labels are computed over the supervised payload.
@@ -383,29 +358,19 @@ class PI052TextTokenizerStep(ProcessorStep):
self._tokenizer = register_paligemma_loc_tokens(AutoTokenizer.from_pretrained(self.tokenizer_name))
return self._tokenizer
# ------------------------------------------------------------------
# Pipeline step
# ------------------------------------------------------------------
def __call__(self, transition: EnvTransition) -> EnvTransition | None:
transition = transition.copy()
complementary = transition.get(TransitionKey.COMPLEMENTARY_DATA, {}) or {}
messages = complementary.get("messages") or []
if not messages:
# No recipe was rendered — caller will fall back to the
# plain Pi0.5 prompt path. We pass the transition through
# unmodified.
# Preserve the transition for the plain PI0.5 prompt fallback.
return transition
tokenizer = self._ensure_tokenizer()
# Normalized proprioceptive state (set by NormalizerProcessorStep, which
# runs before this step). Injected into low-level action prompts so the
# action expert sees proprioception, matching pi05's discretized State:.
# Add normalized proprioception to low-level prompts, matching PI0.5.
state_all = (transition.get(TransitionKey.OBSERVATION) or {}).get(OBS_STATE)
# VQA coords are 0–1000 normalized (Qwen2.5-VL convention) — the
# <loc> conversion is camera-resolution-independent and needs no
# observation lookup here.
# Normalized VQA coordinates need no camera lookup.
if _is_batched_messages(messages):
indices_iter = _sample_indices(complementary.get("index"), len(messages))
encoded = [
@@ -464,9 +429,7 @@ class PI052TextTokenizerStep(ProcessorStep):
sample_idx: int | None = None,
state_row: Any = None,
) -> tuple[Tensor, Tensor, Tensor, Tensor, str]:
# Optional: drop non-target messages per the dropout config.
# Keeps the supervised-target indices stable by re-mapping
# after removal.
# Remap target indices after optional context dropout.
if (
self.plan_dropout_prob
or self.memory_dropout_prob
@@ -480,19 +443,12 @@ class PI052TextTokenizerStep(ProcessorStep):
sample_idx=sample_idx,
)
# Rewrite bbox / keypoint VQA target answers from JSON to
# PaliGemma <loc> text. Coords are 0–1000 normalized so this is
# camera-independent.
# Rewrite normalized VQA answers as PaliGemma <loc> text.
messages = _messages_vqa_to_loc(messages, target_indices)
# Flatten ``say`` tool calls into ``<say>...</say>`` text before
# stripping, so the spoken reply is actually tokenized and
# supervised (PaliGemma's flat prompt has no structured calls).
# Flatten ``say`` calls because PaliGemma receives plain text.
messages = [_strip_blocks(_flatten_say_tool_calls(m)) for m in messages]
# Low-level (action-conditioning) samples get the discretized state
# appended to their user message, mirroring pi05's
# "..., State: {256-bin};" so the action expert sees proprioception.
# Higher-level text streams (subtask/memory generation) stay state-free.
# Add state only to low-level action prompts; keep higher-level streams state-free.
if state_row is not None and any(s == "low_level" for s in message_streams):
state_str = discretize_state_str(state_row)
for m in reversed(messages):
@@ -533,20 +489,13 @@ class PI052TextTokenizerStep(ProcessorStep):
continue
labels[token_pos] = input_ids[token_pos]
# Scan ALL message streams (not just targets): the
# ``low_level_execution`` recipe drops ``target: true`` on
# the assistant to avoid trivial copy-from-user text-CE; the
# flow loss still needs to fire, gated by ``stream: low_level``.
# Scan all streams because low-level flow may intentionally have no text target.
predict_actions = torch.tensor(
bool(any(s == "low_level" for s in message_streams)),
dtype=torch.bool,
)
return input_ids, attention_mask, labels, predict_actions, prompt
# ------------------------------------------------------------------
# Per-component prompt dropout (Pi0.7 §V.E)
# ------------------------------------------------------------------
def _apply_prompt_dropout(
self,
messages: list[dict[str, Any]],
@@ -563,10 +512,7 @@ class PI052TextTokenizerStep(ProcessorStep):
seed = self.dropout_seed
if seed is None:
# Canonical row-index key set by ``BatchProcessor`` /
# ``render_messages_processor``. Falling back to other
# keys silently gave every sample seed=0 → identical
# dropout pattern across the whole epoch.
# Use the canonical row index to avoid identical dropout across an epoch.
seed_src = sample_idx if sample_idx is not None else complementary.get("index", 0)
try:
if hasattr(seed_src, "item"):
src/lerobot/runtime/adapter.py
+6 -36
View File
@@ -12,19 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""Policy adapter base class for the language-conditioned runtime.
"""Policy adapters for the language runtime.
The runtime loop drives the *control algorithm* (throttling, output
rejection, the subtask -> memory cascade, diagnostics) and delegates the
*policy primitives* (act, generate text) to an adapter. :class:`BaseLanguageAdapter`
implements the algorithm once; a policy subclasses it and supplies:
* :meth:`select_action` — observation + language context -> action chunk
* :meth:`generate_text` — a text stream (``kind``) -> decoded string
* :meth:`build_messages` — the prompt for each ``kind``
A policy that needs full control can instead satisfy the
:class:`LanguageConditionedPolicyAdapter` protocol directly.
The base adapter owns generation control and diagnostics while subclasses provide policy-specific actions and text.
"""
from __future__ import annotations
@@ -41,11 +31,7 @@ _SAY_RE = re.compile(r"<\s*say\s*>(.*?)<\s*/\s*say\s*>", re.IGNORECASE | re.DOTA
@dataclass
class GenerationConfig:
"""Text-generation knobs, fixed for the lifetime of an adapter.
These are configuration (set once from the CLI), not per-tick runtime
state — they live on the adapter, never in :class:`RuntimeState`.
"""
"""Text-generation settings fixed for the adapter's lifetime."""
min_new_tokens: int = 0
temperature: float = 0.0
@@ -57,11 +43,7 @@ class GenerationConfig:
@dataclass
class LanguageDiagnostics:
"""Rejection / repeat counters surfaced in the runtime panel.
Keyed by text ``kind`` (``subtask`` / ``memory`` / ...) so the same
accounting works for any cascade shape.
"""
"""Runtime-panel rejection and repeat counters keyed by text kind."""
last_raw: dict[str, str] = field(default_factory=dict)
empty: dict[str, int] = field(default_factory=dict)
@@ -82,8 +64,6 @@ class BaseLanguageAdapter(ABC):
self.diag = LanguageDiagnostics()
self._chunks_until_regen = 0
# --- policy primitives (subclass supplies) ---------------------------
@abstractmethod
def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any:
"""Produce an action chunk from the observation + current language context."""
@@ -98,8 +78,6 @@ class BaseLanguageAdapter(ABC):
) -> str:
"""Generate one text stream (``kind``) and return the decoded string."""
# --- generic control algorithm (runtime calls these) ----------------
def update_language_state(self, observation: dict[str, Any] | None, state: RuntimeState) -> None:
"""Throttled regeneration of the language context (subtask / memory / ...)."""
if self._chunks_until_regen > 0:
@@ -122,16 +100,13 @@ class BaseLanguageAdapter(ABC):
plan, _speech = split_plan_and_say(text)
return "" if looks_like_gibberish(plan) else plan
# --- overridable cascade + shared helpers ---------------------------
def _regenerate_context(self, observation: dict[str, Any] | None, state: RuntimeState) -> None:
"""Default hierarchy: regenerate the subtask, then memory when it changes.
Override for a policy with a different language hierarchy.
"""
if not self.gen.enable_subtask:
# Direct-subtask mode: the operator supplies the subtask; don't
# generate (and thus don't overwrite) it.
# Preserve operator-provided subtasks in direct mode.
return
subtask = self._generate_filtered("subtask", observation, state)
if subtask is None:
@@ -169,12 +144,7 @@ class BaseLanguageAdapter(ABC):
class DirectTaskPolicyAdapter(BaseLanguageAdapter):
"""Adapter for flat policies conditioned directly on the operator's task text.
Policies such as PI0.5 and MolmoAct2 do not expose a language-generation
head. Their preprocessors pack the current task into the model inputs, so
the runtime only needs to request an action chunk.
"""
"""Adapter for flat policies whose preprocessors condition actions on the operator's task."""
def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any:
return self.policy.predict_action_chunk(observation)
src/lerobot/runtime/cli.py
+24 -141
View File
@@ -12,34 +12,9 @@
# 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.
"""Interactive REPL for a language-conditioned robot policy.
"""Interactive CLI for language-conditioned policy rollouts.
Policy-agnostic CLI over :class:`lerobot.runtime.LanguageConditionedRuntime`.
A policy wires it up with :func:`run`, passing an adapter factory
(``policy -> LanguageConditionedPolicyAdapter``); see
``lerobot.scripts.lerobot_language_runtime`` for the entry point.
Stdin is the user channel: type a task, then natural-language
interjections. The runtime prints state changes (plan / subtask /
memory) as they happen.
Examples
--------
No-robot REPL on a Hub checkpoint — useful for sanity-checking text generation::
uv run lerobot-rollout --language \\
--policy.path=<repo-or-dir> \\
--no_robot \\
--task="please clean the kitchen"
With a real robot::
uv run lerobot-rollout --language \\
--policy.path=... \\
--robot.type=so101 --robot.port=/dev/tty.usbmodem...
``--policy.path`` accepts either a local directory or a Hugging Face Hub repo id.
It supports a text-only REPL, real robots, and RoboCasa with local or Hub checkpoints.
"""
from __future__ import annotations
@@ -104,14 +79,7 @@ def _parse_args(argv: list[str] | None = None, *, prog: str | None = None) -> ar
action="store_true",
help="Skip robot connection and open a language-only REPL.",
)
# --- Real-robot mode args ----------------------------------------
# Setting ``--robot.type`` flips the runtime into autonomous mode:
# it connects to the robot, builds an observation provider that
# reads ``robot.get_observation()``, and
# an action executor that postprocesses (denormalises) the policy's
# output and calls ``robot.send_action(...)`` at ``--ctrl_hz``. The
# high-level REPL-style stdin still works in a background thread
# for interjections.
# ``--robot.type`` enables real-time control while stdin remains interactive.
p.add_argument(
"--robot.type",
dest="robot_type",
@@ -152,12 +120,7 @@ def _parse_args(argv: list[str] | None = None, *, prog: str | None = None) -> ar
help="Direct-subtask mode (sim OR robot): your typed text IS the subtask "
"fed to the action expert; the LM subtask generator is disabled.",
)
# --- RoboCasa simulation mode args -------------------------------
# Setting ``--sim`` flips the runtime into simulation mode: instead of
# a real robot it drives a single RoboCasa mujoco scene, feeding the
# eval observation/action pipeline. The operator still types prompts
# (/action <prompt>) that the policy executes inside the chosen scene.
# Mutually exclusive with ``--robot.type``.
# ``--sim`` uses the eval pipeline and is mutually exclusive with a robot.
p.add_argument(
"--sim",
action="store_true",
@@ -359,10 +322,7 @@ def _strip_runtime_owned_language_cols(sample: dict) -> None:
sample.pop(k, None)
# Model-input keys some policies emit OUTSIDE the ``observation.*`` namespace and
# still need at inference. MolmoAct2's processor packs its prompt + images into
# these top-level keys; PI0-family policies never produce them, so keeping the
# allowlist is a no-op for them.
# Non-observation model inputs emitted by processors such as MolmoAct2's.
_MODEL_INPUT_PASSTHROUGH_KEYS = (
"input_ids",
"attention_mask",
@@ -395,12 +355,7 @@ def _load_policy_and_preprocessor(
fp8: bool = False,
device: str | None = None,
) -> tuple[Any, Any, Any]:
"""Load a policy checkpoint (local path or Hub repo id).
When ``load_processors_from_checkpoint`` is set, the pre/post processors
are loaded exactly like ``lerobot-eval``. RoboCasa uses this path so its
normalization and recipe match the checkpoint.
"""
"""Load a local or Hub policy, optionally with its eval processors."""
from lerobot.configs import PreTrainedConfig # noqa: PLC0415
from lerobot.policies.factory import get_policy_class, make_pre_post_processors # noqa: PLC0415
@@ -411,10 +366,7 @@ def _load_policy_and_preprocessor(
if device:
cfg.device = device
# Inference-only overrides (mirror lerobot-eval). torch.compile recompiles
# whenever the prompt length changes (every subtask switch) — catastrophic
# in the interactive runtime — and gradient checkpointing only slows the
# forward pass. Neither is wanted for serving.
# Variable prompts trigger recompilation, and checkpointing only adds inference overhead.
if getattr(cfg, "compile_model", False):
cfg.compile_model = False
if getattr(cfg, "gradient_checkpointing", False):
@@ -453,9 +405,7 @@ def _build_language_rollout_context(args: argparse.Namespace) -> Any:
from lerobot.configs import parser # noqa: PLC0415
from lerobot.rollout import RolloutConfig, build_rollout_context # noqa: PLC0415
# Importing the rollout entry point registers every bundled camera and
# robot config choice used by Draccus. Third-party choices were registered
# by the top-level entry point before reaching this function.
# Import for bundled Draccus camera and robot registrations.
from lerobot.scripts import lerobot_rollout as _rollout_registrations # noqa: F401, PLC0415
rollout_argv = [arg for arg in args.raw_argv if arg.startswith(("--policy.", "--robot."))]
@@ -712,9 +662,7 @@ def _make_state_panel_renderer(
dispatched = int(st.get("actions_dispatched") or 0)
console.print(f" [dim]queued actions: {queue_len} dispatched: {dispatched}[/]")
# Overfit / memorisation diagnostics from the adapter. High repeat
# + fully cycling queue ⇒ stuck on one subtask (memorised a phase);
# climbing gibberish ⇒ LM head collapsed to chat-template salads.
# Surface repeated or rejected generations as overfitting diagnostics.
diag = getattr(runtime.policy_adapter, "diag", None)
if diag is not None:
raw_subtask = diag.last_raw.get("subtask")
@@ -733,9 +681,7 @@ def _make_state_panel_renderer(
if mem_gib:
console.print(f" [dim]gen rejects memory:{mem_gib}[/]")
console.rule(style="cyan")
# Runtime scrollback — log lines pushed from generation steps
# (warnings, gibberish rejections, plan speech). Last N lines,
# oldest first.
# Show recent generation warnings and speech oldest-first.
if scrollback:
for line in scrollback:
console.print(f" [magenta]{line.rstrip()}[/]")
@@ -753,15 +699,7 @@ def _make_state_panel_renderer(
def _silence_noisy_loggers() -> None:
"""Drop chatty third-party loggers down to WARNING.
HuggingFace / httpx / urllib3 emit one log line per HTTP request,
which the REPL has to print between the state block and the
prompt — completely unreadable. We never need that detail in the
REPL and the user can opt back into it via ``-v`` (verbose mode
keeps DEBUG on the lerobot loggers but still gates the noisy ones
here unless they explicitly want them).
"""
"""Keep request-level third-party logs out of the interactive prompt."""
for name in (
"httpcore",
"httpcore.connection",
@@ -781,14 +719,7 @@ def _silence_noisy_loggers() -> None:
):
logging.getLogger(name).setLevel(logging.WARNING)
# The robot's relative-goal-position clamp warning fires *every*
# dispatch tick on a memorised model — the LM is trying to jump
# the wrist far past where max_relative_target allows, so the
# warning floods the panel at ~30 Hz. Promote it from WARNING to
# DEBUG: the dispatch counter on the panel already tells the
# operator the loop is running, and the panel itself shows
# whether motion is happening. If anyone needs the per-action
# clamp detail, ``-v`` puts it back via DEBUG.
# Clamp warnings can fire every control tick and flood the panel.
logging.getLogger("lerobot.robots.utils").setLevel(logging.ERROR)
@@ -824,9 +755,7 @@ def run(
file=sys.stderr,
)
return 2
# Create the sim env subprocess BEFORE the policy initialises CUDA — the
# env worker inherits a corrupt EGL/GL context if forked from a CUDA parent
# (dark/garbled renders). This mirrors eval's make_env-before-make_policy.
# Fork the simulator before CUDA initialization to avoid inherited EGL corruption.
sim_env = None
sim_obs = None
sim_stream_server = None
@@ -877,10 +806,7 @@ def run(
if panel_label is None:
panel_label = str(policy_type or "runtime").upper()
# No startup prompts — the runtime is command-driven. It comes up at
# the command line in ``paused`` mode (robot idle) unless ``--mode``
# forces a mode. The operator drives it with /action, /pause and
# /question.
# Default to idle until the operator supplies a command.
startup_mode = args.mode or "paused"
observation_provider: Callable[[], dict | None] | None = None
@@ -936,10 +862,7 @@ def run(
rerun_log=bool(args.rerun),
get_task=_live_task,
)
# Text-generation knobs are fixed config, passed to the adapter at
# construction — not smuggled through per-tick runtime state. Lets the
# operator try e.g. ``--text_temperature=0.6 --subtask_chunks_per_gen=5``
# on an under-trained checkpoint without recompiling.
# Generation settings belong to the adapter rather than mutable runtime state.
gen_config = GenerationConfig(
min_new_tokens=int(args.text_min_new_tokens or 0),
temperature=float(args.text_temperature or 0.0),
@@ -952,10 +875,6 @@ def run(
policy_adapter=adapter_factory(policy, gen_config),
observation_provider=observation_provider,
action_executor=robot_executor,
# No background event collector — the REPL drives ticks
# synchronously after each user input (REPL mode). Autonomous
# mode runs ``runtime.run()`` in a thread; stdin events are
# injected from the foreground.
event_collector=None,
chunk_hz=args.chunk_hz,
ctrl_hz=args.ctrl_hz,
@@ -971,8 +890,7 @@ def run(
# Let the sim backend read live task/subtask/memory for the video overlay.
if sim_backend is not None:
sim_backend.bind_runtime(runtime)
# Sim runs its control/render loop in the MAIN thread (see
# _run_sim_interactive) — background-thread rendering corrupts EGL.
# Keep EGL rendering on the main thread.
return _run_sim_interactive(
runtime,
sim_backend,
@@ -1008,14 +926,7 @@ def _run_sim_interactive(
panel_label: str = "Runtime",
direct_subtask: bool = False,
) -> int:
"""Main-thread control loop for the RoboCasa sim backend.
The tick loop — and therefore MuJoCo's EGL rendering — runs in the MAIN
thread. Driving the sim render from a background thread intermittently
corrupts the offscreen GL context (dark/garbled frames); main-thread
stepping matches ``lerobot-eval`` and renders cleanly. Stdin is polled
non-blockingly so typed commands still work while the sim runs.
"""
"""Keep RoboCasa rendering on the main thread while polling stdin."""
import select # noqa: PLC0415
import time # noqa: PLC0415
@@ -1028,10 +939,7 @@ def _run_sim_interactive(
runtime.state["current_subtask"] = initial_task if direct_subtask else None
runtime.state["mode"] = "action"
# Clean chat-style prompt. The control loop steps in the MAIN thread (clean
# EGL rendering); the browser live-view shows the rollout, so the terminal
# stays a quiet command line. Nothing is printed mid-step, so typing is never
# clobbered — you can queue the next command any time.
# Keep the terminal quiet while the browser renders the rollout.
_mode_line = (
" Mode: DIRECT subtask (your text drives the action expert as-is)\n"
if direct_subtask
@@ -1082,10 +990,7 @@ def _run_sim_interactive(
runtime.policy.reset()
print("[reset] new kitchen scene", flush=True)
else:
# A bare line is a new command: switch the robot to it
# immediately (clear the in-flight chunk + subtask) and
# force the subtask to regenerate on the very next tick
# (reset the adapter throttle + high-level rate gate).
# Clear queued actions and rearm generation for a new command.
runtime.set_task(cmd)
# Direct mode: the typed text is the subtask itself;
# otherwise clear it so the model regenerates one.
@@ -1101,8 +1006,7 @@ def _run_sim_interactive(
print(f"[running] {cmd}", flush=True)
_prompt()
# One tick in the MAIN thread: subtask/action gen + env.step + render.
# inference_mode matches lerobot-eval's forward context.
# Match lerobot-eval's inference context on the main thread.
if runtime.state.get("mode", "paused") == "action":
with torch.inference_mode():
runtime.step_once()
@@ -1133,23 +1037,14 @@ def _run_robot_interactive(
direct_subtask: bool = False,
panel_label: str = "Runtime",
) -> int:
"""Real-robot interactive loop.
The control loop runs at real-time rates in a background thread
(``runtime.run()`` — a robot must be driven at a steady ``ctrl_hz``), while
the foreground is a clean chat prompt: type a command to run it (generate- or
direct-subtask mode), ``/pause`` / ``/resume`` / ``stop``. Starts PAUSED so
the arm doesn't move until you issue a command.
"""
"""Run steady robot control in the background and commands in the foreground."""
import threading # noqa: PLC0415
import time # noqa: PLC0415
if initial_task:
runtime.set_task(initial_task)
runtime.state["current_subtask"] = initial_task if direct_subtask else None
# A task was given (via --task or the startup picker) => start running it
# immediately. Without an initial task we stay paused until the first
# typed command (which switches to action). No flag needed.
# An explicit initial task starts immediately; otherwise the robot stays paused.
runtime.state["mode"] = "action"
mode_line = (
@@ -1226,15 +1121,7 @@ def _run_robot_interactive(
def _run_repl(
runtime: Any, *, initial_task: str | None, max_ticks: int | None, panel_label: str = "Runtime"
) -> int:
"""Claude-Code-style block REPL.
Each turn redraws a status block (task / subtask / plan / memory)
at the top, prints any robot log lines that came in since the last
turn, then asks for input on a clean ``> `` prompt at the bottom.
No live region, no panel re-renders, no rendering races with HTTP
log lines — just clear-screen + reprint each turn, the way a
chat-style REPL is meant to look.
"""
"""Redraw the status block and logs once per REPL turn."""
try:
from rich.console import Console # noqa: PLC0415
except ImportError:
@@ -1245,8 +1132,6 @@ def _run_repl(
return 2
_redraw = _make_state_panel_renderer(runtime, mode_label="no robot", panel_label=panel_label)
# Keep a local ``console`` just for the styled input prompt; the
# state panel is owned by the shared renderer.
console = Console(highlight=False)
last_logs: list[str] = []
@@ -1268,9 +1153,7 @@ def _run_repl(
if lower in {"stop", "quit", "exit"}:
break
# Command-driven: /action "task", /pause, /question "...",
# /help. ``_handle_slash_command`` runs the VQA query inline
# for /question (single-threaded REPL — no concurrency).
# Slash commands, including VQA questions, run inline.
if _handle_slash_command(runtime, line):
last_logs = list(runtime.state.get("log_lines") or [])
_redraw(last_logs)
src/lerobot/runtime/language_runtime.py
+1 -6
View File
@@ -117,12 +117,7 @@ class RuntimeState:
class LanguageConditionedPolicyAdapter(Protocol):
"""The contract the runtime loop depends on.
:class:`lerobot.runtime.adapter.BaseLanguageAdapter` provides a
batteries-included implementation; a policy can satisfy this protocol
directly for full control.
"""
"""Runtime policy contract, implemented directly or through ``BaseLanguageAdapter``."""
def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any: ...
src/lerobot/runtime/registry.py
+1 -6
View File
@@ -12,12 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""Registry mapping a policy type to its language-runtime adapter.
Kept as import strings (resolved lazily) so ``lerobot-rollout --language``
never imports a policy package until it actually loads that policy — the
same pattern as :mod:`lerobot.policies.factory`.
"""
"""Lazy mapping from policy types to language-runtime adapters."""
from __future__ import annotations
src/lerobot/runtime/rerun_viz.py
+2 -10
View File
@@ -12,13 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""Rerun live visualisation for the interactive runtime (real-robot camera view).
Starts a headless rerun gRPC server + web viewer so a remote operator can watch
the robot's cameras (and state / subtask) over SSH by forwarding two ports and
opening the web viewer in a browser. Logging is best-effort — a rerun failure
never interrupts robot control.
"""
"""Best-effort Rerun camera visualization for local or SSH-forwarded robot rollouts."""
from __future__ import annotations
@@ -41,9 +35,7 @@ def start_rerun(app_name: str = "lerobot_runtime", grpc_port: int = 9876, web_po
url = rr.serve_grpc(grpc_port=grpc_port)
rr.serve_web_viewer(web_port=web_port, open_browser=False, connect_to=url)
_ENABLED = True
# Open the viewer with the data URL as a query param so it auto-connects
# to the gRPC stream (plain http://host:web_port shows only the welcome
# screen — the web app needs the ?url= to know where the data is).
# Include the stream URL so the web viewer connects automatically.
view_url = f"http://localhost:{web_port}/?url={url}"
print(
f"[runtime] rerun live view: {view_url}\n"
src/lerobot/runtime/sim_robocasa.py
+15 -58
View File
@@ -12,18 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""RoboCasa simulation backend for the interactive language runtime.
"""RoboCasa backend for interactive language-conditioned rollouts.
Lets an operator type open-ended prompts (``/action <prompt>``) and have a
language-conditioned policy (e.g. PI052) execute them inside a RoboCasa mujoco
kitchen scene. The observation/action pipeline mirrors ``lerobot-eval`` exactly
so behaviour matches offline evaluation; only the *source* of observations and
the *sink* of actions differ from the real-robot backend, which is left
untouched.
A RoboCasa episode always instantiates a concrete scene (objects + layout) from
its task name, so ``--sim.task`` selects the scene while the prompt typed at the
prompt drives what the policy is asked to do inside it.
It reuses the eval observation/action pipeline while prompts control a persistent selected scene.
"""
from __future__ import annotations
@@ -63,10 +54,7 @@ def _label_panel(img: np.ndarray, label: str) -> np.ndarray:
return img
# RoboCasa's MuJoCo EGL offscreen renderer produces garbled/static frames when
# only ONE worker env is running (reproducible with lerobot-eval --batch_size=1).
# With >=2 workers the renderer is stable. We therefore run the interactive sim
# with a small vec env, drive env 0 with the policy, and ignore the rest.
# Two workers avoid broken single-worker EGL rendering; only env 0 is displayed.
_SIM_N_ENVS = 2
@@ -78,19 +66,13 @@ def create_sim_env(
seed: int | None,
render_size: int = 384,
) -> tuple[Any, dict]:
"""Create + reset a RoboCasa AsyncVectorEnv (n_envs=_SIM_N_ENVS), return (env, obs).
"""Create and reset the vectorized RoboCasa environment before CUDA initializes.
MUST be called BEFORE the policy initialises CUDA in the parent process, so
the forkserver workers don't inherit a CUDA context (which corrupts EGL).
Uses >=2 workers because single-worker EGL rendering is broken on this stack
(garbled frames) — the same reason lerobot-eval renders cleanly only at
batch_size>=2. Only env 0 is driven/displayed.
Two workers keep EGL stable, while only env 0 is driven and displayed.
"""
from lerobot.envs.configs import RoboCasaEnv as RoboCasaEnvConfig # noqa: PLC0415
# Higher-res observation cameras => higher-quality display. The policy is
# unaffected: its preprocessor resizes images to 224 and VISUAL norm is
# identity, so only render cost (not behaviour) changes with render_size.
# The policy resizes inputs, so render_size only affects display quality and cost.
env_cfg = RoboCasaEnvConfig(
task=task,
split=split,
@@ -98,8 +80,7 @@ def create_sim_env(
observation_height=render_size,
observation_width=render_size,
)
# Persistent kitchen: never end/reset on task success, and use a huge horizon
# so the scene doesn't truncate. The user drives it with sequential prompts.
# Keep one kitchen alive across sequential prompts.
envs = env_cfg.create_envs(
n_envs=_SIM_N_ENVS,
use_async_envs=True,
@@ -114,14 +95,7 @@ def create_sim_env(
def start_mjpeg_server(port: int, get_frame: Callable[[], np.ndarray | None]) -> Any:
"""Start an MJPEG server serving frames from ``get_frame()`` on ``port``.
Started early (before the ~60s policy load) so the port listens immediately
and browsers get a page instead of connection-refused. ``get_frame`` returns
the latest annotated frame or None (a "waiting" placeholder is shown until
frames arrive). The server thread only reads/encodes frames — no CUDA/EGL —
so it never affects rendering. Returns the server (for shutdown) or None.
"""
"""Start an MJPEG server that shows a placeholder until ``get_frame`` returns frames."""
import io # noqa: PLC0415
import threading # noqa: PLC0415
import time # noqa: PLC0415
@@ -189,14 +163,9 @@ def start_mjpeg_server(port: int, get_frame: Callable[[], np.ndarray | None]) ->
class RoboCasaSimBackend:
"""Drive a single RoboCasa gym env from the language runtime.
"""Expose a RoboCasa environment through the runtime observation/action contract.
Exposes ``observation_provider`` / ``action_executor`` closures matching the
runtime's injected-callable contract, plus ``disconnect`` so the shared
The runtime cleanup path closes the env and flushes the video.
The env must be created via :func:`create_sim_env` *before* the policy
touches CUDA (see that function's note on the EGL/CUDA fork hazard).
The environment must be created before the policy initializes CUDA.
"""
def __init__(
@@ -216,7 +185,6 @@ class RoboCasaSimBackend:
self.env = env
self._last_obs = last_obs
self._scene_task = task
# Camera views to composite into the display frame (order = left→right).
self._view_cams = view_cams or [
"robot0_agentview_left",
"robot0_eye_in_hand",
@@ -234,8 +202,7 @@ class RoboCasaSimBackend:
self._latest_frame: np.ndarray | None = None
self._stream_server: Any = None
self._reset_count = 0
# State getters wired after the runtime exists (bind_runtime), so the
# video overlay can show the live task/subtask/memory.
# Bind these after runtime construction for live annotations.
self._task_getter: Callable[[], str | None] | None = None
self._subtask_getter: Callable[[], str | None] | None = None
self._memory_getter: Callable[[], str | None] | None = None
@@ -293,9 +260,7 @@ class RoboCasaSimBackend:
except Exception as exc: # noqa: BLE001
logger.warning("[sim] preprocess_observation failed: %s", exc)
return None
# ``task`` feeds the recipe RenderMessagesStep; the PI052 adapter
# overwrites the language tokens with its generated subtask before the
# action forward pass, so this only needs to be present, not exact.
# The adapter later replaces this recipe input with its generated subtask.
obs["task"] = [self._current_task()]
if self.preprocessor is not None:
try:
@@ -317,29 +282,21 @@ class RoboCasaSimBackend:
if action.ndim > 1 and action.shape[0] == 1:
action = action.squeeze(0)
action = action.detach().to("cpu").numpy()
# Only env 0 is policy-driven; tile its action across all workers so
# env.step gets a full (n_envs, action_dim) batch. The extra workers
# exist only to keep MuJoCo's EGL renderer stable (single-worker
# rendering is broken); their rollouts are ignored.
# Tile env 0's action because the extra workers exist only for EGL stability.
action_row = np.asarray(action, dtype=np.float32).reshape(-1)
action_np = np.tile(action_row, (self.env.num_envs, 1))
obs, _reward, terminated, truncated, _info = self.env.step(action_np)
self._last_obs = obs
if self.record:
self._capture_frame()
# AsyncVectorEnv auto-resets a sub-env after it terminates, so the
# scene continues on its own — no manual reset needed here.
# AsyncVectorEnv resets terminated sub-environments automatically.
if bool(np.any(terminated)) or bool(np.any(truncated)):
logger.info("[sim] episode ended — scene auto-reset")
except Exception as exc: # noqa: BLE001
logger.error("[sim] env.step failed: %s", exc, exc_info=True)
def _multiview_frame(self) -> np.ndarray | None:
"""Composite the configured camera views (env 0) side by side, labeled.
Uses the policy's own high-res observation images (env.step already
rendered them), so there's no extra render cost and orientation matches.
"""
"""Label and compose env 0's existing observation views without extra rendering."""
pixels = (self._last_obs or {}).get("pixels")
if not isinstance(pixels, dict) or not pixels:
return None
src/lerobot/scripts/lerobot_language_runtime.py
+1 -6
View File
@@ -15,12 +15,7 @@
"""Compatibility entry point for ``lerobot-language-runtime``.
Policy-agnostic: the runtime resolves the right adapter from the loaded
policy's type via :mod:`lerobot.runtime.registry`. A new
language-conditioned policy just registers its adapter there — no new
script needed. New commands should use ``lerobot-rollout --language`` (or
``lerobot-rollout --sim``); this alias remains so existing scripts do not
break.
New commands should use ``lerobot-rollout --language`` or ``--sim``.
"""
from __future__ import annotations
tests/policies/pi052/test_pi052_attention_masking.py
+7 -23
View File
@@ -56,17 +56,9 @@ def _shifted_ce(logits, labels):
eye = torch.eye(vocab_size, dtype=logits.dtype, device="cuda")
return _shifted_lin_ce(logits.cuda(), eye, labels.cuda()).cpu()
# ---------------------------------------------------------------------------
# A synthetic PI052 prefix layout: [images, prompt-lang, target-lang]
#
# indices 0-1 : 2 image tokens (att = 0)
# indices 2-4 : 3 user-prompt lang (att = 0)
# indices 5-8 : 4 supervised target lang(att = 0 from embed_prefix)
#
# ``text_labels`` covers the 7 language tokens; -100 on the prompt span,
# real ids on the 4-token target span. PaliGemma's prefix has no state
# token (unlike SmolVLA), so the lang span ends at the prefix end.
# ---------------------------------------------------------------------------
# Synthetic prefix: two image tokens, three prompt tokens, and four supervised target tokens.
# Text labels mask the prompt with -100 and cover the target through the prefix end.
N_IMAGE = 2
N_PROMPT = 3
N_TARGET = 4
@@ -95,9 +87,7 @@ def _attends(prefix_att_masks: torch.Tensor) -> torch.Tensor:
def test_mark_sets_att_on_targets_only():
"""Only the supervised target language positions flip to att=1."""
marked = _mark_target_span_causal(
_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END
)
marked = _mark_target_span_causal(_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END)
expected = [False] * PREFIX_LEN
for i in range(LANG_START + N_PROMPT, LANG_END): # target span
expected[i] = True
@@ -107,9 +97,7 @@ def test_mark_sets_att_on_targets_only():
def test_target_tokens_attend_causally_among_themselves():
"""A target token must NOT attend to later targets, but must attend
to earlier ones — genuine causal next-token prediction."""
marked = _mark_target_span_causal(
_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END
)
marked = _mark_target_span_causal(_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END)
attends = _attends(marked)
tgt = range(LANG_START + N_PROMPT, LANG_END)
for i in tgt:
@@ -122,9 +110,7 @@ def test_target_tokens_attend_causally_among_themselves():
def test_target_tokens_attend_prompt_and_images_bidirectionally():
"""Targets keep full visibility of images + the user prompt."""
marked = _mark_target_span_causal(
_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END
)
marked = _mark_target_span_causal(_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END)
attends = _attends(marked)
context = list(range(0, LANG_START + N_PROMPT)) # images + prompt
for i in range(LANG_START + N_PROMPT, LANG_END):
@@ -136,9 +122,7 @@ def test_non_target_subtask_stays_bidirectional():
"""A flow-only / non-target language span (all -100 labels) leaves the
mask untouched — the action expert reads it bidirectionally."""
all_ignored = torch.full((1, N_PROMPT + N_TARGET), -100, dtype=torch.long)
marked = _mark_target_span_causal(
_embed_prefix_att_masks(), all_ignored, LANG_START, LANG_END
)
marked = _mark_target_span_causal(_embed_prefix_att_masks(), all_ignored, LANG_START, LANG_END)
assert torch.equal(marked, _embed_prefix_att_masks())
tests/policies/pi052/test_pi052_fast_action_loss.py
+5 -13
View File
@@ -18,7 +18,7 @@
import pytest
import torch
from torch.nn import functional as F
from torch.nn import functional as F # noqa: N812
pytest.importorskip("transformers")
pytest.importorskip("liger_kernel")
@@ -40,9 +40,7 @@ def _fast_ce(logits, action_tokens, action_code_mask, predict_actions_t):
vocab_size = logits.shape[-1]
eye = torch.eye(vocab_size, dtype=logits.dtype, device="cuda")
predict = predict_actions_t.cuda() if predict_actions_t is not None else None
loss = _fast_lin_ce(
logits.cuda(), eye, action_tokens.cuda(), action_code_mask.cuda(), predict
)
loss = _fast_lin_ce(logits.cuda(), eye, action_tokens.cuda(), action_code_mask.cuda(), predict)
return loss.cpu()
@@ -67,9 +65,7 @@ def test_fast_ce_supervises_only_discrete_action_codes():
reduction="mean",
)
# Looser tolerance: the fused Triton kernel (GPU) differs from CPU eager
# F.cross_entropy at the ~1e-7 level, which exceeds the default rtol on
# these very small (~1e-4) losses.
# Allow the fused GPU kernel's ~1e-7 difference on small losses.
assert torch.allclose(loss, expected, atol=1e-5, rtol=1e-3)
@@ -77,9 +73,7 @@ def test_fast_ce_masks_non_action_samples():
"""Recipe samples with predict_actions=False do not contribute FAST loss."""
vocab_size = 8
action_tokens = torch.tensor([[1, 2, 3, 4], [1, 2, 5, 6]])
action_code_mask = torch.tensor(
[[False, False, True, True], [False, False, True, True]]
)
action_code_mask = torch.tensor([[False, False, True, True], [False, False, True, True]])
predict_actions = torch.tensor([True, False])
logits = torch.zeros(2, action_tokens.shape[1], vocab_size)
@@ -96,9 +90,7 @@ def test_fast_ce_masks_non_action_samples():
reduction="mean",
)
# Looser tolerance: the fused Triton kernel (GPU) differs from CPU eager
# F.cross_entropy at the ~1e-7 level, which exceeds the default rtol on
# these very small (~1e-4) losses.
# Allow the fused GPU kernel's ~1e-7 difference on small losses.
assert torch.allclose(loss, expected, atol=1e-5, rtol=1e-3)
tests/policies/pi052/test_pi052_text_processor.py
+1 -8
View File
@@ -63,18 +63,11 @@ def test_flatten_leaves_messages_without_tool_calls_untouched():
def test_flatten_drops_non_say_tool_calls_but_keeps_content():
weather = {"type": "function", "function": {"name": "check_weather", "arguments": {}}}
out = _flatten_say_tool_calls(
{"role": "assistant", "content": "plan only", "tool_calls": [weather]}
)
out = _flatten_say_tool_calls({"role": "assistant", "content": "plan only", "tool_calls": [weather]})
assert out["content"] == "plan only"
assert "tool_calls" not in out
# ---------------------------------------------------------------------------
# EOS-termination supervision
# ---------------------------------------------------------------------------
def test_format_messages_appends_eos_to_target_turns_only():
msgs = [
{"role": "user", "content": "pick cube"},