From 0fe31bfae1001b1f743fc285824d487b95318cda Mon Sep 17 00:00:00 2001
From: Pepijn <pepijn@huggingface.co>
Date: Wed, 15 Jul 2026 18:17:23 +0200
Subject: [PATCH] fix pi052 runtime and training safety

---
 src/lerobot/policies/pi052/__init__.py        |  15 +-
 .../policies/pi052/fit_fast_tokenizer.py      |  26 +-
 src/lerobot/policies/pi052/modeling_pi052.py  | 709 +++++-------------
 src/lerobot/runtime/cli.py                    |  15 +-
 src/lerobot/runtime/language_runtime.py       |  68 +-
 src/lerobot/runtime/sim_robocasa.py           |  46 +-
 src/lerobot/scripts/lerobot_train.py          |  28 +-
 src/lerobot/utils/io_utils.py                 |  74 +-
 src/lerobot/utils/logging_utils.py            |  22 +
 .../policies/pi052/test_pi052_bucketed_ce.py  |  66 +-
 tests/policies/pi052/test_pi052_import.py     |  27 +
 .../pi0_fast/test_pi0_fast_tokenizer_fit.py   |  23 +
 tests/runtime/test_language_runtime.py        |  35 +-
 tests/runtime/test_sim_robocasa.py            |  21 +
 tests/utils/test_logging_utils.py             |  18 +
 tests/utils/test_rerun_visualization.py       | 311 ++++++++
 16 files changed, 861 insertions(+), 643 deletions(-)
 create mode 100644 tests/policies/pi052/test_pi052_import.py
 create mode 100644 tests/utils/test_rerun_visualization.py

