style: compact comments in language runtime

2026-07-24 18:26:11 +00:00 · 2026-07-15 13:52:52 +02:00
parent 1eed8df1c4
commit 7c125c0028
20 changed files with 205 additions and 917 deletions
@@ -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
+"""π0.5 with recipe-driven language supervision and hierarchical inference.
 inference recipe on lerobot.
-Extends :class:`lerobot.policies.pi05.PI05Policy` with:
+PI052 adds supervised PaliGemma text generation, prompt dropout, and autoregressive inference to PI0.5.
 * 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.
 """
 from .configuration_pi052 import PI052Config
@@ -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
+"""PI0.5 with its PaliGemma text head enabled for hierarchical language/action training.
 hierarchical inference recipe.
-Same architecture as the existing ``PI05Policy`` (PaliGemma 2B VLM +
+The runtime generates high-level text and low-level flow-matched actions at separate rates.
 ~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.
 """
 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.
+    """PI0.5 configuration for recipe-driven text and action supervision."""
    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``).
    """
    # 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 --------------------------------------------------------
+    # Balance frequent recipe text supervision against the paper's α=10 flow weight.
    # 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.
    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) ---------------------------
+    # Optional context dropout improves tolerance to missing or stale language state.
    # 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.
    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 ------------------
+    # FAST adds discrete-action CE to the text and flow objectives from 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.
    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 --------------------------------
+    # Knowledge insulation detaches VLM K/V from action-loss gradients (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: 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 --------------------------------------------
+    # Boost sparse text-head updates while retaining PI0.5's optimizer schedule.
    # pi052 inherits π0.5's openpi-validated optimizer config (peak LR
    # 2.5e-5, cosine→2.5e-6, 1k warmup, AdamW (0.9, 0.95), wd=0.01,
    # grad_clip=1.0). The only place pi052 needs to diverge from pi05
    # is the LM-head LR multiplier: pi05 has no text supervision so the
    # head doesn't get gradients; pi052 always has text supervision
    # (subtask / memory / VQA) via the recipe, and under KI the LM head
    # only sees gradients on ~30–45% of the batch (the text-CE mask
    # share of the recipe). Under aggressive cosine decay this is too
    # weak to keep the head pinned, so it drifts back toward PaliGemma's
    # pretrained ``<loc>`` first-token bias. 5x is the documented fix
    # (see ``PI05Config.lm_head_lr_scale`` docstring); the wiring is
    # already in ``PI05Policy.get_optim_params`` — it splits the LM head
    # + tied ``embed_tokens`` into their own param group while sharing
    # the same cosine lambda, so the 5x ratio is preserved across decay.
    lm_head_lr_scale: float = 5.0
-    # Separate LRs for the VLM backbone vs the action expert (paper §III.B).
+    # Scale pretrained backbone and new action-expert groups independently; 1.0 preserves legacy behavior.
    # 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).
    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
+    # Reuse each VLM prefix across independent denoising draws; 1 restores single-draw flow.
    # 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.
    flow_num_repeats: int = 5
-    # PaLM-style z-loss on text CE. Penalises the log-partition function
+    # PaLM-style z-loss stabilizes large-vocabulary CE; 0 disables it.
    # ``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.
    text_ce_z_loss_weight: float = 1e-4
-    # Liger Triton kernels (rope + geglu + layer_norm) are now patched
+    # Liger patches are optional, process-global, and idempotent.
    # 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``.
    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
+    # Keep serialized PI052 AdamW options local because PI05Config lacks them.
    # 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.
    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
+        # Override PI0.5's expert-only default when training text.
        # 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``.
        if self.text_loss_weight > 0 and self.unfreeze_lm_head:
            self.train_expert_only = False
        if self.flow_num_repeats < 1:
@@ -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.
+    """Print supervised text predictions and token accuracy for up to ``n_samples`` rows."""
-
+    # Unwrap distributed wrappers that do not proxy custom policy methods.
    Opt-in via ``LEROBOT_DEBUG_PREDS_EVERY=<step_interval>``. Only the
    policy types that expose ``debug_text_predictions`` participate
    (currently PI052); others are silently skipped. Pretty-prints up to
    ``n_samples`` samples from the current batch, showing the prompt,
    every supervised position's (label, prediction, ✓/✗), and a
    per-sample token-accuracy summary — the cheapest "is text training
    actually learning anything" signal.
    """
    # Accelerator/DDP wraps the policy in a ``module`` attribute and
    # doesn't proxy custom methods through, so a naive
    # ``hasattr(policy, "debug_text_predictions")`` returns False on the
    # wrapper — and the helper would silently no-op. Walk through any
    # ``.module`` indirection (DDP, FSDP, ``accelerator.prepare`` wrappers)
    # to reach the raw policy that actually defines the method.
    inner = policy
    while hasattr(inner, "module") and not hasattr(inner, "debug_text_predictions"):
        inner = inner.module
