admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-22 09:21:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix pi052 runtime and training safety

Browse Source
This commit is contained in:
Pepijn
2026-07-15 18:17:23 +02:00
parent 3cec067795
commit 0fe31bfae1
16 changed files with 861 additions and 643 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/policies/pi052/__init__.py
+2 -13
View File
@@ -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"]
src/lerobot/policies/pi052/fit_fast_tokenizer.py
+13 -13
View File
@@ -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,
)
src/lerobot/policies/pi052/modeling_pi052.py
+191 -518
View File
@@ -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": [],
}
src/lerobot/runtime/cli.py
+11 -4
View File
@@ -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:
src/lerobot/runtime/language_runtime.py
+45 -23
View File
@@ -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:
src/lerobot/runtime/sim_robocasa.py
+27 -19
View File
@@ -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."""
src/lerobot/scripts/lerobot_train.py
+8 -20
View File
@@ -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:
src/lerobot/utils/io_utils.py
+45 -29
View File
@@ -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:
src/lerobot/utils/logging_utils.py
+22
View File
@@ -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()
tests/policies/pi052/test_pi052_bucketed_ce.py
+65 -1
View File
@@ -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])
tests/policies/pi052/test_pi052_import.py
+27
View File
@@ -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)
tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py
+23
View File
@@ -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 = []
tests/runtime/test_language_runtime.py
+32 -3
View File
@@ -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) == []
tests/runtime/test_sim_robocasa.py
+21
View File
@@ -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]
tests/utils/test_logging_utils.py
+18
View File
@@ -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")
tests/utils/test_rerun_visualization.py
+311
View File
@@ -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