diff --git a/src/lerobot/policies/pi052/__init__.py b/src/lerobot/policies/pi052/__init__.py
index 7d24d772b..d9857c5c2 100644
--- a/src/lerobot/policies/pi052/__init__.py
+++ b/src/lerobot/policies/pi052/__init__.py
@@ -12,19 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""π0.5 with recipe-driven language supervision and hierarchical inference.
-
-PI052 adds supervised PaliGemma text generation, prompt dropout, and autoregressive inference to PI0.5.
-"""
+"""PI052 configuration; model and processors are imported lazily by their factories."""
 
 from .configuration_pi052 import PI052Config
-from .modeling_pi052 import PI052Policy
-from .processor_pi052 import make_pi052_pre_post_processors
-from .text_processor_pi052 import PI052TextTokenizerStep
 
-__all__ = [
-    "PI052Config",
-    "PI052Policy",
-    "PI052TextTokenizerStep",
-    "make_pi052_pre_post_processors",
-]
+__all__ = ["PI052Config"]
diff --git a/src/lerobot/policies/pi052/fit_fast_tokenizer.py b/src/lerobot/policies/pi052/fit_fast_tokenizer.py
index f694288c1..4e4501b9e 100644
--- a/src/lerobot/policies/pi052/fit_fast_tokenizer.py
+++ b/src/lerobot/policies/pi052/fit_fast_tokenizer.py
@@ -34,6 +34,10 @@ logger = logging.getLogger(__name__)
 _CACHE_SENTINEL = "processor_config.json"
 
 
+def _is_local_leader() -> bool:
+    return int(os.environ.get("LOCAL_RANK", "0")) == 0
+
+
 def _dataset_signature(
     dataset_repo_id: str,
     base_tokenizer_name: str,
@@ -109,8 +113,8 @@ def fit_fast_tokenizer(
         )
         return str(out_dir)
 
-    # Only the local main process writes the tokenizer; other ranks wait on the cache sentinel.
-    is_leader = int(os.environ.get("RANK", "0")) == 0 and int(os.environ.get("LOCAL_RANK", "0")) == 0
+    # Each node fits its node-local cache once; its other local ranks wait.
+    is_leader = _is_local_leader()
     if not is_leader:
         timeout_s = 1800.0  # 30 min — covers ~1024-sample fits on cold caches
         start = time.monotonic()
@@ -243,14 +247,10 @@ def resolve_fast_tokenizer(config: Any, dataset_repo_id: str | None) -> str:
     if not getattr(config, "auto_fit_fast_tokenizer", False) or dataset_repo_id is None:
         return config.action_tokenizer_name
 
-    try:
-        return fit_fast_tokenizer(
-            dataset_repo_id=dataset_repo_id,
-            cache_dir=Path(config.fast_tokenizer_cache_dir).expanduser(),
-            base_tokenizer_name=config.action_tokenizer_name,
-            n_samples=config.fast_tokenizer_fit_samples,
-            chunk_size=config.chunk_size,
-        )
-    except Exception as exc:  # noqa: BLE001
-        logger.warning("FAST tokenizer fit failed (%s); using %r instead.", exc, config.action_tokenizer_name)
-        return config.action_tokenizer_name
+    return fit_fast_tokenizer(
+        dataset_repo_id=dataset_repo_id,
+        cache_dir=Path(config.fast_tokenizer_cache_dir).expanduser(),
+        base_tokenizer_name=config.action_tokenizer_name,
+        n_samples=config.fast_tokenizer_fit_samples,
+        chunk_size=config.chunk_size,
+    )
diff --git a/src/lerobot/policies/pi052/modeling_pi052.py b/src/lerobot/policies/pi052/modeling_pi052.py
index 3b8cc9090..072ac33b9 100644
--- a/src/lerobot/policies/pi052/modeling_pi052.py
+++ b/src/lerobot/policies/pi052/modeling_pi052.py
@@ -14,21 +14,21 @@
 
 """PI0.5 with joint flow/text training and hierarchical language inference."""
 
-# ruff: noqa: N806, N812
-
 from __future__ import annotations
 
+import json
 import logging
-import math
 import types
 from collections import deque
 from contextlib import nullcontext
 from pathlib import Path
-from typing import Any, TypedDict, Unpack
+from typing import Any, Unpack
 
 import torch
-from torch import Tensor, nn
-from torch.nn import functional as F  # noqa: N812
+from safetensors.torch import load_file
+from torch import Tensor
+from torch.nn import functional
+from transformers.utils import cached_file
 
 from lerobot.configs import PreTrainedConfig
 from lerobot.utils.constants import (
@@ -40,223 +40,98 @@ from lerobot.utils.constants import (
 )
 from lerobot.utils.import_utils import require_package
 
-from ..pi05.configuration_pi05 import PI05Config
-from ..pi_gemma import PaliGemmaWithExpertModel, get_gemma_config
+from ..pi05.modeling_pi05 import (
+    ActionSelectKwargs,
+    PI05Policy,
+    PI05Pytorch as PI05PytorchBase,
+    create_sinusoidal_pos_embedding,
+    make_att_2d_masks,
+)
 from ..pretrained import PreTrainedPolicy, T
-from ..rtc.modeling_rtc import RTCProcessor
 from .configuration_pi052 import PI052Config
 
 logger = logging.getLogger(__name__)
 
-
-# Generic dual-expert transformer helpers live in ``lerobot.policies.pi_gemma``.
+_SAFETENSORS_FILE = "model.safetensors"
+_SAFETENSORS_INDEX = "model.safetensors.index.json"
 
 
-class ActionSelectKwargs(TypedDict, total=False):
-    inference_delay: int | None
-    prev_chunk_left_over: Tensor | None
-    execution_horizon: int | None
+def _resolve_weight_files(
+    pretrained_name_or_path: str | Path,
+    *,
+    force_download: bool,
+    resume_download: bool | None,
+    proxies: dict | None,
+    token: str | bool | None,
+    cache_dir: str | Path | None,
+    local_files_only: bool,
+    revision: str | None,
+) -> list[Path]:
+    model_id = str(pretrained_name_or_path)
+    local_dir = Path(model_id)
+    load_kwargs = {
+        "revision": revision,
+        "cache_dir": cache_dir,
+        "force_download": force_download,
+        "resume_download": resume_download,
+        "proxies": proxies,
+        "token": token,
+        "local_files_only": local_files_only,
+    }
 
-
-def get_safe_dtype(target_dtype, device_type):
-    """Get a safe dtype for the given device type."""
-    if device_type == "mps" and target_dtype == torch.float64:
-        return torch.float32
-    if device_type == "cpu":
-        # CPU doesn't support bfloat16, use float32 instead
-        if target_dtype == torch.bfloat16:
-            return torch.float32
-        if target_dtype == torch.float64:
-            return torch.float64
-    return target_dtype
-
-
-def create_sinusoidal_pos_embedding(  # see openpi `create_sinusoidal_pos_embedding` (exact copy)
-    time: torch.Tensor, dimension: int, min_period: float, max_period: float, device="cpu"
-) -> Tensor:
-    """Computes sine-cosine positional embedding vectors for scalar positions."""
-    if dimension % 2 != 0:
-        raise ValueError(f"dimension ({dimension}) must be divisible by 2")
-
-    if time.ndim != 1:
-        raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.")
-
-    dtype = get_safe_dtype(torch.float64, device.type)
-    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
-    period = min_period * (max_period / min_period) ** fraction
-
-    # Compute the outer product
-    scaling_factor = 1.0 / period * 2 * math.pi
-    sin_input = scaling_factor[None, :] * time[:, None]
-    return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
-
-
-def sample_beta(alpha, beta, bsize, device):  # see openpi `sample_beta` (exact copy)
-    # Beta sampling uses _sample_dirichlet which isn't implemented for MPS, so sample on CPU
-    alpha_t = torch.tensor(alpha, dtype=torch.float32)
-    beta_t = torch.tensor(beta, dtype=torch.float32)
-    dist = torch.distributions.Beta(alpha_t, beta_t)
-    return dist.sample((bsize,)).to(device)
-
-
-def make_att_2d_masks(pad_masks, att_masks):  # see openpi `make_att_2d_masks` (exact copy)
-    """Copied from big_vision.
-
-    Tokens can attend to valid inputs tokens which have a cumulative mask_ar
-    smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to
-    setup several types of attention, for example:
-
-      [[1 1 1 1 1 1]]: pure causal attention.
-
-      [[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between
-          themselves and the last 3 tokens have a causal attention. The first
-          entry could also be a 1 without changing behaviour.
-
-      [[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a
-          block can attend all previous blocks and all tokens on the same block.
-
-    Args:
-      input_mask: bool[B, N] true if its part of the input, false if padding.
-      mask_ar: int32[B, N] mask that's 1 where previous tokens cannot depend on
-        it and 0 where it shares the same attention mask as the previous token.
-    """
-    if att_masks.ndim != 2:
-        raise ValueError(att_masks.ndim)
-    if pad_masks.ndim != 2:
-        raise ValueError(pad_masks.ndim)
-
-    cumsum = torch.cumsum(att_masks, dim=1)
-    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
-    pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
-    return att_2d_masks & pad_2d_masks
-
-
-def pad_vector(vector, new_dim):
-    """Pad the last dimension of a vector to new_dim with zeros.
-
-    Can be (batch_size x sequence_length x features_dimension)
-    or (batch_size x features_dimension)
-    """
-    if vector.shape[-1] >= new_dim:
-        return vector
-    return F.pad(vector, (0, new_dim - vector.shape[-1]))
-
-
-def resize_with_pad_torch(  # see openpi `resize_with_pad_torch` (exact copy)
-    images: torch.Tensor,
-    height: int,
-    width: int,
-    mode: str = "bilinear",
-) -> torch.Tensor:
-    """PyTorch version of resize_with_pad. Resizes an image to a target height and width without distortion
-    by padding with black. If the image is float32, it must be in the range [-1, 1].
-
-    Args:
-        images: Tensor of shape [*b, h, w, c] or [*b, c, h, w]
-        height: Target height
-        width: Target width
-        mode: Interpolation mode ('bilinear', 'nearest', etc.)
-
-    Returns:
-        Resized and padded tensor with same shape format as input
-    """
-    # Check if input is in channels-last format [*b, h, w, c] or channels-first [*b, c, h, w]
-    if images.shape[-1] <= 4:  # Assume channels-last format
-        channels_last = True
-        if images.dim() == 3:
-            images = images.unsqueeze(0)  # Add batch dimension
-        images = images.permute(0, 3, 1, 2)  # [b, h, w, c] -> [b, c, h, w]
+    if local_dir.is_dir():
+        index_path = local_dir / _SAFETENSORS_INDEX
+        single_path = local_dir / _SAFETENSORS_FILE
     else:
-        channels_last = False
-        if images.dim() == 3:
-            images = images.unsqueeze(0)  # Add batch dimension
-
-    batch_size, channels, cur_height, cur_width = images.shape
-
-    # Calculate resize ratio
-    ratio = max(cur_width / width, cur_height / height)
-    resized_height = int(cur_height / ratio)
-    resized_width = int(cur_width / ratio)
-
-    # Resize
-    resized_images = F.interpolate(
-        images,
-        size=(resized_height, resized_width),
-        mode=mode,
-        align_corners=False if mode == "bilinear" else None,
-    )
-
-    # Handle dtype-specific clipping
-    if images.dtype == torch.uint8:
-        resized_images = torch.round(resized_images).clamp(0, 255).to(torch.uint8)
-    elif images.dtype == torch.float32:
-        resized_images = resized_images.clamp(0.0, 1.0)
-    else:
-        raise ValueError(f"Unsupported image dtype: {images.dtype}")
-
-    # Calculate padding
-    pad_h0, remainder_h = divmod(height - resized_height, 2)
-    pad_h1 = pad_h0 + remainder_h
-    pad_w0, remainder_w = divmod(width - resized_width, 2)
-    pad_w1 = pad_w0 + remainder_w
-
-    # Pad
-    constant_value = 0 if images.dtype == torch.uint8 else 0.0
-    padded_images = F.pad(
-        resized_images,
-        (pad_w0, pad_w1, pad_h0, pad_h1),  # left, right, top, bottom
-        mode="constant",
-        value=constant_value,
-    )
-
-    # Convert back to original format if needed
-    if channels_last:
-        padded_images = padded_images.permute(0, 2, 3, 1)  # [b, c, h, w] -> [b, h, w, c]
-
-    return padded_images
-
-
-class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
-    """Core PI05 PyTorch model."""
-
-    def __init__(self, config: PI05Config, rtc_processor: RTCProcessor | None = None):
-        super().__init__()
-        self.config = config
-        self.rtc_processor = rtc_processor
-
-        paligemma_config = get_gemma_config(config.paligemma_variant)
-        action_expert_config = get_gemma_config(config.action_expert_variant)
-
-        if config.image_resolution[0] != config.image_resolution[1]:
-            raise ValueError(
-                f"PaliGemma expects square image resolution, invalid resolution: {config.image_resolution}"
-            )
-
-        self.paligemma_with_expert = PaliGemmaWithExpertModel(
-            paligemma_config,
-            action_expert_config,
-            use_adarms=[False, True],
-            precision=config.dtype,
-            image_size=config.image_resolution[0],
-            freeze_vision_encoder=config.freeze_vision_encoder,
-            train_expert_only=config.train_expert_only,
+        resolved_index = cached_file(
+            model_id,
+            _SAFETENSORS_INDEX,
+            _raise_exceptions_for_missing_entries=False,
+            **load_kwargs,
         )
+        index_path = Path(resolved_index) if resolved_index is not None else None
+        single_path = None
+        if index_path is None:
+            resolved_file = cached_file(model_id, _SAFETENSORS_FILE, **load_kwargs)
+            single_path = Path(resolved_file) if resolved_file is not None else None
 
-        self.action_in_proj = nn.Linear(config.max_action_dim, action_expert_config.width)
-        self.action_out_proj = nn.Linear(action_expert_config.width, config.max_action_dim)
+    if index_path is None or not index_path.is_file():
+        if single_path is None or not single_path.is_file():
+            raise FileNotFoundError(f"No {_SAFETENSORS_FILE} found in {model_id!r}.")
+        return [single_path]
 
-        self.time_mlp_in = nn.Linear(action_expert_config.width, action_expert_config.width)
-        self.time_mlp_out = nn.Linear(action_expert_config.width, action_expert_config.width)
+    index = json.loads(index_path.read_text())
+    shard_names = sorted(set(index.get("weight_map", {}).values()))
+    if not shard_names:
+        raise ValueError(f"Invalid safetensors index without a weight_map: {index_path}")
+    if local_dir.is_dir():
+        files = [local_dir / name for name in shard_names]
+    else:
+        files = []
+        for name in shard_names:
+            resolved_file = cached_file(model_id, name, **load_kwargs)
+            if resolved_file is None:
+                raise FileNotFoundError(f"Checkpoint shard {name!r} not found in {model_id!r}.")
+            files.append(Path(resolved_file))
+    missing = [str(path) for path in files if not path.is_file()]
+    if missing:
+        raise FileNotFoundError(f"Missing checkpoint shards: {missing}")
+    return files
 
-        # Initialize gradient checkpointing flag
-        self.gradient_checkpointing_enabled = False
 
-        # Compile model if requested
-        if config.compile_model:
-            torch.set_float32_matmul_precision("high")
-            self.sample_actions = torch.compile(self.sample_actions, mode=config.compile_mode)
-            # Also compile the main forward pass used during training
-            self.forward = torch.compile(self.forward, mode=config.compile_mode)
+def _load_weight_files(files: list[Path]) -> dict[str, Tensor]:
+    state_dict: dict[str, Tensor] = {}
+    for path in files:
+        shard = load_file(path)
+        overlap = state_dict.keys() & shard.keys()
+        if overlap:
+            raise ValueError(f"Duplicate checkpoint keys in {path}: {sorted(overlap)[:5]}")
+        state_dict.update(shard)
+    return state_dict
+
+
+class PI05Pytorch(PI05PytorchBase):  # see openpi `PI0Pytorch`
+    """Core PI05 PyTorch model."""
 
     def gradient_checkpointing_enable(self):
         """Enable gradient checkpointing for memory optimization."""
@@ -276,17 +151,6 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = False
         logging.info("Disabled gradient checkpointing for PI05Pytorch model")
 
-    def _rtc_enabled(self):
-        return self.config.rtc_config is not None and self.config.rtc_config.enabled
-
-    def _apply_checkpoint(self, func, *args, **kwargs):
-        """Helper method to apply gradient checkpointing if enabled."""
-        if self.gradient_checkpointing_enabled and self.training:
-            return torch.utils.checkpoint.checkpoint(
-                func, *args, use_reentrant=False, preserve_rng_state=False, **kwargs
-            )
-        return func(*args, **kwargs)
-
     def _prepare_attention_masks_4d(self, att_2d_masks, dtype=None):
         """Helper method to prepare 4D attention masks for transformer."""
         att_2d_masks_4d = att_2d_masks[:, None, :, :]
@@ -295,22 +159,6 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             result = result.to(dtype=dtype)
         return result
 
-    def sample_noise(self, shape, device):
-        return torch.normal(
-            mean=0.0,
-            std=1.0,
-            size=shape,
-            dtype=torch.float32,
-            device=device,
-        )
-
-    def sample_time(self, bsize, device):
-        time_beta = sample_beta(
-            self.config.time_sampling_beta_alpha, self.config.time_sampling_beta_beta, bsize, device
-        )
-        time = time_beta * self.config.time_sampling_scale + self.config.time_sampling_offset
-        return time.to(dtype=torch.float32, device=device)
-
     def embed_prefix(
         self, images, img_masks, tokens, masks
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
@@ -374,9 +222,9 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
 
         def time_mlp_func(time_emb):
             x = self.time_mlp_in(time_emb)
-            x = F.silu(x)
+            x = functional.silu(x)
             x = self.time_mlp_out(x)
-            return F.silu(x)
+            return functional.silu(x)
 
         time_emb = self._apply_checkpoint(time_mlp_func, time_emb)
         action_time_emb = action_emb
@@ -449,7 +297,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
 
         v_t = self._apply_checkpoint(action_out_proj_func, suffix_out)
 
-        return F.mse_loss(u_t, v_t, reduction="none")
+        return functional.mse_loss(u_t, v_t, reduction="none")
 
     @torch.no_grad()  # see openpi `sample_actions` (slightly adapted)
     def sample_actions(
@@ -610,11 +458,13 @@ def _enable_hf_kernels() -> None:
     logger.info("PI052: HF kernels (Liger) enabled — rope, geglu fused.")
 
 
-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."""
+def _reduce_action_loss(per_sample: Tensor, predict_actions_t: Tensor | None, reduction: str) -> Tensor:
+    """Mask non-action samples and apply the requested batch reduction."""
     if predict_actions_t is None:
-        return per_sample.mean()
+        return per_sample if reduction == "none" else per_sample.mean()
     mask = predict_actions_t.to(per_sample.dtype)
+    if reduction == "none":
+        return per_sample * mask
     return (per_sample * mask).sum() / mask.sum().clamp(min=1.0)
 
 
@@ -631,7 +481,7 @@ def _lin_ce_small(
     """Small-N linear CE on materialized logits (see ``_lin_ce_flat``)."""
     logits = (flat_hidden @ lm_head_weight.t()).float()
     n_valid = (flat_labels != -100).sum().clamp(min=1)
-    loss = F.cross_entropy(logits, flat_labels, ignore_index=-100, reduction="sum") / n_valid
+    loss = functional.cross_entropy(logits, flat_labels, ignore_index=-100, reduction="sum") / n_valid
     if z_loss_weight > 0:
         lse = torch.logsumexp(logits, dim=-1)
         valid = (flat_labels != -100).to(lse.dtype)
@@ -666,9 +516,9 @@ def _lin_ce_flat(
 
         if compact_rows == 0:
             return _lin_ce_flat(
-                F.pad(compact_hidden, (0, 0, 0, 1)),
+                functional.pad(compact_hidden, (0, 0, 0, 1)),
                 lm_head_weight,
-                F.pad(compact_labels, (0, 1), value=-100),
+                functional.pad(compact_labels, (0, 1), value=-100),
                 z_loss_weight,
                 compiled=compiled,
             )
@@ -686,8 +536,8 @@ def _lin_ce_flat(
                 labels_chunk = compact_labels[start:end]
                 pad_rows = chunk_rows - rows
                 if pad_rows:
-                    hidden_chunk = F.pad(hidden_chunk, (0, 0, 0, pad_rows))
-                    labels_chunk = F.pad(labels_chunk, (0, pad_rows), value=-100)
+                    hidden_chunk = functional.pad(hidden_chunk, (0, 0, 0, pad_rows))
+                    labels_chunk = functional.pad(labels_chunk, (0, pad_rows), value=-100)
                 chunk_loss = _lin_ce_flat(
                     hidden_chunk,
                     lm_head_weight,
@@ -721,10 +571,24 @@ def _shifted_lin_ce(
     labels: Tensor,
     z_loss_weight: float = 0.0,
     compiled: bool = False,
+    reduction: str = "mean",
 ) -> Tensor:
     """Compute next-token CE through the shape-aware linear-CE dispatcher."""
     shift_hidden = hidden[:, :-1, :].contiguous()
     shift_labels = labels[:, 1:].contiguous().long()
+    if reduction == "none":
+        return torch.stack(
+            [
+                _lin_ce_flat(
+                    sample_hidden.to(lm_head_weight.dtype),
+                    lm_head_weight,
+                    sample_labels,
+                    z_loss_weight,
+                    compiled=compiled,
+                )
+                for sample_hidden, sample_labels in zip(shift_hidden, shift_labels, strict=True)
+            ]
+        )
     batch_size, target_length, hidden_size = shift_hidden.shape
     flat_hidden = shift_hidden.reshape(batch_size * target_length, hidden_size)
     flat_labels = shift_labels.reshape(batch_size * target_length)
@@ -755,6 +619,7 @@ def _fast_lin_ce(
     action_code_mask: Tensor,
     predict_actions_t: Tensor | None,
     compiled: bool = False,
+    reduction: str = "mean",
 ) -> Tensor:
     """Compute FAST token CE over the enabled action-code positions."""
     shift_hidden = hidden[:, :-1, :].contiguous()
@@ -766,6 +631,18 @@ def _fast_lin_ce(
     # Encode the mask with ignore_index to avoid a host sync and preserve graph capture.
     shift_targets = torch.where(shift_valid, shift_targets, torch.full_like(shift_targets, -100))
 
+    if reduction == "none":
+        return torch.stack(
+            [
+                _lin_ce_flat(
+                    sample_hidden.to(lm_head_weight.dtype),
+                    lm_head_weight,
+                    sample_labels,
+                    compiled=compiled,
+                )
+                for sample_hidden, sample_labels in zip(shift_hidden, shift_targets, strict=True)
+            ]
+        )
     batch_size, target_length, hidden_size = shift_hidden.shape
     flat_hidden = shift_hidden.reshape(batch_size * target_length, hidden_size).to(lm_head_weight.dtype)
     flat_labels = shift_targets.reshape(batch_size * target_length)
@@ -790,7 +667,9 @@ def _get_flex_kernel_options(device: torch.device) -> dict | None:
         return None
     device_index = device.index if device.index is not None else torch.cuda.current_device()
     if device_index not in _flex_kernel_options:
-        smem = torch.cuda.get_device_properties(device_index).shared_memory_per_block_optin
+        smem = torch.cuda.get_device_properties(
+            device_index
+        ).shared_memory_per_block_optin  # spellchecker:disable-line
         _flex_kernel_options[device_index] = _FLEX_SHRUNK_TILES if smem < 128 * 1024 else None
     return _flex_kernel_options[device_index]
 
@@ -1206,13 +1085,11 @@ def _paligemma_forward_ki(
     return [outputs_embeds[0], outputs_embeds[1]], None
 
 
-class PI052Policy(PreTrainedPolicy):
+class PI052Policy(PI05Policy):
     """π0.5 with the PaliGemma LM head re-enabled.
 
-    Self-contained: the PI0.5 backbone (PaliGemmaWithExpertModel / PI05Pytorch)
-    is vendored in ``pi05_backbone.py`` and the PI05Policy wrapper logic is
-    inlined directly here, so this policy does not depend on or inherit from
-    ``lerobot.policies.pi05`` (which stays identical to ``main``).
+    It inherits unchanged PI0.5 policy behavior and replaces the core model with
+    the joint flow/text implementation below.
     """
 
     config_class = PI052Config
@@ -1222,9 +1099,8 @@ class PI052Policy(PreTrainedPolicy):
         # Patch before constructing Gemma/SigLIP layers; the operation is optional and idempotent.
         _enable_hf_kernels()
 
-        # ---- inlined PI05Policy.__init__ ----------------------------------
         require_package("transformers", extra="pi")
-        super().__init__(config)
+        PreTrainedPolicy.__init__(self, config)
         config.validate_features()
         self.config = config
         self.init_rtc_processor()
@@ -1233,7 +1109,6 @@ class PI052Policy(PreTrainedPolicy):
             self.model.gradient_checkpointing_enable()
         self.model.to(config.device)
         self.reset()
-        # ---- end inlined PI05Policy.__init__ ------------------------------
 
         # Re-enable layers PI0.5 freezes when text supervision is requested.
         if config.text_loss_weight > 0 and config.unfreeze_lm_head:
@@ -1341,6 +1216,8 @@ class PI052Policy(PreTrainedPolicy):
         reduction: str = "mean",
     ) -> tuple[Tensor, dict]:
         """Compute the enabled flow, text and FAST training losses."""
+        if reduction not in {"mean", "none"}:
+            raise ValueError(f"Unsupported loss reduction: {reduction!r}")
         text_labels = batch.get("text_labels")
         predict_actions_t = batch.get("predict_actions")
 
@@ -1389,14 +1266,15 @@ class PI052Policy(PreTrainedPolicy):
                 action_mask=action_mask if run_fast else None,
                 action_code_mask=action_code_mask if run_fast else None,
                 predict_actions_t=predict_actions_t,
+                reduction=reduction,
             )
-            loss_dict["flow_loss"] = flow_loss.detach()
+            loss_dict["flow_loss"] = flow_loss.detach().mean()
             total = self.config.flow_loss_weight * flow_loss
             if text_loss is not None:
-                loss_dict["text_loss"] = text_loss.detach()
+                loss_dict["text_loss"] = text_loss.detach().mean()
                 total = total + self.config.text_loss_weight * text_loss
             if fast_loss is not None:
-                loss_dict["fast_action_loss"] = fast_loss.detach()
+                loss_dict["fast_action_loss"] = fast_loss.detach().mean()
                 total = total + self.config.fast_action_loss_weight * fast_loss
         elif run_text or run_fast:
             text_loss, fast_loss = self._compute_text_and_fast_loss(
@@ -1406,13 +1284,14 @@ class PI052Policy(PreTrainedPolicy):
                 action_mask=action_mask if run_fast else None,
                 action_code_mask=action_code_mask if run_fast else None,
                 predict_actions_t=predict_actions_t,
+                reduction=reduction,
             )
             if text_loss is not None:
-                loss_dict["text_loss"] = text_loss.detach()
+                loss_dict["text_loss"] = text_loss.detach().mean()
                 weighted = self.config.text_loss_weight * text_loss
                 total = weighted if total is None else total + weighted
             if fast_loss is not None:
-                loss_dict["fast_action_loss"] = fast_loss.detach()
+                loss_dict["fast_action_loss"] = fast_loss.detach().mean()
                 weighted = self.config.fast_action_loss_weight * fast_loss
                 total = weighted if total is None else total + weighted
 
@@ -1426,9 +1305,7 @@ class PI052Policy(PreTrainedPolicy):
             )
 
         # Keep metrics detached on-device until logging to avoid extra CUDA synchronization.
-        loss_dict["loss"] = total.detach() if total.dim() == 0 else float("nan")
-        if reduction == "none":
-            return total.expand(batch[OBS_LANGUAGE_TOKENS].shape[0]), loss_dict
+        loss_dict["loss"] = total.detach().mean()
         return total, loss_dict
 
     def _compute_all_losses_fused(
@@ -1439,6 +1316,7 @@ class PI052Policy(PreTrainedPolicy):
         action_mask: Tensor | None,
         action_code_mask: Tensor | None,
         predict_actions_t: Tensor | None = None,
+        reduction: str = "mean",
     ) -> tuple[Tensor, Tensor | None, Tensor | None]:
         """Compute flow, text and FAST losses from one shared prefix."""
         # ---- preamble (mirrors PI05Pytorch.forward) ------------------
@@ -1496,6 +1374,7 @@ class PI052Policy(PreTrainedPolicy):
                 predict_actions_t,
                 num_repeats,
                 suppress_prefix_grads=suppress_prefix_grads,
+                reduction=reduction,
             )
         else:
             prefix_out, flow_loss = self._combined_prefix_and_flow(
@@ -1507,10 +1386,17 @@ class PI052Policy(PreTrainedPolicy):
                 fast_len,
                 predict_actions_t,
                 suppress_prefix_grads=suppress_prefix_grads,
+                reduction=reduction,
             )
 
         text_loss, fast_loss = self._prefix_ce_losses(
-            prefix_out, text_labels, action_tokens, action_code_mask, fast_len, predict_actions_t
+            prefix_out,
+            text_labels,
+            action_tokens,
+            action_code_mask,
+            fast_len,
+            predict_actions_t,
+            reduction,
         )
         return flow_loss, text_loss, fast_loss
 
@@ -1524,6 +1410,7 @@ class PI052Policy(PreTrainedPolicy):
         fast_len: int,
         predict_actions_t: Tensor | None,
         suppress_prefix_grads: bool = False,
+        reduction: str = "mean",
     ) -> tuple[Tensor, Tensor]:
         """Run the single-repeat combined prefix and action path."""
         from lerobot.utils.constants import ACTION  # noqa: PLC0415
@@ -1576,13 +1463,13 @@ class PI052Policy(PreTrainedPolicy):
         # ---- flow loss (mirrors PI05Pytorch.forward) ----------------
         suffix_out_slice = suffix_out[:, -self.model.config.chunk_size :].to(dtype=torch.float32)
         v_t = self.model.action_out_proj(suffix_out_slice)
-        flow_per_dim = F.mse_loss(u_t, v_t, reduction="none")
+        flow_per_dim = functional.mse_loss(u_t, v_t, reduction="none")
         # Truncate to the actual action dimensionality (PI05 pads
         # internally to max_action_dim).
         original_action_dim = self.config.output_features[ACTION].shape[0]
         flow_per_dim = flow_per_dim[:, :, :original_action_dim]
         per_sample_flow = flow_per_dim.mean(dim=(1, 2))
-        flow_loss = _mask_per_sample(per_sample_flow, predict_actions_t)
+        flow_loss = _reduce_action_loss(per_sample_flow, predict_actions_t, reduction)
         return prefix_out, flow_loss
 
     def _ki_forward_kwargs(self, suppress_prefix_grads: bool = False, flex_masks=None) -> dict[str, Any]:
@@ -1609,6 +1496,7 @@ class PI052Policy(PreTrainedPolicy):
         predict_actions_t: Tensor | None,
         num_repeats: int,
         suppress_prefix_grads: bool = False,
+        reduction: str = "mean",
     ) -> tuple[Tensor, Tensor]:
         """Run K independent action draws against one shared VLM prefix."""
         from lerobot.utils.constants import ACTION  # noqa: PLC0415
@@ -1711,9 +1599,9 @@ class PI052Policy(PreTrainedPolicy):
         original_action_dim = self.config.output_features[ACTION].shape[0]
         v_t = model.action_out_proj(suffix_out.to(dtype=torch.float32))
         v_t = v_t.view(batch_size, k, chunk, -1)  # (B, k, chunk, motor)
-        flow_per_dim = F.mse_loss(u_t, v_t, reduction="none")[..., :original_action_dim]
+        flow_per_dim = functional.mse_loss(u_t, v_t, reduction="none")[..., :original_action_dim]
         per_sample_flow = flow_per_dim.mean(dim=(1, 2, 3))
-        flow_loss = _mask_per_sample(per_sample_flow, predict_actions_t)
+        flow_loss = _reduce_action_loss(per_sample_flow, predict_actions_t, reduction)
         return prefix_out, flow_loss
 
     def _prefix_ce_losses(
@@ -1724,6 +1612,7 @@ class PI052Policy(PreTrainedPolicy):
         action_code_mask: Tensor | None,
         fast_len: int,
         predict_actions_t: Tensor | None,
+        reduction: str = "mean",
     ) -> tuple[Tensor | None, Tensor | None]:
         """Compute enabled text and FAST losses from the shared prefix output."""
         lm_head = self.model.paligemma_with_expert.paligemma.lm_head
@@ -1742,6 +1631,7 @@ class PI052Policy(PreTrainedPolicy):
                 text_labels,
                 z_loss_weight=getattr(self.config, "text_ce_z_loss_weight", 0.0),
                 compiled=self.config.use_compiled_text_ce,
+                reduction=reduction,
             )
 
         fast_loss: Tensor | None = None
@@ -1754,6 +1644,7 @@ class PI052Policy(PreTrainedPolicy):
                 action_code_mask,
                 predict_actions_t,
                 compiled=self.config.use_compiled_text_ce,
+                reduction=reduction,
             )
 
         return text_loss, fast_loss
@@ -1766,6 +1657,7 @@ class PI052Policy(PreTrainedPolicy):
         action_mask: Tensor | None,
         action_code_mask: Tensor | None,
         predict_actions_t: Tensor | None = None,
+        reduction: str = "mean",
     ) -> tuple[Tensor | None, Tensor | None]:
         """Single prefix forward → text CE + FAST CE.
 
@@ -1848,6 +1740,7 @@ class PI052Policy(PreTrainedPolicy):
                 text_labels,
                 z_loss_weight=getattr(self.config, "text_ce_z_loss_weight", 0.0),
                 compiled=self.config.use_compiled_text_ce,
+                reduction=reduction,
             )
 
         fast_loss: Tensor | None = None
@@ -1860,6 +1753,7 @@ class PI052Policy(PreTrainedPolicy):
                 action_code_mask,
                 predict_actions_t,
                 compiled=self.config.use_compiled_text_ce,
+                reduction=reduction,
             )
 
         return text_loss, fast_loss
@@ -2249,16 +2143,10 @@ class PI052Policy(PreTrainedPolicy):
         strict: bool = True,
         **kwargs,
     ) -> T:
-        """Override the from_pretrained method to handle key remapping and display important disclaimer."""
-        print(
-            "The PI05 model is a direct port of the OpenPI implementation. \n"
-            "This implementation follows the original OpenPI structure for compatibility. \n"
-            "Original implementation: https://github.com/Physical-Intelligence/openpi"
-        )
+        """Load a PI05/PI052 checkpoint, including sharded safetensors checkpoints."""
         if pretrained_name_or_path is None:
             raise ValueError("pretrained_name_or_path is required")
 
-        # Use provided config if available, otherwise create default config
         if config is None:
             config = PreTrainedConfig.from_pretrained(
                 pretrained_name_or_path=pretrained_name_or_path,
@@ -2272,157 +2160,40 @@ class PI052Policy(PreTrainedPolicy):
                 **kwargs,
             )
 
-        # Initialize model without loading weights
-        # Check if dataset_stats were provided in kwargs
         model = cls(config, **kwargs)
+        files = _resolve_weight_files(
+            pretrained_name_or_path,
+            force_download=force_download,
+            resume_download=resume_download,
+            proxies=proxies,
+            token=token,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+            revision=revision,
+        )
+        fixed_state_dict = model._fix_pytorch_state_dict_keys(_load_weight_files(files), model.config)
+        remapped_state_dict = {
+            key if key.startswith("model.") else f"model.{key}": value
+            for key, value in fixed_state_dict.items()
+        }
 
-        # Load state dict (expects keys with "model." prefix)
-        try:
-            print(f"Loading model from: {pretrained_name_or_path}")
-            try:
-                from transformers.utils import cached_file
-
-                resolved_file = cached_file(
-                    pretrained_name_or_path,
-                    "model.safetensors",
-                    cache_dir=kwargs.get("cache_dir"),
-                    force_download=kwargs.get("force_download", False),
-                    resume_download=kwargs.get("resume_download"),
-                    proxies=kwargs.get("proxies"),
-                    token=kwargs.get("token"),
-                    revision=kwargs.get("revision"),
-                    local_files_only=kwargs.get("local_files_only", False),
-                )
-                from safetensors.torch import load_file
-
-                original_state_dict = load_file(resolved_file)
-                print("✓ Loaded state dict from model.safetensors")
-            except Exception as e:
-                print(f"Could not load state dict from remote files: {e}")
-                print("Returning model without loading pretrained weights")
-                return model
-
-            # First, fix any key differences (see openpi model.py, _fix_pytorch_state_dict_keys)
-            fixed_state_dict = model._fix_pytorch_state_dict_keys(original_state_dict, model.config)
-
-            # Then add "model." prefix for all keys that don't already have it
-            remapped_state_dict = {}
-            remap_count = 0
-
-            for key, value in fixed_state_dict.items():
-                if not key.startswith("model."):
-                    new_key = f"model.{key}"
-                    remapped_state_dict[new_key] = value
-                    remap_count += 1
-                else:
-                    remapped_state_dict[key] = value
-
-            if remap_count > 0:
-                print(f"Remapped {remap_count} state dict keys")
-
-            lm_head_key = "model.paligemma_with_expert.paligemma.lm_head.weight"
-            embed_tokens_key = (
-                "model.paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"
-            )
-            if lm_head_key not in remapped_state_dict and embed_tokens_key in remapped_state_dict:
-                remapped_state_dict[lm_head_key] = remapped_state_dict[embed_tokens_key].clone().float()
-                print("Initialized PaliGemma lm_head from language token embeddings")
-            elif lm_head_key in remapped_state_dict:
-                remapped_state_dict[lm_head_key] = remapped_state_dict[lm_head_key].float()
-
-            # Load the remapped state dict into the model
-            missing_keys, unexpected_keys = model.load_state_dict(remapped_state_dict, strict=strict)
+        lm_head_key = "model.paligemma_with_expert.paligemma.lm_head.weight"
+        embed_tokens_key = "model.paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"
+        if lm_head_key not in remapped_state_dict and embed_tokens_key in remapped_state_dict:
+            remapped_state_dict[lm_head_key] = remapped_state_dict[embed_tokens_key].clone().float()
+        elif lm_head_key in remapped_state_dict:
+            remapped_state_dict[lm_head_key] = remapped_state_dict[lm_head_key].float()
 
+        missing_keys, unexpected_keys = model.load_state_dict(remapped_state_dict, strict=strict)
+        if not strict:
             if missing_keys:
-                print(f"Missing keys when loading state dict: {len(missing_keys)} keys")
-                if len(missing_keys) <= 5:
-                    for key in missing_keys:
-                        print(f"  - {key}")
-                else:
-                    for key in missing_keys[:5]:
-                        print(f"  - {key}")
-                    print(f"  ... and {len(missing_keys) - 5} more")
-
+                logger.warning("Missing PI052 checkpoint keys: %s", missing_keys)
             if unexpected_keys:
-                print(f"Unexpected keys when loading state dict: {len(unexpected_keys)} keys")
-                if len(unexpected_keys) <= 5:
-                    for key in unexpected_keys:
-                        print(f"  - {key}")
-                else:
-                    for key in unexpected_keys[:5]:
-                        print(f"  - {key}")
-                    print(f"  ... and {len(unexpected_keys) - 5} more")
-
-            if not missing_keys and not unexpected_keys:
-                print("All keys loaded successfully!")
-
-        except Exception as e:
-            print(f"Warning: Could not load state dict: {e}")
-
+                logger.warning("Unexpected PI052 checkpoint keys: %s", unexpected_keys)
+        model.to(config.device)
+        model.eval()
         return model
 
-    def _fix_pytorch_state_dict_keys(
-        self, state_dict, model_config
-    ):  # see openpi `BaseModelConfig, _fix_pytorch_state_dict_keys`
-        """Fix state dict keys to match current model architecture."""
-        import re
-
-        fixed_state_dict = {}
-
-        for key, value in state_dict.items():
-            new_key = key
-
-            # Handle layer norm structure changes: .weight -> .dense.weight + .dense.bias
-            # For gemma expert layers
-            if re.match(
-                r"paligemma_with_expert\.gemma_expert\.model\.layers\.\d+\.(input_layernorm|post_attention_layernorm)\.weight",
-                key,
-            ):
-                # Check if the model actually has adaRMS enabled for the expert
-                expert_uses_adarms = getattr(
-                    self.model.paligemma_with_expert.gemma_expert.config, "use_adarms", False
-                )
-                if expert_uses_adarms:
-                    logging.warning(f"Skipping layer norm key (adaRMS mismatch): {key}")
-                    continue
-
-            if re.match(r"paligemma_with_expert\.gemma_expert\.model\.norm\.weight", key):
-                # Check if the model actually has adaRMS enabled for the expert
-                expert_uses_adarms = getattr(
-                    self.model.paligemma_with_expert.gemma_expert.config, "use_adarms", False
-                )
-                if expert_uses_adarms:
-                    logging.warning(f"Skipping norm key (adaRMS mismatch): {key}")
-                    continue
-
-            # Handle MLP naming changes for pi05
-            # pi05 model expects time_mlp_*, but checkpoint might have action_time_mlp_*
-            if key.startswith("action_time_mlp_in."):
-                new_key = key.replace("action_time_mlp_in.", "time_mlp_in.")
-            elif key.startswith("action_time_mlp_out."):
-                new_key = key.replace("action_time_mlp_out.", "time_mlp_out.")
-            # Also handle state_proj which shouldn't exist in pi05
-            if key.startswith("state_proj."):
-                logging.warning(f"Skipping state_proj key in pi05 mode: {key}")
-                continue
-
-            # Handle vision tower embedding layer potential differences
-            if "patch_embedding" in key:
-                # Some checkpoints might have this, but current model expects different structure
-                logging.warning(f"Vision embedding key might need handling: {key}")
-
-            if (
-                key == "model.paligemma_with_expert.paligemma.lm_head.weight"
-                or key == "paligemma_with_expert.paligemma.lm_head.weight"
-            ):
-                fixed_state_dict[
-                    "model.paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"
-                ] = value.clone()
-
-            fixed_state_dict[new_key] = value
-
-        return fixed_state_dict
-
     def get_optim_params(self):
         """Return policy parameters, optionally split into LR-scaled groups.
 
@@ -2491,93 +2262,6 @@ class PI052Policy(PreTrainedPolicy):
         )
         return groups
 