@@ -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*
+Training invokes this automatically when FAST loss and automatic fitting are enabled.
 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.
 """
 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``
+# ``ProcessorMixin.save_pretrained`` writes this shared cache sentinel.
 # 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.
 _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;
+    # Only the local main process writes the tokenizer; other ranks wait on the cache sentinel.
    # 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``.
    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
+    # Read action columns directly to avoid video decoding and bound memory to sampled episodes.
    # 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.
    rng = np.random.default_rng(seed)
    actions_buf: list[np.ndarray] = []
-    # Resolve the dataset's data parquet shards directly, sidestepping
+    # Read v3 parquet shards directly to avoid split lookup failures and repeated metadata parsing.
    # ``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.
    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
+    # Load only episode indices and fixed-width actions across all shards.
    # 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.
    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;
+    # Normalize Arrow action representations into an (N, D) array.
    # ``to_numpy(zero_copy_only=False)`` produces an object array of
    # 1-D NumPy actions either way, which we stack into (N, D).
    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.
+    # Sort once because episode order is only guaranteed within each shard.
    # ``eps`` is monotonically increasing within each parquet shard but
    # we make no assumption across shards — sort once and group.
    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.
+    # Match training-time quantile normalization so FAST sees the same bounded action space.
    #
    # 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.
    flat = actions.reshape(-1, actions.shape[-1])
    q01 = np.quantile(flat, 0.01, axis=0)
    q99 = np.quantile(flat, 0.99, axis=0)
@@ -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
+The opt-in swap calibrates once on a real observation and falls back to BF16 when unavailable.
 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
 """
 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
+        # Fold fixed RMSNorm weights into the GEMM to avoid FP8 activation outliers.
-        # the kernel channel_scale = 1/(1+w). This is exact (it just moves the
+        # Adaptive RMSNorm instead uses an identity channel scale.
        # 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).
        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,
