fix pi052 runtime and training safety

2026-07-24 10:16:09 +00:00 · 2026-07-15 18:17:23 +02:00
parent 3cec067795
commit 0fe31bfae1
16 changed files with 861 additions and 643 deletions
@@ -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 configuration; model and processors are imported lazily by their factories."""
 PI052 adds supervised PaliGemma text generation, prompt dropout, and autoregressive inference to PI0.5.
 """
 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__ = [
+__all__ = ["PI052Config"]
    "PI052Config",
    "PI052Policy",
    "PI052TextTokenizerStep",
    "make_pi052_pre_post_processors",
 ]
@@ -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.
+    # Each node fits its node-local cache once; its other local ranks wait.
-    is_leader = int(os.environ.get("RANK", "0")) == 0 and int(os.environ.get("LOCAL_RANK", "0")) == 0
+    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(
-        return fit_fast_tokenizer(
+        dataset_repo_id=dataset_repo_id,
-            dataset_repo_id=dataset_repo_id,
+        cache_dir=Path(config.fast_tokenizer_cache_dir).expanduser(),
-            cache_dir=Path(config.fast_tokenizer_cache_dir).expanduser(),
+        base_tokenizer_name=config.action_tokenizer_name,
-            base_tokenizer_name=config.action_tokenizer_name,
+        n_samples=config.fast_tokenizer_fit_samples,
-            n_samples=config.fast_tokenizer_fit_samples,
+        chunk_size=config.chunk_size,
-            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
@@ -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 safetensors.torch import load_file
-from torch.nn import functional as F  # noqa: N812
+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 ..pi05.modeling_pi05 import (
-from ..pi_gemma import PaliGemmaWithExpertModel, get_gemma_config
+    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__)
-
+_SAFETENSORS_FILE = "model.safetensors"
-# Generic dual-expert transformer helpers live in ``lerobot.policies.pi_gemma``.
+_SAFETENSORS_INDEX = "model.safetensors.index.json"
-class ActionSelectKwargs(TypedDict, total=False):
+def _resolve_weight_files(
-    inference_delay: int | None
+    pretrained_name_or_path: str | Path,
-    prev_chunk_left_over: Tensor | None
+    *,
-    execution_horizon: int | None
+    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,
    }
-
+    if local_dir.is_dir():
-def get_safe_dtype(target_dtype, device_type):
+        index_path = local_dir / _SAFETENSORS_INDEX
-    """Get a safe dtype for the given device type."""
+        single_path = local_dir / _SAFETENSORS_FILE
    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]
    else:
-        channels_last = False
+        resolved_index = cached_file(
-        if images.dim() == 3:
+            model_id,
-            images = images.unsqueeze(0)  # Add batch dimension
+            _SAFETENSORS_INDEX,
-
+            _raise_exceptions_for_missing_entries=False,
-    batch_size, channels, cur_height, cur_width = images.shape
+            **load_kwargs,
    # 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,
        )
        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)
+    if index_path is None or not index_path.is_file():
-        self.action_out_proj = nn.Linear(action_expert_config.width, config.max_action_dim)
+        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)
+    index = json.loads(index_path.read_text())
-        self.time_mlp_out = nn.Linear(action_expert_config.width, action_expert_config.width)
+    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
+def _load_weight_files(files: list[Path]) -> dict[str, Tensor]:
-        if config.compile_model:
+    state_dict: dict[str, Tensor] = {}
-            torch.set_float32_matmul_precision("high")
+    for path in files:
-            self.sample_actions = torch.compile(self.sample_actions, mode=config.compile_mode)
+        shard = load_file(path)
-            # Also compile the main forward pass used during training
+        overlap = state_dict.keys() & shard.keys()
-            self.forward = torch.compile(self.forward, mode=config.compile_mode)
+        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:
+def _reduce_action_loss(per_sample: Tensor, predict_actions_t: Tensor | None, reduction: str) -> Tensor:
-    """Mean over samples where ``predict_actions_t`` is True, else over all."""
+    """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))
+                    hidden_chunk = functional.pad(hidden_chunk, (0, 0, 0, pad_rows))
-                    labels_chunk = F.pad(labels_chunk, (0, pad_rows), value=-100)
+                    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)
+    It inherits unchanged PI0.5 policy behavior and replaces the core model with
-    is vendored in ``pi05_backbone.py`` and the PI05Policy wrapper logic is
+    the joint flow/text implementation below.
    inlined directly here, so this policy does not depend on or inherit from
    ``lerobot.policies.pi05`` (which stays identical to ``main``).
    """
    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")