-    def init_rtc_processor(self):
-        """Initialize RTC processor if RTC is enabled in config."""
-        self.rtc_processor = None
-
-        # Create processor if config provided
-        # If RTC is not enabled - we can still track the denoising data
-        if self.config.rtc_config is not None:
-            self.rtc_processor = RTCProcessor(self.config.rtc_config)
-
-            model_value = getattr(self, "model", None)
-            if model_value is not None:
-                model_value.rtc_processor = self.rtc_processor
-
-    def _rtc_enabled(self) -> bool:
-        return self.config.rtc_config is not None and self.config.rtc_config.enabled
-
-    def _preprocess_images(self, batch: dict[str, Tensor]) -> tuple[list[Tensor], list[Tensor]]:
-        """Preprocess images for the model.
-
-        Images from LeRobot are typically in [B, C, H, W] format and normalized to [0, 1].
-        PaliGemma expects images in [B, C, H, W] format and normalized to [-1, 1].
-        """
-        images = []
-        img_masks = []
-
-        # Get device from model parameters
-        device = next(self.parameters()).device
-
-        present_img_keys = [key for key in self.config.image_features if key in batch]
-        missing_img_keys = [key for key in self.config.image_features if key not in batch]
-
-        if len(present_img_keys) == 0:
-            raise ValueError(
-                f"All image features are missing from the batch. At least one expected. "
-                f"(batch: {batch.keys()}) (image_features: {self.config.image_features})"
-            )
-
-        # Preprocess image features present in the batch
-        for key in present_img_keys:
-            img = batch[key]
-
-            # Ensure tensor is on the same device as the model
-            if img.device != device:
-                img = img.to(device)
-
-            # Ensure float32 dtype for consistency
-            if img.dtype != torch.float32:
-                img = img.to(torch.float32)
-
-            # from openpi preprocess_observation_pytorch: Handle both [B, C, H, W] and [B, H, W, C] formats
-            is_channels_first = img.shape[1] == 3  # Check if channels are in dimension 1
-
-            if is_channels_first:
-                # Convert [B, C, H, W] to [B, H, W, C] for processing
-                img = img.permute(0, 2, 3, 1)
-
-            # from openpi preprocess_observation_pytorch: Resize with padding if needed
-            if img.shape[1:3] != self.config.image_resolution:
-                img = resize_with_pad_torch(img, *self.config.image_resolution)
-
-            # Normalize from [0,1] to [-1,1] as expected by siglip
-            img = img * 2.0 - 1.0
-
-            # from openpi preprocess_observation_pytorch: Convert back to [B, C, H, W] format if it was originally channels-first
-            if is_channels_first:
-                img = img.permute(0, 3, 1, 2)  # [B, H, W, C] -> [B, C, H, W]
-
-            images.append(img)
-            # Create mask (all ones for real images)
-            bsize = img.shape[0]
-            mask = torch.ones(bsize, dtype=torch.bool, device=device)
-            img_masks.append(mask)
-
-        # Create image features not present in the batch as fully 0 padded images
-        for _num_empty_cameras in range(len(missing_img_keys)):
-            img = torch.ones_like(img) * -1  # Padded with -1 for SigLIP
-            mask = torch.zeros_like(mask)  # Mask is zero for empty cameras
-            images.append(img)
-            img_masks.append(mask)
-
-        return images, img_masks
-
-    def prepare_action(self, batch):
-        """Pad action"""
-        actions = pad_vector(batch[ACTION], self.config.max_action_dim)
-        return actions
-
     @torch.no_grad()
     def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor:
         """Predict a chunk of actions given environment observations."""