+            x_fp8,
-            self.input_scale, self.gate_up_scale, self.hidden_scale, self.down_scale,
+            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,
+            x_fp8,
-            self.input_scale, self.up_scale, self.hidden_scale, self.down_scale,
+            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 = (
+    gemma_layers = list(model.paligemma_with_expert.gemma_expert.model.layers) + list(
-        list(model.paligemma_with_expert.gemma_expert.model.layers)
+        model.paligemma_with_expert.paligemma.model.language_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
@@ -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.
+"""PI052 adapter for the policy-agnostic language 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`.
 """
 from .pi052_adapter import PI052PolicyAdapter
@@ -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.
+"""PI052 actions and text generation for the generic language 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`.
 """
 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
+        # Match the training prompt by conditioning on both subtask and discretized state.
        # training and lerobot-eval's low-level prompt ("{subtask}, State: {..};").
        # Without the state the action expert is off-distribution.
        content = subtask
        obs_state = observation.get(OBS_STATE)
        if isinstance(obs_state, torch.Tensor) and obs_state.numel() > 0:
@@ -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
+# ruff: noqa: N806, N812
 ``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.
 """
 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
+# Generic dual-expert transformer helpers live in ``lerobot.policies.pi_gemma``.
 # transformer (PaliGemmaWithExpertModel, sdpa_attention_forward,
 # compute_layer_complete, get_gemma_config) lives 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
+            # GemmaTextScaledWordEmbedding already applies sqrt(hidden_size); do not scale twice.
            # (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.
            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
+        # Build the constant suffix mask on-device to avoid a per-step host sync.
        # rather than via torch.tensor(python_list, device=cuda), which is a
        # host->device sync on every denoise step.
        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
+        # The model already checkpoints each layer; an outer checkpoint would duplicate recomputation.
        # ``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.
        (_, 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
+        # Precompute timesteps on-device to avoid a host sync per denoising step.
        # rebuilding a tensor from a Python float every step
        # (``torch.tensor(time, device=cuda)`` is a host->device sync ×num_steps).
        times = torch.tensor([1.0 + s * dt for s in range(num_steps)], dtype=torch.float32, device=device)
        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
+        # Crop appended suffix K/V after each step instead of copying the read-only prefix cache.
        # (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).
        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
+# FAST tokens occupy the high vocabulary range and must be masked during text generation.
 # tokenizer maps a FAST BPE id ``t`` to the PaliGemma vocab id
 # ``vocab_size - 1 - fast_skip_tokens - t`` (see ``TokenizerProcessorStep``),
 # so action tokens occupy the top ``_FAST_ACTION_VOCAB_SIZE`` ids below the
 # ``fast_skip_tokens`` margin. The upper part collides with the reserved
 # ``<loc>`` block; the lower part sits just under it and otherwise leaks into
 # generated text as high-codepoint gibberish (the action-trained LM head puts
 # heavy mass on these ids), so ``select_message`` masks it.
 _FAST_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
+    # Keep Liger optional until the training path needs it.
    # without liger-kernel — the call site only fires from the training
    # forward, which always pulls in the kernel.
    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
+    # Encode the mask with ignore_index to avoid a host sync and preserve graph capture.
    # ``.any().item()`` sync — Liger returns 0.0 when every position
    # is ignored, preserving graph capture for CUDA graphs.
    shift_targets = torch.where(shift_valid, shift_targets, torch.full_like(shift_targets, -100))
    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 keeps the forward equivalent while detaching VLM K/V for action-query gradients.
 # 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.
 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),
+    # ``_gated_residual`` is LeRobot's adaRMSNorm helper, not a Transformers symbol.
    # not part of HF's ``modeling_gemma``. pi05's own layer code imports
    # it from ``pi_gemma`` — mirror that here.
    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
+    # SDPA requires each fp32-generated mask slice to match its query dtype.
    # 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.
    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
+        # Patch before constructing Gemma/SigLIP layers; the operation is optional and idempotent.
        # 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.
        _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
+        # Re-enable layers PI0.5 freezes when text supervision is requested.
        # freezes a few terminal layers when ``train_expert_only`` is
        # the (default) True. We re-enable the head if the user
        # wants text supervision.
        if config.text_loss_weight > 0 and config.unfreeze_lm_head:
            self._unfreeze_lm_head()
-        # Knowledge insulation: bind a custom ``forward`` on the
+        # Bind knowledge insulation per instance so stock PI0.5 policies remain unchanged.
        # 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.
        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
+        # Size per-environment inference state lazily; the scalar mirrors env 0 for compatibility.
        # 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).
        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 —
+        # Discover terminal text layers dynamically across Transformers versions.
        # find them dynamically by walking up from the LM head so we
        # don't hard-code module names that may drift 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
+        # Delegate only unannotated batches; PI0.5 ignores recipe action-routing masks.
        # (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``).
        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
+        # Compute the host-side action-routing decision once for both flow and FAST.
        # it needs a host-side bool (one CUDA sync); compute it once and reuse for both flow and FAST
        # instead of syncing twice.
        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
-        # ------------------------------------------------------------
+        # Flow uses one fused prefix/suffix pass; text-only batches skip the suffix.
        # 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.
        # ------------------------------------------------------------
        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
+        # Keep metrics detached on-device until logging to avoid extra CUDA synchronization.
        # them to python floats only on logging steps (see update_policy's log_metrics gate).
        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
+        # Make supervised text causal rather than a bidirectional copy task.
        # genuine next-token prediction, not a bidirectional copy task
        # (see ``_mark_target_span_causal``).
        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);
+            # Gemma embedding already applies its hidden-size scale.
            # do not scale FAST action tokens again (would double-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 ----
+        # Amortized flow reuses one VLM prefix across fresh denoising targets.
        # ``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.
        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
+        # Block suffix-to-FAST attention to prevent trivial action leakage.
        # 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.
        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
+            # Position flow parallel to FAST so its RoPE offsets match inference without 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.
            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
+            # Broadcast each sample's timestep conditioning across its action chunk.
            # block's chunk tokens so each block carries its own timestep.
            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] ----
+        # Each action block attends to the non-FAST prefix and itself, never other blocks.
        # 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.
        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
+        # Restart every independent flow block after the non-FAST prefix to match inference RoPE.
        # (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.
        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(...)``
+            # Liger avoids materializing the full vocabulary logits tensor.
            # materialisation; the kernel multiplies on-the-fly and
            # never holds the full (B, T, 257k) 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
+        # Make supervised text causal before appending FAST tokens.
        # ``_mark_target_span_causal``) before the FAST tokens are
        # appended — same fix as ``_compute_all_losses_fused``.
        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
+        # Cache the image-heavy prefix; disabling the cache retains the full-recompute parity path.
        # 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 = None
        backbone = self.model.paligemma_with_expert
        lm_head = backbone.paligemma.lm_head
-        # ``_prepare_attention_masks_4d`` always returns fp32 (0.0 / -inf
+        # Use q_proj's dtype because norms and embeddings may remain fp32 while SDPA queries are bf16.
        # 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.
        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
+                # Run the full bidirectional prefix initially or whenever caching is disabled.
                # full bidirectional-prefix forward. ``current_*`` already grow
                # in the no-cache fallback below.
                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
+                # Incremental decoding feeds only the last token while retaining prefix padding masks.
                # valid cached key (``current_pad`` already includes this token),
                # so pad positions in the prefix stay masked just like the
                # recompute path.
                step_embs = new_emb
                att_2d = current_pad[:, None, :]
                att_2d_4d = self.model._prepare_attention_masks_4d(att_2d, dtype=backbone_dtype)
@@ -1894,10 +1730,7 @@ class PI052Policy(PreTrainedPolicy):
            if special_ids and len(generated) < min_new_tokens:
                for sid in special_ids:
                    logits_step[..., sid] = float("-inf")
-            # Mask FAST action tokens that fall *below* the ``<loc>`` block.
+            # Suppress FAST-only vocabulary that otherwise leaks into generated text.
            # 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.
            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 +
+            # Both paths track valid keys, but only recompute retains full embedding history.
            # position ids). Only the recompute path needs the full embedding /
            # block-attention history re-fed each step.
            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:
+        # Mirror training by appending the already normalized state to low-level prompts.
        # "User: {subtask}, State: {256-bin};"). batch state is already
        # normalized by the eval preprocessor's NormalizerProcessorStep.
        state_all = batch.get(OBS_STATE)
-        # Decide whether to (re)generate subtasks this chunk or hold the last
+        # Hold subtasks for the configured interval to match training and avoid rapid replanning.
        # 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).
        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
+        # Generate and batch one independently conditioned subtask per environment.
        # 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).
        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
+        # Feed the generated subtask verbatim, matching low-level training.
        # model's own generated subtask verbatim (this is exactly what the
        # ``low_level_execution`` recipe did at training — ``user: ${subtask}``).
        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
+        # Reuse the last valid subtask, or derive an initial imperative, when generation fails.
        # 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.
        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)
-    # ------------------------------------------------------------------
+    # PI0.5 flow-only fallback for unannotated batches.
    # 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).
    # ------------------------------------------------------------------
    @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
+        # Keep the tied LM projection and embeddings in the same optimizer group.
        # head group — boosting only the projection without the embedding pulls
        # them apart and breaks the tie PaliGemma was pre-trained with.
        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
+        # Guard before first-observation FP8 calibration to prevent recursive prediction.
        # and swap the MLPs in place. Guard set before the call so the calibration
        # forward (which re-enters predict_action_chunk) does not recurse.
        if self.config.use_flashrt_fp8_mlp and not getattr(self, "_fp8_applied", False):
            self._fp8_applied = True
            self.apply_flashrt_fp8_mlp(batch)
@@ -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:
+Without a recipe it delegates to the standard PI0.5 pipeline.
    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.
 """
 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 directly to keep optional language dependencies out of ``lerobot.processor``.
 # import it directly.
 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).
+    # Add FAST action-token supervision only when explicitly enabled.
    # 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.
    if getattr(config, "enable_fast_action_loss", False):
-        # Per Pertsch et al. 2025 (FAST [64], π0.5 §III.C): fit the
+        # Fit once on this dataset and cache by dataset, base tokenizer, and sample count.
        # 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).
        action_tokenizer_path = config.action_tokenizer_name
-        if (
+        if getattr(config, "auto_fit_fast_tokenizer", False) and dataset_repo_id is not None:
            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(
@@ -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
+PaliGemma is not chat-trained, so messages use explicit role delimiters instead of a chat template.
 ``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"``.
 """
 from __future__ import annotations
@@ -182,11 +166,7 @@ def _sample_indices(value: Any, batch_size: int) -> list[int | None]:
    return [int(value)] * batch_size
-# VQA spatial answers → PaliGemma <loc> format (PI052 only).
+# Convert normalized Qwen2.5-VL coordinates to PaliGemma's resolution-independent <loc> range.
 # Dataset JSON uses Qwen2.5-VL's 0–1000 *normalized* grounding coords (not
 # pixels — verified empirically); PaliGemma's <locNNNN> vocab is [0, 1023], so
 # ``loc_idx = round(coord / 1000 * 1023)`` is resolution-independent. Converted
 # here, not in the dataset, so the raw JSON stays backbone-agnostic.
 _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
+        # Supervise the explicit role format used again during generation.
        # role tokens at generation time, which is fine for greedy
        # decoding because the chat template is implicit in the
        # supervised target span.
        header = f"{role.capitalize()}: "
-        # A supervised target turn ends with EOS so the model learns to
+        # Include EOS only in supervised target spans so generation learns to stop.
        # terminate; the span below covers content + EOS. Non-target
        # turns (and inference) carry no EOS.
        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
+            # Preserve the transition for the plain PI0.5 prompt fallback.
            # plain Pi0.5 prompt path. We pass the transition through
            # unmodified.
            return transition
        tokenizer = self._ensure_tokenizer()
-        # Normalized proprioceptive state (set by NormalizerProcessorStep, which
+        # Add normalized proprioception to low-level prompts, matching PI0.5.
        # runs before this step). Injected into low-level action prompts so the
        # action expert sees proprioception, matching pi05's discretized State:.
        state_all = (transition.get(TransitionKey.OBSERVATION) or {}).get(OBS_STATE)
-        # VQA coords are 0–1000 normalized (Qwen2.5-VL convention) — the
+        # Normalized VQA coordinates need no camera lookup.
        # <loc> conversion is camera-resolution-independent and needs no
        # observation lookup here.
        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.
+        # Remap target indices after optional context dropout.
        # Keeps the supervised-target indices stable by re-mapping
        # after removal.
        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
+        # Rewrite normalized VQA answers as PaliGemma <loc> text.
        # PaliGemma <loc> text. Coords are 0–1000 normalized so this is
        # camera-independent.
        messages = _messages_vqa_to_loc(messages, target_indices)
-        # Flatten ``say`` tool calls into ``<say>...</say>`` text before
+        # Flatten ``say`` calls because PaliGemma receives plain text.
        # stripping, so the spoken reply is actually tokenized and
        # supervised (PaliGemma's flat prompt has no structured calls).
        messages = [_strip_blocks(_flatten_say_tool_calls(m)) for m in messages]
-        # Low-level (action-conditioning) samples get the discretized state
+        # Add state only to low-level action prompts; keep higher-level streams state-free.
        # 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.
        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
+        # Scan all streams because low-level flow may intentionally have no text target.
        # ``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``.
        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`` /
+            # Use the canonical row index to avoid identical dropout across an epoch.
            # ``render_messages_processor``. Falling back to other
            # keys silently gave every sample seed=0 → identical
            # dropout pattern across the whole epoch.
            seed_src = sample_idx if sample_idx is not None else complementary.get("index", 0)
            try:
                if hasattr(seed_src, "item"):
@@ -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
+The base adapter owns generation control and diagnostics while subclasses provide policy-specific actions and text.
 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.
 """
 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.
+    """Text-generation settings fixed for the adapter's lifetime."""
    These are configuration (set once from the CLI), not per-tick runtime
    state — they live on the adapter, never in :class:`RuntimeState`.
    """
    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.
+    """Runtime-panel rejection and repeat counters keyed by text kind."""
    Keyed by text ``kind`` (``subtask`` / ``memory`` / ...) so the same
    accounting works for any cascade shape.
    """
    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
+            # Preserve operator-provided subtasks in direct mode.
            # generate (and thus don't overwrite) it.
            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.
+    """Adapter for flat policies whose preprocessors condition actions on the operator's task."""
    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.
    """
    def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any:
        return self.policy.predict_action_chunk(observation)
@@ -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`.
+It supports a text-only REPL, real robots, and RoboCasa with local or Hub checkpoints.
 A policy wires it up with :func:`run`, passing an adapter factory
 (``policy -> LanguageConditionedPolicyAdapter``); see
 ``lerobot.scripts.lerobot_language_runtime`` for the entry point.
 Stdin is the user channel: type a task, then natural-language
 interjections. The runtime prints state changes (plan / subtask /
 memory) as they happen.
 Examples
 --------
 No-robot REPL on a Hub checkpoint — useful for sanity-checking text generation::
    uv run lerobot-rollout --language \\
        --policy.path=<repo-or-dir> \\
        --no_robot \\
        --task="please clean the kitchen"
 With a real robot::
    uv run lerobot-rollout --language \\
        --policy.path=... \\
        --robot.type=so101 --robot.port=/dev/tty.usbmodem...
 ``--policy.path`` accepts either a local directory or a Hugging Face Hub repo id.
 """
 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 ----------------------------------------
+    # ``--robot.type`` enables real-time control while stdin remains interactive.
    # 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.
    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 -------------------------------
+    # ``--sim`` uses the eval pipeline and is mutually exclusive with a robot.
    # Setting ``--sim`` flips the runtime into simulation mode: instead of
    # a real robot it drives a single RoboCasa mujoco scene, feeding the
    # eval observation/action pipeline. The operator still types prompts
    # (/action <prompt>) that the policy executes inside the chosen scene.
    # Mutually exclusive with ``--robot.type``.
    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
+# Non-observation model inputs emitted by processors such as MolmoAct2's.
 # 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.
 _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).
+    """Load a local or Hub policy, optionally with its eval processors."""
    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.
    """
    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
+    # Variable prompts trigger recompilation, and checkpointing only adds inference overhead.
    # 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.
    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
+    # Import for bundled Draccus camera and robot registrations.
    # robot config choice used by Draccus. Third-party choices were registered
    # by the top-level entry point before reaching this function.
    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
+        # Surface repeated or rejected generations as overfitting diagnostics.
        # + fully cycling queue ⇒ stuck on one subtask (memorised a phase);
        # climbing gibberish ⇒ LM head collapsed to chat-template salads.
        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
+        # Show recent generation warnings and speech oldest-first.
        # (warnings, gibberish rejections, plan speech). Last N lines,
        # 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.
+    """Keep request-level third-party logs out of the interactive prompt."""
    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).
    """
    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*
+    # Clamp warnings can fire every control tick and flood the panel.
    # 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.
    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
+    # Fork the simulator before CUDA initialization to avoid inherited EGL corruption.
    # 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.
    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
+    # Default to idle until the operator supplies a command.
    # the command line in ``paused`` mode (robot idle) unless ``--mode``
    # forces a mode. The operator drives it with /action, /pause and
    # /question.
    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
+    # Generation settings belong to the adapter rather than mutable runtime state.
    # 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.
    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
+        # Keep EGL rendering on the main thread.
        # _run_sim_interactive) — background-thread rendering corrupts EGL.
        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.
+    """Keep RoboCasa rendering on the main thread while polling stdin."""
    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.
    """
    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
+    # Keep the terminal quiet while the browser renders the rollout.
    # 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.
    _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
+                            # Clear queued actions and rearm generation for a new command.
                            # 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).
                            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.
+            # Match lerobot-eval's inference context on the main thread.
            # inference_mode matches lerobot-eval's forward context.
            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.
+    """Run steady robot control in the background and commands in the foreground."""
    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.
    """
    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
+        # An explicit initial task starts immediately; otherwise the robot stays paused.
        # immediately. Without an initial task we stay paused until the first
        # typed command (which switches to action). No flag needed.
        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.
+    """Redraw the status block and logs once per REPL turn."""
    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.
    """
    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 "...",
+            # Slash commands, including VQA questions, run inline.
            # /help. ``_handle_slash_command`` runs the VQA query inline
            # for /question (single-threaded REPL — no concurrency).
            if _handle_slash_command(runtime, line):
                last_logs = list(runtime.state.get("log_lines") or [])
                _redraw(last_logs)
@@ -117,12 +117,7 @@ class RuntimeState:
 class LanguageConditionedPolicyAdapter(Protocol):
-    """The contract the runtime loop depends on.
+    """Runtime policy contract, implemented directly or through ``BaseLanguageAdapter``."""
    :class:`lerobot.runtime.adapter.BaseLanguageAdapter` provides a
    batteries-included implementation; a policy can satisfy this protocol
    directly for full control.
    """
    def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any: ...
@@ -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.
+"""Lazy mapping from policy types to language-runtime adapters."""
 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`.
 """
 from __future__ import annotations
@@ -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).
+"""Best-effort Rerun camera visualization for local or SSH-forwarded robot rollouts."""
 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.
 """
 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
+        # Include the stream URL so the web viewer connects automatically.
        # 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).
        view_url = f"http://localhost:{web_port}/?url={url}"
        print(
            f"[runtime] rerun live view: {view_url}\n"
@@ -12,18 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""RoboCasa simulation backend for the interactive language runtime.
+"""RoboCasa backend for interactive language-conditioned rollouts.
-Lets an operator type open-ended prompts (``/action <prompt>``) and have a
+It reuses the eval observation/action pipeline while prompts control a persistent selected scene.
 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.
 """
 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
+# Two workers avoid broken single-worker EGL rendering; only env 0 is displayed.
 # 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.
 _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
+    Two workers keep EGL stable, while only env 0 is driven and displayed.
    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.
    """
    from lerobot.envs.configs import RoboCasaEnv as RoboCasaEnvConfig  # noqa: PLC0415
-    # Higher-res observation cameras => higher-quality display. The policy is
+    # The policy resizes inputs, so render_size only affects display quality and cost.
    # unaffected: its preprocessor resizes images to 224 and VISUAL norm is
    # identity, so only render cost (not behaviour) changes with render_size.
    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
+    # Keep one kitchen alive across sequential prompts.
    # so the scene doesn't truncate. The user drives it with 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``.
+    """Start an MJPEG server that shows a placeholder until ``get_frame`` returns frames."""
    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.
    """
    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
+    The environment must be created before the policy initializes CUDA.
    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).
    """
    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
+        # Bind these after runtime construction for live annotations.
        # video overlay can show the live task/subtask/memory.
        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
+        # The adapter later replaces this recipe input with its generated subtask.
        # overwrites the language tokens with its generated subtask before the
        # action forward pass, so this only needs to be present, not exact.
        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
+            # Tile env 0's action because the extra workers exist only for EGL stability.
            # 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.
            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
+            # AsyncVectorEnv resets terminated sub-environments automatically.
            # scene continues on its own — no manual reset needed here.
            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.
+        """Label and compose env 0's existing observation views without extra rendering."""
        Uses the policy's own high-res observation images (env.step already
        rendered them), so there's no extra render cost and orientation matches.
        """
        pixels = (self._last_obs or {}).get("pixels")
        if not isinstance(pixels, dict) or not pixels:
            return None
@@ -15,12 +15,7 @@
 """Compatibility entry point for ``lerobot-language-runtime``.