+        loss_dict["loss"] = total.detach().mean()
        if reduction == "none":
            return total.expand(batch[OBS_LANGUAGE_TOKENS].shape[0]), loss_dict
        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."""
+        """Load a PI05/PI052 checkpoint, including sharded safetensors checkpoints."""
        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"
        )
        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)
+        lm_head_key = "model.paligemma_with_expert.paligemma.lm_head.weight"
-        try:
+        embed_tokens_key = "model.paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"
-            print(f"Loading model from: {pretrained_name_or_path}")
+        if lm_head_key not in remapped_state_dict and embed_tokens_key in remapped_state_dict:
-            try:
+            remapped_state_dict[lm_head_key] = remapped_state_dict[embed_tokens_key].clone().float()
-                from transformers.utils import cached_file
+        elif lm_head_key in remapped_state_dict:
-
+            remapped_state_dict[lm_head_key] = remapped_state_dict[lm_head_key].float()
                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)
        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")
+                logger.warning("Missing PI052 checkpoint keys: %s", missing_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")
            if unexpected_keys:
-                print(f"Unexpected keys when loading state dict: {len(unexpected_keys)} keys")
+                logger.warning("Unexpected PI052 checkpoint keys: %s", unexpected_keys)
-                if len(unexpected_keys) <= 5:
+        model.to(config.device)
-                    for key in unexpected_keys:
+        model.eval()
                        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}")
        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": [],
        }
@@ -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")
+    lock = getattr(runtime.state, "lock", nullcontext())
-    if hasattr(queue, "clear"):
+    with lock:
-        queue.clear()
+        queue = runtime.state.get("action_queue")
        if hasattr(queue, "clear"):
            queue.clear()
 def _handle_slash_command(runtime: Any, line: str) -> bool:
@@ -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)
+        with self.lock:
-        if previous == value:
+            previous = self.language_context.get(key)
-            return False
+            if previous == value:
-        if value is None:
+                return False
-            self.language_context.pop(key, None)
+            if value is None:
-        else:
+                self.language_context.pop(key, None)
-            self.language_context[key] = value
+            else:
-        if label is not None and value:
+                self.language_context[key] = value
-            self.log(f"  {label}: {value}")
+            self.revision += 1
-        return True
+            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):
+        with self.lock:
-            setattr(self, key, value)
+            if hasattr(self, key):
-        else:
+                if key == "mode" and self.mode != value:
-            self.extra[key] = 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
+        with self.state.lock:
-        self.state.log(f"Task: {task}")
+            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:
+        with self.state.lock:
-            return
+            if self.state.mode != "action" or not self.state.task:
-        if self.state.action_queue:
+                return
-            return
+            if self.state.action_queue:
-        if self.state.tick is None or not self._chunk_gate.due(self.state.tick, force=force):
+                return
-            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:
@@ -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
+        writer = self._video_writer
-
+        self._video_writer = None
        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))
        try:
-            write_video(str(path), np.stack(self._frames), fps)
+            writer.close()
-            logger.info("[sim] wrote video (%d frames) to %s", len(self._frames), path)
+            logger.info("[sim] wrote video (%d frames) to %s", writer.frames_written, self._video_path)
-            print(f"[runtime] sim video saved to {path}", flush=True)
+            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."""
@@ -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.materialize_tensors()
        train_metrics.grad_norm = grad_norm.item()
        train_metrics.update_s = time.perf_counter() - start_time
        # 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:
@@ -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
+    writer = StreamingVideoWriter(video_path, fps)
-
+    try:
    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"
        for frame_array in stacked_frames:
-            if height != orig_height or width != orig_width:
+            writer.add_frame(frame_array)
-                frame_array = frame_array[:height, :width]
+    finally:
-            frame = av.VideoFrame.from_ndarray(frame_array, format="rgb24")
+        writer.close()
            for packet in stream.encode(frame):
                container.mux(packet)
        for packet in stream.encode():
            container.mux(packet)
 def deserialize_json_into_object[T: JsonLike](fpath: Path, obj: T) -> T:
@@ -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()
@@ -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])
@@ -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)
@@ -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 = []
@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from lerobot.runtime import (
+import threading
-    LanguageConditionedRuntime,
+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) == []
@@ -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]
@@ -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")
@@ -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