@@ -2640,14 +2324,3 @@ class PI052Policy(PreTrainedPolicy):
             loss = losses.mean()
             loss_dict["loss"] = loss.item()
             return loss, loss_dict
-
-    def _get_default_peft_targets(self) -> dict[str, any]:
-        """Return default PEFT target modules for PI0.5 fine-tuning."""
-        common_projections = (
-            "state_proj|action_in_proj|action_out_proj|action_time_mlp_in|action_time_mlp_out"
-        )
-        target_modules = rf"(.*\.gemma_expert\..*\.self_attn\.(q|v)_proj|model\.({common_projections}))"
-        return {
-            "target_modules": target_modules,
-            "modules_to_save": [],
-        }
diff --git a/src/lerobot/runtime/cli.py b/src/lerobot/runtime/cli.py
index 7ff08079a..a848a2614 100644
--- a/src/lerobot/runtime/cli.py
+++ b/src/lerobot/runtime/cli.py
@@ -23,6 +23,7 @@ import argparse
 import logging
 import sys
 from collections.abc import Callable
+from contextlib import nullcontext
 from typing import Any
 
 from .adapter import GenerationConfig
@@ -389,7 +390,11 @@ def _load_policy_and_preprocessor(
     policy = policy_cls.from_pretrained(policy_path, config=cfg)
     policy.to(cfg.device)
     if load_processors_from_checkpoint:
-        preprocessor, postprocessor = make_pre_post_processors(cfg, pretrained_path=cfg.pretrained_path)
+        preprocessor, postprocessor = make_pre_post_processors(
+            cfg,
+            pretrained_path=cfg.pretrained_path,
+            preprocessor_overrides={"device_processor": {"device": str(cfg.device)}},
+        )
 
     policy.eval()
     return policy, preprocessor, postprocessor
@@ -540,9 +545,11 @@ def _strip_quotes(text: str) -> str:
 
 def _clear_action_queue(runtime: Any) -> None:
     """Drop any queued action chunk so nothing fires while paused."""
-    queue = runtime.state.get("action_queue")
-    if hasattr(queue, "clear"):
-        queue.clear()
+    lock = getattr(runtime.state, "lock", nullcontext())
+    with lock:
+        queue = runtime.state.get("action_queue")
+        if hasattr(queue, "clear"):
+            queue.clear()
 
 
 def _handle_slash_command(runtime: Any, line: str) -> bool:
diff --git a/src/lerobot/runtime/language_runtime.py b/src/lerobot/runtime/language_runtime.py
index c56f78afe..419028c90 100644
--- a/src/lerobot/runtime/language_runtime.py
+++ b/src/lerobot/runtime/language_runtime.py
@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import logging
+import threading
 import time
 from collections import deque
 from collections.abc import Callable
@@ -41,6 +42,8 @@ class RuntimeState:
     actions_dispatched: int = 0
     action_deadline: float | None = None
     extra: dict[str, Any] = field(default_factory=dict)
+    revision: int = 0
+    lock: Any = field(default_factory=threading.RLock, repr=False)
 
     def emit(self, event_name: str) -> None:
         self.events.add(event_name)
@@ -55,16 +58,18 @@ class RuntimeState:
         self.log_lines.append(line)
 
     def set_context(self, key: str, value: str | None, *, label: str | None = None) -> bool:
-        previous = self.language_context.get(key)
-        if previous == value:
-            return False
-        if value is None:
-            self.language_context.pop(key, None)
-        else:
-            self.language_context[key] = value
-        if label is not None and value:
-            self.log(f"  {label}: {value}")
-        return True
+        with self.lock:
+            previous = self.language_context.get(key)
+            if previous == value:
+                return False
+            if value is None:
+                self.language_context.pop(key, None)
+            else:
+                self.language_context[key] = value
+            self.revision += 1
+            if label is not None and value:
+                self.log(f"  {label}: {value}")
+            return True
 
     def get(self, key: str, default: Any = None) -> Any:
         try:
@@ -87,10 +92,13 @@ class RuntimeState:
         raise KeyError(key)
 
     def __setitem__(self, key: str, value: Any) -> None:
-        if hasattr(self, key):
-            setattr(self, key, value)
-        else:
-            self.extra[key] = value
+        with self.lock:
+            if hasattr(self, key):
+                if key == "mode" and self.mode != value:
+                    self.revision += 1
+                setattr(self, key, value)
+            else:
+                self.extra[key] = value
 
 
 class LanguageConditionedPolicyAdapter(Protocol):
@@ -179,8 +187,11 @@ class LanguageConditionedRuntime:
         return getattr(self.policy_adapter, "policy", self.policy_adapter)
 
     def set_task(self, task: str) -> None:
-        self.state.task = task
-        self.state.log(f"Task: {task}")
+        with self.state.lock:
+            if self.state.task != task:
+                self.state.revision += 1
+            self.state.task = task
+            self.state.log(f"Task: {task}")
 
     def stop(self) -> None:
         self._stop = True
@@ -255,12 +266,14 @@ class LanguageConditionedRuntime:
         self.state.extra["recent_interjection"] = None
 
     def maybe_enqueue_action_chunk(self, *, force: bool = False) -> None:
-        if self.state.mode != "action" or not self.state.task:
-            return
-        if self.state.action_queue:
-            return
-        if self.state.tick is None or not self._chunk_gate.due(self.state.tick, force=force):
-            return
+        with self.state.lock:
+            if self.state.mode != "action" or not self.state.task:
+                return
+            if self.state.action_queue:
+                return
+            if self.state.tick is None or not self._chunk_gate.due(self.state.tick, force=force):
+                return
+            revision = self.state.revision
         observation = self._current_observation()
         if observation is None:
             return
@@ -270,7 +283,16 @@ class LanguageConditionedRuntime:
             logger.warning("select_action failed: %s", exc, exc_info=logger.isEnabledFor(logging.DEBUG))
             self.state.log(f"  [warn] select_action failed: {type(exc).__name__}: {exc}")
             return
-        self._enqueue_chunk(chunk)
+        with self.state.lock:
+            if (
+                self.state.revision != revision
+                or self.state.mode != "action"
+                or self.state.stop
+                or self._stop
+            ):
+                logger.info("Discarded an action chunk invalidated during inference.")
+                return
+            self._enqueue_chunk(chunk)
 
     def _enqueue_chunk(self, chunk: Any) -> None:
         if chunk is None:
diff --git a/src/lerobot/runtime/sim_robocasa.py b/src/lerobot/runtime/sim_robocasa.py
index 9b5fdfcae..d1a5380b1 100644
--- a/src/lerobot/runtime/sim_robocasa.py
+++ b/src/lerobot/runtime/sim_robocasa.py
@@ -21,12 +21,14 @@ from __future__ import annotations
 
 import logging
 from collections.abc import Callable
+from datetime import datetime
 from pathlib import Path
 from typing import Any
 
 import numpy as np
 import torch
 
+from lerobot.utils.io_utils import StreamingVideoWriter
 from lerobot.utils.video_annotation import annotate_frame
 
 logger = logging.getLogger(__name__)
@@ -197,7 +199,8 @@ class RoboCasaSimBackend:
         self.record = record
         self.output_dir = Path(output_dir) if output_dir else Path("outputs/runtime_sim")
 
-        self._frames: list[np.ndarray] = []
+        self._video_writer: StreamingVideoWriter | None = None
+        self._video_path: Path | None = None
         self._live_counter = 0
         self._latest_frame: np.ndarray | None = None
         self._stream_server: Any = None
@@ -287,8 +290,7 @@ class RoboCasaSimBackend:
             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()
+            self._capture_frame()
             # AsyncVectorEnv resets terminated sub-environments automatically.
             if bool(np.any(terminated)) or bool(np.any(truncated)):
                 logger.info("[sim] episode ended — scene auto-reset")
@@ -334,9 +336,23 @@ class RoboCasaSimBackend:
             frame,
             (("Task", self._current_task()), ("Subtask", subtask), ("Memory", memory)),
         )
