From 7c125c00287f23c4b740c7c652a9aab7c9afe5dc Mon Sep 17 00:00:00 2001 From: Pepijn Date: Wed, 15 Jul 2026 13:52:52 +0200 Subject: [PATCH] style: compact comments in language runtime --- src/lerobot/policies/pi052/__init__.py | 16 +- .../policies/pi052/configuration_pi052.py | 125 +------- src/lerobot/policies/pi052/debug_utils.py | 18 +- .../policies/pi052/fit_fast_tokenizer.py | 85 +----- src/lerobot/policies/pi052/flashrt_fp8.py | 73 +++-- .../policies/pi052/inference/__init__.py | 8 +- .../policies/pi052/inference/pi052_adapter.py | 12 +- src/lerobot/policies/pi052/modeling_pi052.py | 289 +++--------------- src/lerobot/policies/pi052/processor_pi052.py | 44 +-- .../policies/pi052/text_processor_pi052.py | 82 +---- src/lerobot/runtime/adapter.py | 42 +-- src/lerobot/runtime/cli.py | 165 ++-------- src/lerobot/runtime/language_runtime.py | 7 +- src/lerobot/runtime/registry.py | 7 +- src/lerobot/runtime/rerun_viz.py | 12 +- src/lerobot/runtime/sim_robocasa.py | 73 +---- .../scripts/lerobot_language_runtime.py | 7 +- .../pi052/test_pi052_attention_masking.py | 30 +- .../pi052/test_pi052_fast_action_loss.py | 18 +- .../pi052/test_pi052_text_processor.py | 9 +- 20 files changed, 205 insertions(+), 917 deletions(-) diff --git a/src/lerobot/policies/pi052/__init__.py b/src/lerobot/policies/pi052/__init__.py index dc3928bdf..7d24d772b 100644 --- a/src/lerobot/policies/pi052/__init__.py +++ b/src/lerobot/policies/pi052/__init__.py @@ -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 diff --git a/src/lerobot/policies/pi052/configuration_pi052.py b/src/lerobot/policies/pi052/configuration_pi052.py index b65f1cbae..2afd76ef4 100644 --- a/src/lerobot/policies/pi052/configuration_pi052.py +++ b/src/lerobot/policies/pi052/configuration_pi052.py @@ -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 ```` 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: diff --git a/src/lerobot/policies/pi052/debug_utils.py b/src/lerobot/policies/pi052/debug_utils.py index c0ae0c59e..569dfbed0 100644 --- a/src/lerobot/policies/pi052/debug_utils.py +++ b/src/lerobot/policies/pi052/debug_utils.py @@ -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=``. 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 diff --git a/src/lerobot/policies/pi052/fit_fast_tokenizer.py b/src/lerobot/policies/pi052/fit_fast_tokenizer.py index 298f84909..e3d103519 100644 --- a/src/lerobot/policies/pi052/fit_fast_tokenizer.py +++ b/src/lerobot/policies/pi052/fit_fast_tokenizer.py @@ -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) diff --git a/src/lerobot/policies/pi052/flashrt_fp8.py b/src/lerobot/policies/pi052/flashrt_fp8.py index e4496f9e0..91e721fba 100644 --- a/src/lerobot/policies/pi052/flashrt_fp8.py +++ b/src/lerobot/policies/pi052/flashrt_fp8.py @@ -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 diff --git a/src/lerobot/policies/pi052/inference/__init__.py b/src/lerobot/policies/pi052/inference/__init__.py index 15a9b721f..8e8f1b3c7 100644 --- a/src/lerobot/policies/pi052/inference/__init__.py +++ b/src/lerobot/policies/pi052/inference/__init__.py @@ -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 diff --git a/src/lerobot/policies/pi052/inference/pi052_adapter.py b/src/lerobot/policies/pi052/inference/pi052_adapter.py index 1cbc0bf97..72737c4e8 100644 --- a/src/lerobot/policies/pi052/inference/pi052_adapter.py +++ b/src/lerobot/policies/pi052/inference/pi052_adapter.py @@ -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: diff --git a/src/lerobot/policies/pi052/modeling_pi052.py b/src/lerobot/policies/pi052/modeling_pi052.py index d3061d3b4..082bf5e2a 100644 --- a/src/lerobot/policies/pi052/modeling_pi052.py +++ b/src/lerobot/policies/pi052/modeling_pi052.py @@ -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 -# ```` 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 ```` 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) diff --git a/src/lerobot/policies/pi052/processor_pi052.py b/src/lerobot/policies/pi052/processor_pi052.py index a3fe73c0f..6a33ff7ce 100644 --- a/src/lerobot/policies/pi052/processor_pi052.py +++ b/src/lerobot/policies/pi052/processor_pi052.py @@ -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( diff --git a/src/lerobot/policies/pi052/text_processor_pi052.py b/src/lerobot/policies/pi052/text_processor_pi052.py index bb41622e0..d3036a692 100644 --- a/src/lerobot/policies/pi052/text_processor_pi052.py +++ b/src/lerobot/policies/pi052/text_processor_pi052.py @@ -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 format (PI052 only). -# Dataset JSON uses Qwen2.5-VL's 0–1000 *normalized* grounding coords (not -# pixels — verified empirically); PaliGemma's 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 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 - # 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 text. Coords are 0–1000 normalized so this is - # camera-independent. + # Rewrite normalized VQA answers as PaliGemma text. messages = _messages_vqa_to_loc(messages, target_indices) - # Flatten ``say`` tool calls into ``...`` 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"): diff --git a/src/lerobot/runtime/adapter.py b/src/lerobot/runtime/adapter.py index 56ef4c187..1c499a594 100644 --- a/src/lerobot/runtime/adapter.py +++ b/src/lerobot/runtime/adapter.py @@ -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) diff --git a/src/lerobot/runtime/cli.py b/src/lerobot/runtime/cli.py index 76e9b1ace..152752bee 100644 --- a/src/lerobot/runtime/cli.py +++ b/src/lerobot/runtime/cli.py @@ -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= \\ - --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 ) 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) diff --git a/src/lerobot/runtime/language_runtime.py b/src/lerobot/runtime/language_runtime.py index 6c4bda8c2..dbb1b7263 100644 --- a/src/lerobot/runtime/language_runtime.py +++ b/src/lerobot/runtime/language_runtime.py @@ -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: ... diff --git a/src/lerobot/runtime/registry.py b/src/lerobot/runtime/registry.py index d297cf656..98cb7e492 100644 --- a/src/lerobot/runtime/registry.py +++ b/src/lerobot/runtime/registry.py @@ -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 diff --git a/src/lerobot/runtime/rerun_viz.py b/src/lerobot/runtime/rerun_viz.py index 1da696fcf..c7b81deba 100644 --- a/src/lerobot/runtime/rerun_viz.py +++ b/src/lerobot/runtime/rerun_viz.py @@ -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" diff --git a/src/lerobot/runtime/sim_robocasa.py b/src/lerobot/runtime/sim_robocasa.py index e97874db8..d6cf82704 100644 --- a/src/lerobot/runtime/sim_robocasa.py +++ b/src/lerobot/runtime/sim_robocasa.py @@ -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 ``) 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 diff --git a/src/lerobot/scripts/lerobot_language_runtime.py b/src/lerobot/scripts/lerobot_language_runtime.py index c197326da..708fdeb0d 100644 --- a/src/lerobot/scripts/lerobot_language_runtime.py +++ b/src/lerobot/scripts/lerobot_language_runtime.py @@ -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 diff --git a/tests/policies/pi052/test_pi052_attention_masking.py b/tests/policies/pi052/test_pi052_attention_masking.py index c12412376..96c98fb91 100644 --- a/tests/policies/pi052/test_pi052_attention_masking.py +++ b/tests/policies/pi052/test_pi052_attention_masking.py @@ -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()) diff --git a/tests/policies/pi052/test_pi052_fast_action_loss.py b/tests/policies/pi052/test_pi052_fast_action_loss.py index de99b25d6..a32a66597 100644 --- a/tests/policies/pi052/test_pi052_fast_action_loss.py +++ b/tests/policies/pi052/test_pi052_fast_action_loss.py @@ -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) diff --git a/tests/policies/pi052/test_pi052_text_processor.py b/tests/policies/pi052/test_pi052_text_processor.py index 77695e12e..f56ed16e8 100644 --- a/tests/policies/pi052/test_pi052_text_processor.py +++ b/tests/policies/pi052/test_pi052_text_processor.py @@ -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"},