-Policy-agnostic: the runtime resolves the right adapter from the loaded
+New commands should use ``lerobot-rollout --language`` or ``--sim``.
 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.
 """
 from __future__ import annotations
@@ -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]
+# 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.
 #   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.
 # ---------------------------------------------------------------------------
 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(
+    marked = _mark_target_span_causal(_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END)
        _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(
+    marked = _mark_target_span_causal(_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END)
        _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(
+    marked = _mark_target_span_causal(_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END)
        _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(
+    marked = _mark_target_span_causal(_embed_prefix_att_masks(), all_ignored, LANG_START, LANG_END)
        _embed_prefix_att_masks(), all_ignored, LANG_START, LANG_END
    )
    assert torch.equal(marked, _embed_prefix_att_masks())
@@ -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(
+    loss = _fast_lin_ce(logits.cuda(), eye, action_tokens.cuda(), action_code_mask.cuda(), predict)
        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
+    # Allow the fused GPU kernel's ~1e-7 difference on small losses.
    # F.cross_entropy at the ~1e-7 level, which exceeds the default rtol on
    # these very small (~1e-4) 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(
+    action_code_mask = torch.tensor([[False, False, True, True], [False, False, True, True]])
        [[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
+    # Allow the fused GPU kernel's ~1e-7 difference on small losses.
    # F.cross_entropy at the ~1e-7 level, which exceeds the default rtol on
    # these very small (~1e-4) losses.
    assert torch.allclose(loss, expected, atol=1e-5, rtol=1e-3)
@@ -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(
+    out = _flatten_say_tool_calls({"role": "assistant", "content": "plan only", "tool_calls": [weather]})
        {"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"},