-        self._frames.append(annotated)
         self._latest_frame = annotated  # served by the live MJPEG stream
         self._write_live_frame(annotated)
+        if self.record:
+            self._write_recording_frame(annotated)
+
+    def _write_recording_frame(self, frame: np.ndarray) -> None:
+        try:
+            if self._video_writer is None:
+                self.output_dir.mkdir(parents=True, exist_ok=True)
+                stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+                self._video_path = self.output_dir / f"sim_{stamp}.mp4"
+                fps = int((getattr(self.env, "metadata", None) or {}).get("render_fps", 20))
+                self._video_writer = StreamingVideoWriter(self._video_path, fps)
+            self._video_writer.add_frame(frame)
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("[sim] video encoding failed: %s", exc)
+            self.record = False
 
     def _write_live_frame(self, frame: np.ndarray) -> None:
         """Write a rolling latest.png every few frames for live viewing over SSH.
@@ -345,8 +361,6 @@ class RoboCasaSimBackend:
         the rollout in near-real-time without a GUI window. Written atomically
         (temp + replace) so a reader never sees a half-written file.
         """
-        if not self.record:
-            return
         self._live_counter += 1
         if self._live_counter % 3 != 0:
             return
@@ -363,22 +377,16 @@ class RoboCasaSimBackend:
             logger.debug("[sim] live frame write failed: %s", exc)
 
     def _flush_video(self) -> None:
-        if not self.record or not self._frames:
+        if self._video_writer is None:
             return
-        from datetime import datetime  # noqa: PLC0415
-
-        from lerobot.utils.io_utils import write_video  # noqa: PLC0415
-
-        self.output_dir.mkdir(parents=True, exist_ok=True)
-        stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        path = self.output_dir / f"sim_{stamp}.mp4"
-        fps = int((getattr(self.env, "metadata", None) or {}).get("render_fps", 20))
+        writer = self._video_writer
+        self._video_writer = None
         try:
-            write_video(str(path), np.stack(self._frames), fps)
-            logger.info("[sim] wrote video (%d frames) to %s", len(self._frames), path)
-            print(f"[runtime] sim video saved to {path}", flush=True)
+            writer.close()
+            logger.info("[sim] wrote video (%d frames) to %s", writer.frames_written, self._video_path)
+            print(f"[runtime] sim video saved to {self._video_path}", flush=True)
         except Exception as exc:  # noqa: BLE001
-            logger.warning("[sim] write_video failed: %s", exc)
+            logger.warning("[sim] video close failed: %s", exc)
 
     def attach_stream_server(self, server: Any) -> None:
         """Attach an already-running MJPEG server so disconnect() can stop it."""
diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py
index 428f1ff8e..1d54af604 100644
--- a/src/lerobot/scripts/lerobot_train.py
+++ b/src/lerobot/scripts/lerobot_train.py
@@ -24,6 +24,7 @@ import os
 import sys
 import time
 from contextlib import nullcontext
+from datetime import timedelta
 from pprint import pformat
 from typing import TYPE_CHECKING, Any
 
@@ -171,11 +172,12 @@ def update_policy(
     train_metrics.lr = optimizer.param_groups[0]["lr"]
     if torch.cuda.is_available():
         train_metrics.gpu_mem_gb = torch.cuda.max_memory_allocated() / (1024**3)
-    # Materialize GPU metrics only when logging to avoid synchronizing every step.
+    train_metrics.accumulate_tensor("loss", loss)
+    train_metrics.accumulate_tensor("grad_norm", grad_norm)
+    train_metrics.update_s = time.perf_counter() - start_time
+    # Synchronize accumulated GPU metrics only when logging.
     if log_metrics:
-        train_metrics.loss = loss.item()
-        train_metrics.grad_norm = grad_norm.item()
-        train_metrics.update_s = time.perf_counter() - start_time
+        train_metrics.materialize_tensors()
         # Materialize detached loss components during the same logging synchronization.
         if output_dict:
             output_dict = {
@@ -208,7 +210,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
 
     require_package("accelerate", extra="training")
     from accelerate import Accelerator
-    from accelerate.utils import DistributedDataParallelKwargs, DistributedType
+    from accelerate.utils import DistributedDataParallelKwargs, DistributedType, InitProcessGroupKwargs
 
     cfg.validate()
 
@@ -217,10 +219,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     # We set step_scheduler_with_optimizer=False to prevent accelerate from adjusting the lr_scheduler steps based on the num_processes
     # We set find_unused_parameters=True to handle models with conditional computation
     if accelerator is None:
-        from datetime import timedelta
-
-        from accelerate.utils import InitProcessGroupKwargs
-
         # Static graphs restore DDP overlap when conditional parameter usage is stable.
         # Environment flags retain the existing defaults.
         ddp_find_unused = os.environ.get("LEROBOT_DDP_FIND_UNUSED", "1") == "1"
@@ -336,7 +334,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
 
     active_cfg = cfg.trainable_config
     processor_pretrained_path = active_cfg.pretrained_path
-    # Build PI052 processors from the current config so recipe and FAST labels are generated.
+    # A weight checkpoint may contain PI05 or differently configured PI052 processors.
     if cfg.policy.type == "pi052" and processor_pretrained_path is not None and not cfg.resume:
         logging.warning(
             "pi052 is loading pretrained weights from %s, but building processors from the current "
@@ -344,16 +342,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             processor_pretrained_path,
         )
         processor_pretrained_path = None
-    if (
-        getattr(active_cfg, "use_relative_actions", False)
-        and processor_pretrained_path is not None
-        and not cfg.resume
-    ):
-        logging.warning(
-            "use_relative_actions=true with pretrained processors can skip relative transforms if "
-            "the checkpoint processors do not define them. Building processors from current policy config."
-        )
-        processor_pretrained_path = None
 
     processor_kwargs = {}
     if (processor_pretrained_path and not cfg.resume) or not processor_pretrained_path:
diff --git a/src/lerobot/utils/io_utils.py b/src/lerobot/utils/io_utils.py
index e037b412c..7deeb2d70 100644
--- a/src/lerobot/utils/io_utils.py
+++ b/src/lerobot/utils/io_utils.py
@@ -23,6 +23,46 @@ logger = logging.getLogger(__name__)
 JsonLike = str | int | float | bool | None | list["JsonLike"] | dict[str, "JsonLike"] | tuple["JsonLike", ...]
 
 
+class StreamingVideoWriter:
+    """Incrementally encode RGB frames to an MP4 without retaining them in memory."""
+
+    def __init__(self, video_path: str | Path, fps: int) -> None:
+        from .import_utils import require_package
+
+        require_package("av", extra="av-dep")
+        import av
+
+        self._av = av
+        self._container = av.open(str(video_path), mode="w")
+        self._stream = self._container.add_stream("libx264", rate=fps)
+        self._shape: tuple[int, int] | None = None
+        self.frames_written = 0
+
+    def add_frame(self, frame_array) -> None:
+        orig_height, orig_width = frame_array.shape[:2]
+        height = orig_height - orig_height % 2
+        width = orig_width - orig_width % 2
+        if self._shape is None:
+            self._shape = (height, width)
+            self._stream.width = width
+            self._stream.height = height
+            self._stream.pix_fmt = "yuv420p"
+        elif self._shape != (height, width):
+            raise ValueError(f"Video frame shape changed from {self._shape} to {(height, width)}")
+        frame = self._av.VideoFrame.from_ndarray(frame_array[:height, :width], format="rgb24")
+        for packet in self._stream.encode(frame):
+            self._container.mux(packet)
+        self.frames_written += 1
+
+    def close(self) -> None:
+        if self._container is None:
+            return
+        for packet in self._stream.encode():
+            self._container.mux(packet)
+        self._container.close()
+        self._container = None
+
+
 def load_json(fpath: Path) -> Any:
     """Load data from a JSON file.
 
@@ -58,36 +98,12 @@ def write_video(video_path: str | Path, stacked_frames: list, fps: int) -> None:
         stacked_frames: List of HWC uint8 numpy arrays (RGB).
         fps: Frames per second for the output video.
     """
-    from .import_utils import require_package
-
-    require_package("av", extra="av-dep")
-    import av
-
-    with av.open(str(video_path), mode="w") as container:
-        orig_height, orig_width = stacked_frames[0].shape[:2]
-        # yuv420p requires even dimensions; crop by one pixel if needed
-        height = orig_height if orig_height % 2 == 0 else orig_height - 1
-        width = orig_width if orig_width % 2 == 0 else orig_width - 1
-        if height != orig_height or width != orig_width:
-            logger.warning(
-                "Frame dimensions %dx%d are not even; cropping to %dx%d for yuv420p compatibility.",
-                orig_width,
-                orig_height,
-                width,
-                height,
-            )
-        stream = container.add_stream("libx264", rate=fps)
-        stream.width = width
-        stream.height = height
-        stream.pix_fmt = "yuv420p"
+    writer = StreamingVideoWriter(video_path, fps)
+    try:
         for frame_array in stacked_frames:
-            if height != orig_height or width != orig_width:
-                frame_array = frame_array[:height, :width]
-            frame = av.VideoFrame.from_ndarray(frame_array, format="rgb24")
-            for packet in stream.encode(frame):
-                container.mux(packet)
-        for packet in stream.encode():
-            container.mux(packet)
+            writer.add_frame(frame_array)
+    finally:
+        writer.close()
 
 
 def deserialize_json_into_object[T: JsonLike](fpath: Path, obj: T) -> T:
diff --git a/src/lerobot/utils/logging_utils.py b/src/lerobot/utils/logging_utils.py
index eb02b4ba6..d2f26e3f2 100644
--- a/src/lerobot/utils/logging_utils.py
+++ b/src/lerobot/utils/logging_utils.py
@@ -104,6 +104,8 @@ class MetricsTracker:
         "episodes",
         "epochs",
         "accelerator",
+        "_tensor_sums",
+        "_tensor_counts",
     ]
 
     def __init__(
@@ -129,6 +131,8 @@ class MetricsTracker:
         self.episodes = self.samples / self._avg_samples_per_ep
         self.epochs = self.samples / self._num_frames
         self.accelerator = accelerator
+        self._tensor_sums: dict[str, torch.Tensor] = {}
+        self._tensor_counts: dict[str, int] = {}
 
     def __getattr__(self, name: str) -> int | dict[str, AverageMeter] | AverageMeter | Any:
         if name in self.__dict__:
@@ -156,6 +160,22 @@ class MetricsTracker:
         self.episodes = self.samples / self._avg_samples_per_ep
         self.epochs = self.samples / self._num_frames
 
+    def accumulate_tensor(self, name: str, value: torch.Tensor) -> None:
+        """Accumulate a detached metric on-device until the next logging step."""
+        if name not in self.metrics:
+            raise KeyError(f"Unknown metric {name!r}.")
+        value = value.detach()
+        self._tensor_sums[name] = self._tensor_sums.get(name, torch.zeros_like(value)) + value
+        self._tensor_counts[name] = self._tensor_counts.get(name, 0) + 1
+
+    def materialize_tensors(self) -> None:
+        """Transfer pending tensor averages to their meters with one sync per metric."""
+        for name, total in self._tensor_sums.items():
+            count = self._tensor_counts[name]
+            self.metrics[name].update((total / count).item(), n=count)
+        self._tensor_sums.clear()
+        self._tensor_counts.clear()
+
     def reduce_across_ranks(self) -> None:
         """
         Synchronises the running averages of every metric whose ``reduction`` is not ``"none"``
@@ -227,3 +247,5 @@ class MetricsTracker:
         """Resets average meters."""
         for m in self.metrics.values():
             m.reset()
+        self._tensor_sums.clear()
+        self._tensor_counts.clear()
diff --git a/tests/policies/pi052/test_pi052_bucketed_ce.py b/tests/policies/pi052/test_pi052_bucketed_ce.py
index 8e966bdd7..4cbb93cb5 100644
--- a/tests/policies/pi052/test_pi052_bucketed_ce.py
+++ b/tests/policies/pi052/test_pi052_bucketed_ce.py
@@ -19,7 +19,71 @@ import torch
 
 pytest.importorskip("transformers")
 
-from lerobot.policies.pi052.modeling_pi052 import _lin_ce_flat
+from lerobot.policies.pi052.modeling_pi052 import _lin_ce_flat, _shifted_lin_ce
+
+
+def test_shifted_ce_none_retains_distinct_per_sample_losses():
+    hidden = torch.tensor(
+        [
+            [[8.0, 0.0], [0.0, 8.0], [0.0, 0.0]],
+            [[0.0, 8.0], [8.0, 0.0], [0.0, 0.0]],
+        ]
+    )
+    labels = torch.tensor([[0, 0, 1], [0, 0, 1]])
+    losses = _shifted_lin_ce(hidden, torch.eye(2), labels, reduction="none")
+
+    assert losses.shape == (2,)
+    assert losses[0] < losses[1]
+
+
+def test_checkpoint_resolution_forwards_explicit_hub_options(monkeypatch, tmp_path):
+    import lerobot.policies.pi052.modeling_pi052 as modeling_pi052
+
+    checkpoint = tmp_path / "model.safetensors"
+    checkpoint.touch()
+    calls = []
+
+    def fake_cached_file(model_id, filename, **kwargs):
+        calls.append((model_id, filename, kwargs))
+        return None if filename.endswith("index.json") else str(checkpoint)
+
+    monkeypatch.setattr(modeling_pi052, "cached_file", fake_cached_file)
+    files = modeling_pi052._resolve_weight_files(
+        "org/model",
+        force_download=True,
+        resume_download=True,
+        proxies={"https": "proxy"},
+        token="secret",
+        cache_dir=tmp_path / "cache",
+        local_files_only=True,
+        revision="commit",
+    )
+
+    assert files == [checkpoint]
+    for _model_id, _filename, kwargs in calls:
+        assert kwargs["revision"] == "commit"
+        assert kwargs["cache_dir"] == tmp_path / "cache"
+        assert kwargs["force_download"] is True
+        assert kwargs["resume_download"] is True
+        assert kwargs["proxies"] == {"https": "proxy"}
+        assert kwargs["token"] == "secret"
+        assert kwargs["local_files_only"] is True
+
+
+def test_checkpoint_resolution_rejects_local_directory_without_weights(tmp_path):
+    import lerobot.policies.pi052.modeling_pi052 as modeling_pi052
+
+    with pytest.raises(FileNotFoundError, match="model.safetensors"):
+        modeling_pi052._resolve_weight_files(
+            tmp_path,
+            force_download=False,
+            resume_download=None,
+            proxies=None,
+            token=None,
+            cache_dir=None,
+            local_files_only=False,
+            revision=None,
+        )
 
 
 @pytest.mark.parametrize("z_loss_weight", [0.0, 1e-4])
diff --git a/tests/policies/pi052/test_pi052_import.py b/tests/policies/pi052/test_pi052_import.py
new file mode 100644
index 000000000..b298d39fe
--- /dev/null
+++ b/tests/policies/pi052/test_pi052_import.py
@@ -0,0 +1,27 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import subprocess
+import sys
+
+
+def test_pi052_config_import_does_not_load_model_or_dataset_processor():
+    code = """
+import sys
+from lerobot.policies import PI052Config
+assert PI052Config.__name__ == "PI052Config"
+assert "lerobot.policies.pi052.modeling_pi052" not in sys.modules
+assert "lerobot.policies.pi052.processor_pi052" not in sys.modules
+"""
+    subprocess.run([sys.executable, "-c", code], check=True)
diff --git a/tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py b/tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py
index 7f928a32b..790bbfe4c 100644
--- a/tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py
+++ b/tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py
@@ -16,6 +16,8 @@
 
 from types import SimpleNamespace
 
+import pytest
+
 from lerobot.policies import factory
 from lerobot.policies.pi0_fast.configuration_pi0_fast import PI0FastConfig
 from lerobot.policies.pi052 import fit_fast_tokenizer as fit_module
@@ -48,6 +50,27 @@ def test_pi0_fast_resolves_dataset_specific_tokenizer(monkeypatch, tmp_path):
     }
 
 
+def test_fast_fit_failure_is_not_silently_replaced(monkeypatch, tmp_path):
+    config = PI0FastConfig(auto_fit_fast_tokenizer=True, fast_tokenizer_cache_dir=str(tmp_path))
+    monkeypatch.setattr(
+        fit_module,
+        "fit_fast_tokenizer",
+        lambda **kwargs: (_ for _ in ()).throw(RuntimeError("fit failed")),
+    )
+
+    with pytest.raises(RuntimeError, match="fit failed"):
+        fit_module.resolve_fast_tokenizer(config, "user/dataset")
+
+
+def test_each_node_uses_its_local_rank_zero_as_fit_leader(monkeypatch):
+    monkeypatch.setenv("RANK", "8")
+    monkeypatch.setenv("LOCAL_RANK", "0")
+    assert fit_module._is_local_leader()
+
+    monkeypatch.setenv("LOCAL_RANK", "1")
+    assert not fit_module._is_local_leader()
+
+
 def test_pretrained_pi0_fast_overrides_only_fitted_tokenizer(monkeypatch):
     config = PI0FastConfig(auto_fit_fast_tokenizer=True)
     calls = []
diff --git a/tests/runtime/test_language_runtime.py b/tests/runtime/test_language_runtime.py
index af21b5e20..38ffc77c5 100644
--- a/tests/runtime/test_language_runtime.py
+++ b/tests/runtime/test_language_runtime.py
@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from lerobot.runtime import (
-    LanguageConditionedRuntime,
-)
+import threading
+import time
+
+from lerobot.runtime import LanguageConditionedRuntime, Tick
 
 
 class FakeAdapter:
@@ -69,3 +70,31 @@ def test_runtime_handles_user_interjection():
 
     assert "please say ok" in adapter.interjections
     assert runtime.state.language_context["plan"] == "new plan"
+
+
+def test_prompt_change_discards_in_flight_action_chunk():
+    started = threading.Event()
+    release = threading.Event()
+
+    class BlockingAdapter(FakeAdapter):
+        def select_action(self, observation, state):
+            started.set()
+            assert release.wait(timeout=2)
+            return ["stale"]
+
+    runtime = LanguageConditionedRuntime(
+        policy_adapter=BlockingAdapter(),
+        observation_provider=lambda: {"observation.state": 1},
+    )
+    runtime.set_task("old task")
+    runtime.state.tick = Tick(index=1, monotonic_seconds=time.monotonic())
+    inference = threading.Thread(target=runtime.maybe_enqueue_action_chunk, kwargs={"force": True})
+    inference.start()
+    assert started.wait(timeout=2)
+
+    runtime.set_task("new task")
+    release.set()
+    inference.join(timeout=2)
+
+    assert not inference.is_alive()
+    assert list(runtime.state.action_queue) == []
diff --git a/tests/runtime/test_sim_robocasa.py b/tests/runtime/test_sim_robocasa.py
index 3b788df9e..5f60bfd6e 100644
--- a/tests/runtime/test_sim_robocasa.py
+++ b/tests/runtime/test_sim_robocasa.py
@@ -17,6 +17,7 @@ from types import SimpleNamespace
 
 import numpy as np
 
+from lerobot.runtime.sim_robocasa import RoboCasaSimBackend
 from lerobot.utils.video_annotation import annotate_frame
 
 
@@ -59,3 +60,23 @@ def test_overlay_draws_each_label_once(monkeypatch):
     assert all(color == (255, 255, 255) and thickness == 1 for _, color, thickness in put_text_calls)
     assert len(rectangle_calls) == 1
     assert not np.shares_memory(annotated, frame)
+
+
+def test_capture_updates_live_frame_when_recording_is_disabled(monkeypatch):
+    backend = object.__new__(RoboCasaSimBackend)
+    frame = np.full((8, 8, 3), 42, dtype=np.uint8)
+    written = []
+    backend.record = False
+    backend.runtime_state = None
+    backend._multiview_frame = lambda: frame
+    backend._current_task = lambda: "task"
+    backend._subtask_getter = None
+    backend._memory_getter = None
+    backend._latest_frame = None
+    backend._write_live_frame = written.append
+    monkeypatch.setattr("lerobot.runtime.sim_robocasa.annotate_frame", lambda image, labels: image)
+
+    backend._capture_frame()
+
+    assert backend._latest_frame is frame
+    assert written == [frame]
diff --git a/tests/utils/test_logging_utils.py b/tests/utils/test_logging_utils.py
index aa851bd2a..3f08b9089 100644
--- a/tests/utils/test_logging_utils.py
+++ b/tests/utils/test_logging_utils.py
@@ -37,6 +37,12 @@ class MockAccelerator:
             return self._reduce_fn(tensor, reduction)
         return tensor
 
+    def gather(self, tensor):
+        if self._reduce_fn is None:
+            return tensor.repeat(self.num_processes)
+        reduced = self._reduce_fn(tensor, "max")
+        return torch.cat([tensor.repeat(self.num_processes - 1), reduced])
+
 
 def test_average_meter_initialization():
     meter = AverageMeter("loss", ":.2f")
@@ -168,6 +174,18 @@ def test_metrics_tracker_reset_averages(mock_metrics):
     assert tracker.accuracy.avg == 0.0
 
 
+def test_metrics_tracker_materializes_full_tensor_window(mock_metrics):
+    tracker = MetricsTracker(batch_size=2, num_frames=10, num_episodes=2, metrics=mock_metrics)
+    tracker.accumulate_tensor("loss", torch.tensor(1.0))
+    tracker.accumulate_tensor("loss", torch.tensor(3.0))
+
+    assert tracker.loss.count == 0
+    tracker.materialize_tensors()
+
+    assert tracker.loss.avg == pytest.approx(2.0)
+    assert tracker.loss.count == 2
+
+
 def test_average_meter_invalid_reduction():
     with pytest.raises(ValueError):
         AverageMeter("loss", reduction="median")
diff --git a/tests/utils/test_rerun_visualization.py b/tests/utils/test_rerun_visualization.py
new file mode 100644
index 000000000..d4c3e6999
--- /dev/null
+++ b/tests/utils/test_rerun_visualization.py
@@ -0,0 +1,311 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib
+import sys
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+pytest.importorskip("rerun", reason="rerun-sdk is required (install lerobot[viz])")
+
+from lerobot.types import TransitionKey
+from lerobot.utils.constants import OBS_STATE
+
+
+@pytest.fixture
+def mock_rerun(monkeypatch):
+    """
+    Provide a mock `rerun` module (and `rerun.blueprint` submodule) so tests don't
+    depend on the real library. Also reload the module-under-test so it binds to
+    this mock `rr`.
+    """
+    calls = []
+    blueprints = []
+
+    class DummyScalar:
+        def __init__(self, value):
+            # Scalars may be built from a single float or from a 1D array batch.
+            self.value = value
+
+    class DummyImage:
+        def __init__(self, arr):
+            self.arr = arr
+
+        def compress(self, *a, **k):
+            return self
+
+    class DummyDepthImage:
+        def __init__(self, arr, meter=None, colormap=None):
+            self.arr = arr
+            self.meter = meter
+            self.colormap = colormap
+
+    def dummy_log(key, obj=None, **kwargs):
+        # Accept either positional `obj` or keyword `entity` and record remaining kwargs.
+        if obj is None and "entity" in kwargs:
+            obj = kwargs.pop("entity")
+        calls.append((key, obj, kwargs))
+
+    def dummy_send_blueprint(blueprint, *a, **k):
+        blueprints.append(blueprint)
+
+    # Mock the `rerun.blueprint` submodule used to build the layout.
+    dummy_rrb = SimpleNamespace(
+        Spatial2DView=lambda origin=None, name=None: SimpleNamespace(
+            kind="Spatial2DView", origin=origin, name=name
+        ),
+        TimeSeriesView=lambda name=None, contents=None: SimpleNamespace(
+            kind="TimeSeriesView", name=name, contents=contents
+        ),
+        Grid=lambda *views: SimpleNamespace(kind="Grid", views=list(views)),
+        Blueprint=lambda root: SimpleNamespace(kind="Blueprint", root=root),
+    )
+
+    dummy_rr = SimpleNamespace(
+        __name__="rerun",
+        __package__="rerun",
+        __spec__=SimpleNamespace(name="rerun", submodule_search_locations=None),
+        Scalars=DummyScalar,
+        Image=DummyImage,
+        DepthImage=DummyDepthImage,
+        components=SimpleNamespace(Colormap=SimpleNamespace(Viridis="viridis")),
+        log=dummy_log,
+        send_blueprint=dummy_send_blueprint,
+        init=lambda *a, **k: None,
+        spawn=lambda *a, **k: None,
+        blueprint=dummy_rrb,
+    )
+
+    # Inject fake modules into sys.modules (both `rerun` and `rerun.blueprint`).
+    monkeypatch.setitem(sys.modules, "rerun", dummy_rr)
+    monkeypatch.setitem(sys.modules, "rerun.blueprint", dummy_rrb)
+
+    # Now import and reload the module under test, to bind to our rerun mock
+    import lerobot.utils.rerun_visualization as rv
+
+    importlib.reload(rv)
+
+    # Expose the reloaded module, the call recorder and the captured blueprints
+    yield rv, calls, blueprints
+
+
+def _keys(calls):
+    """Helper to extract just the keys logged to rr.log"""
+    return [k for (k, _obj, _kw) in calls]
+
+
+def _obj_for(calls, key):
+    """Find the first object logged under a given key."""
+    for k, obj, _kw in calls:
+        if k == key:
+            return obj
+    raise KeyError(f"Key {key} not found in calls: {calls}")
+
+
+def _kwargs_for(calls, key):
+    for k, _obj, kw in calls:
+        if k == key:
+            return kw
+    raise KeyError(f"Key {key} not found in calls: {calls}")
+
+
+def _views_by_kind(blueprint, kind):
+    """Return the views of a given kind from the (single) blueprint's grid."""
+    return [v for v in blueprint.root.views if v.kind == kind]
+
+
+def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
+    rv, calls, blueprints = mock_rerun
+
+    # Build EnvTransition dict
+    obs = {
+        f"{OBS_STATE}.temperature": np.float32(25.0),
+        # CHW image should be converted to HWC for rr.Image
+        "observation.camera": np.zeros((3, 10, 20), dtype=np.uint8),
+    }
+    act = {
+        "action.throttle": 0.7,
+        # 1D array should be logged as a single Scalars batch under one entity path
+        "action.vector": np.array([1.0, 2.0], dtype=np.float32),
+    }
+    transition = {
+        TransitionKey.OBSERVATION: obs,
+        TransitionKey.ACTION: act,
+    }
+
+    # Extract observation and action data from transition like in the real call sites
+    obs_data = transition.get(TransitionKey.OBSERVATION, {})
+    action_data = transition.get(TransitionKey.ACTION, {})
+    rv.log_rerun_data(observation=obs_data, action=action_data)
+
+    # We expect:
+    # - observation.state.temperature -> Scalars
+    # - observation.camera -> Image (HWC) with static=True
+    # - action.throttle -> Scalars
+    # - action.vector -> single Scalars batch (no per-element suffix)
+    expected_keys = {
+        f"{OBS_STATE}.temperature",
+        "observation.camera",
+        "action.throttle",
+        "action.vector",
+    }
+    assert set(_keys(calls)) == expected_keys
+
+    # Check scalar types and values
+    temp_obj = _obj_for(calls, f"{OBS_STATE}.temperature")
+    assert type(temp_obj).__name__ == "DummyScalar"
+    assert float(temp_obj.value) == pytest.approx(25.0)
+
+    throttle_obj = _obj_for(calls, "action.throttle")
+    assert type(throttle_obj).__name__ == "DummyScalar"
+    assert float(throttle_obj.value) == pytest.approx(0.7)
+
+    # 1D vector logged as a single batched Scalars under one entity path
+    vec = _obj_for(calls, "action.vector")
+    assert type(vec).__name__ == "DummyScalar"
+    np.testing.assert_allclose(np.asarray(vec.value), [1.0, 2.0])
+
+    # Check image handling: CHW -> HWC
+    img_obj = _obj_for(calls, "observation.camera")
+    assert type(img_obj).__name__ == "DummyImage"
+    assert img_obj.arr.shape == (10, 20, 3)  # transposed
+    assert _kwargs_for(calls, "observation.camera").get("static", False) is True  # static=True for images
+
+    # A blueprint should have been built and sent exactly once, and cached on the function.
+    assert len(blueprints) == 1
+    assert rv.log_rerun_data.blueprint is blueprints[0]
+
+    bp = blueprints[0]
+    # One spatial view per image path
+    spatial_views = _views_by_kind(bp, "Spatial2DView")
+    assert {v.origin for v in spatial_views} == {"observation.camera"}
+
+    # One time-series view each for observation and action scalars
+    ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
+    assert set(ts_views) == {"observation", "action"}
+    assert ts_views["observation"].contents == [f"{OBS_STATE}.temperature"]
+    assert ts_views["action"].contents == ["action.throttle", "action.vector"]
+
+
+def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
+    rv, calls, blueprints = mock_rerun
+
+    # First dict without prefixes treated as observation
+    # Second dict without prefixes treated as action
+    obs_plain = {
+        "temp": 1.5,
+        # Already HWC image => should stay as-is
+        "img": np.zeros((5, 6, 3), dtype=np.uint8),
+        "none": None,  # should be skipped
+    }
+    act_plain = {
+        "throttle": 0.3,
+        "vec": np.array([9, 8, 7], dtype=np.float32),
+    }
+
+    # Extract observation and action data from list like the old function logic did
+    # First dict was treated as observation, second as action
+    rv.log_rerun_data(observation=obs_plain, action=act_plain)
+
+    # Expected keys with auto-prefixes. The 1D vector is a single batched Scalars.
+    expected = {
+        "observation.temp",
+        "observation.img",
+        "action.throttle",
+        "action.vec",
+    }
+    logged = set(_keys(calls))
+    assert logged == expected
+
+    # Scalars
+    t = _obj_for(calls, "observation.temp")
+    assert type(t).__name__ == "DummyScalar"
+    assert float(t.value) == pytest.approx(1.5)
+
+    throttle = _obj_for(calls, "action.throttle")
+    assert type(throttle).__name__ == "DummyScalar"
+    assert float(throttle.value) == pytest.approx(0.3)
+
+    # Image stays HWC
+    img = _obj_for(calls, "observation.img")
+    assert type(img).__name__ == "DummyImage"
+    assert img.arr.shape == (5, 6, 3)
+    assert _kwargs_for(calls, "observation.img").get("static", False) is True
+
+    # Vector logged as a single batched Scalars under one entity path
+    vec = _obj_for(calls, "action.vec")
+    assert type(vec).__name__ == "DummyScalar"
+    np.testing.assert_allclose(np.asarray(vec.value), [9, 8, 7])
+
+    # Blueprint sent once with the expected view layout
+    assert len(blueprints) == 1
+    bp = blueprints[0]
+    spatial_views = _views_by_kind(bp, "Spatial2DView")
+    assert {v.origin for v in spatial_views} == {"observation.img"}
+    ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
+    assert ts_views["observation"].contents == ["observation.temp"]
+    assert ts_views["action"].contents == ["action.throttle", "action.vec"]
+
+
+def test_log_rerun_data_kwargs_only(mock_rerun):
+    rv, calls, blueprints = mock_rerun
+
+    rv.log_rerun_data(
+        observation={"observation.temp": 10.0, "observation.gray": np.zeros((8, 8, 1), dtype=np.uint8)},
+        action={"action.a": 1.0},
+    )
+
+    keys = set(_keys(calls))
+    assert "observation.temp" in keys
+    assert "observation.gray" in keys
+    assert "action.a" in keys
+
+    temp = _obj_for(calls, "observation.temp")
+    assert type(temp).__name__ == "DummyScalar"
+    assert float(temp.value) == pytest.approx(10.0)
+
+    img = _obj_for(calls, "observation.gray")
+    assert type(img).__name__ == "DummyDepthImage"  # single-channel -> DepthImage
+    assert img.arr.shape == (8, 8, 1)  # remains HWC
+    assert _kwargs_for(calls, "observation.gray").get("static", False) is True
+
+    a = _obj_for(calls, "action.a")
+    assert type(a).__name__ == "DummyScalar"
+    assert float(a.value) == pytest.approx(1.0)
+
+    # Blueprint sent once, with a spatial view for the image and time-series views for scalars
+    assert len(blueprints) == 1
+    bp = blueprints[0]
+    assert {v.origin for v in _views_by_kind(bp, "Spatial2DView")} == {"observation.gray"}
+    ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
+    assert ts_views["observation"].contents == ["observation.temp"]
+    assert ts_views["action"].contents == ["action.a"]
+
+
+def test_log_rerun_data_blueprint_sent_only_once(mock_rerun):
+    """The blueprint is built from the first call and not resent on subsequent calls."""
+    rv, calls, blueprints = mock_rerun
+
+    rv.log_rerun_data(observation={"temp": 1.0}, action={"a": 2.0})
+    assert len(blueprints) == 1
+    first_blueprint = rv.log_rerun_data.blueprint
+
+    rv.log_rerun_data(observation={"temp": 3.0}, action={"a": 4.0})
+    # Still only one blueprint, and the cached one is unchanged.
+    assert len(blueprints) == 1
+    assert rv.log_rerun_data.blueprint is first_blueprint