diff --git a/src/lerobot/policies/common/flow_matching.py b/src/lerobot/policies/common/flow_matching.py new file mode 100644 index 000000000..f66b0aacf --- /dev/null +++ b/src/lerobot/policies/common/flow_matching.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python + +# 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. + +"""Flow-matching sampling primitives shared across policies. + +Canonical versions of the beta-distributed timestep sampler and the forward-Euler +denoising loop (with its real-time-chunking hook) that the openpi-derived policies +(pi0, pi05, smolvla, eo1) historically each carried a copy of. All functions are +stateless; adopting them does not affect checkpoints. +""" + +from collections.abc import Callable +from typing import TYPE_CHECKING + +import torch +from torch import Tensor + +if TYPE_CHECKING: + from lerobot.policies.rtc.modeling_rtc import RTCProcessor + + +def sample_beta(alpha: float, beta: float, bsize: int, device) -> Tensor: # see openpi (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 sample_noise(shape, device) -> Tensor: + """Standard-normal float32 noise, the flow-matching x_1 sample.""" + return torch.normal( + mean=0.0, + std=1.0, + size=shape, + dtype=torch.float32, + device=device, + ) + + +def sample_time_beta(bsize: int, device, *, alpha: float, beta: float, scale: float, offset: float) -> Tensor: + """Beta-distributed flow-matching timesteps: ``Beta(alpha, beta) * scale + offset`` (openpi convention).""" + time_beta = sample_beta(alpha, beta, bsize, device) + time = time_beta * scale + offset + return time.to(dtype=torch.float32, device=device) + + +def euler_integrate( + denoise_fn: Callable[[Tensor, Tensor], Tensor], + noise: Tensor, + num_steps: int, + *, + rtc_processor: "RTCProcessor | None" = None, + rtc_enabled: bool = False, + inference_delay: int | None = None, + prev_chunk_left_over: Tensor | None = None, + execution_horizon: int | None = None, +) -> Tensor: + """Forward-Euler integration of a velocity field from t=1 (noise) to t=0 (actions). + + This is the openpi sampling loop: ``dt = -1/num_steps``, ``time = 1.0 + step*dt``, + ``x_t <- x_t + dt * v_t``, with the optional real-time-chunking (RTC) guidance hook + wrapping the velocity computation and debug tracking after each step. + + Args: + denoise_fn: Computes the velocity ``v_t`` from ``(x_t, time_tensor)`` where + ``time_tensor`` is a float32 tensor of shape ``(batch_size,)``. The returned + velocity must have the same shape and dtype as ``x_t``. + noise: Initial sample ``x_1`` of shape ``(batch_size, ...)``. + num_steps: Number of Euler steps. + rtc_processor: Optional RTC processor. Debug tracking fires whenever it is set and + has debugging enabled, even if RTC guidance itself is disabled (this mirrors + the historical per-policy loops). + rtc_enabled: Whether to route the velocity computation through + ``rtc_processor.denoise_step`` (requires ``rtc_processor``). + inference_delay: RTC guidance parameter, forwarded verbatim. + prev_chunk_left_over: RTC guidance parameter, forwarded verbatim. + execution_horizon: RTC guidance parameter, forwarded verbatim. + """ + bsize = noise.shape[0] + device = noise.device + + dt = -1.0 / num_steps + x_t = noise + for step in range(num_steps): + time = 1.0 + step * dt + time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize) + + def denoise_step_partial_call(input_x_t, current_timestep=time_tensor): + return denoise_fn(input_x_t, current_timestep) + + if rtc_enabled: + v_t = rtc_processor.denoise_step( + x_t=x_t, + prev_chunk_left_over=prev_chunk_left_over, + inference_delay=inference_delay, + time=time, + original_denoise_step_partial=denoise_step_partial_call, + execution_horizon=execution_horizon, + ) + else: + v_t = denoise_step_partial_call(x_t) + + x_t = x_t + dt * v_t + + if rtc_processor is not None and rtc_processor.is_debug_enabled(): + rtc_processor.track(time=time, x_t=x_t, v_t=v_t) + + return x_t diff --git a/src/lerobot/policies/common/vla_utils.py b/src/lerobot/policies/common/vla_utils.py new file mode 100644 index 000000000..b728a49f4 --- /dev/null +++ b/src/lerobot/policies/common/vla_utils.py @@ -0,0 +1,243 @@ +#!/usr/bin/env python + +# 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. + +"""Helpers shared by the openpi-derived VLA policies (pi0, pi05, pi0_fast, smolvla, eo1, xvla). + +These are the canonical versions of functions that historically were copy-pasted per +policy. They are pure (no parameters, no module state), so importing them from here +instead of a policy-local copy has no effect on checkpoints. +""" + +import math +from typing import TYPE_CHECKING + +import torch +import torch.nn.functional as F # noqa: N812 +from torch import Tensor + +from lerobot.utils.constants import OPENPI_ATTENTION_MASK_VALUE +from lerobot.utils.device_utils import get_safe_dtype +from lerobot.utils.import_utils import _transformers_available, require_package + +if TYPE_CHECKING or _transformers_available: + from transformers import DynamicCache +else: + DynamicCache = None + + +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 make_att_2d_masks(pad_masks: Tensor, att_masks: Tensor) -> Tensor: # see openpi (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 prepare_attention_masks_4d(att_2d_masks: Tensor, dtype: torch.dtype | None = None) -> Tensor: + """Expand boolean 2D attention masks to the additive 4D layout expected by transformers. + + Valid positions become 0.0 and masked positions the large negative openpi constant. + """ + att_2d_masks_4d = att_2d_masks[:, None, :, :] + result = torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE) + if dtype is not None: + result = result.to(dtype=dtype) + return result + + +def clone_past_key_values(past_key_values): + """Clone the DynamicCache returned by prefix prefill for compiled denoising.""" + if DynamicCache is None: + require_package("transformers", extra="transformers-dep") + + return DynamicCache( + tuple( + (keys.clone(), values.clone(), sliding_window) for keys, values, sliding_window in past_key_values + ) + ) + + +def pad_vector(vector: Tensor, new_dim: int, *, truncate: bool = False) -> Tensor: + """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) + + With ``truncate=False`` (openpi behavior), vectors whose last dimension is already + >= new_dim are returned unchanged. With ``truncate=True`` (xVLA behavior), the last + dimension is truncated to exactly ``new_dim`` (which may be 0). + """ + if vector.shape[-1] == new_dim: + return vector + if not truncate: + if vector.shape[-1] >= new_dim: + return vector + return F.pad(vector, (0, new_dim - vector.shape[-1])) + shape = list(vector.shape) + current_dim = shape[-1] + shape[-1] = new_dim + new_vector = vector.new_zeros(*shape) + length = min(current_dim, new_dim) + new_vector[..., :length] = vector[..., :length] + return new_vector + + +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]. + + Padding is centered (openpi convention). For the top-left-padding variant used by + smolvla/xvla, see :func:`resize_with_pad`. + + 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 + 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 + + +def resize_with_pad(img: torch.Tensor, height: int, width: int, *, pad_value: float) -> torch.Tensor: + """Resize a (b, c, h, w) image without distortion, padding on the LEFT and TOP. + + This is the smolvla/xvla convention. For the centered-padding openpi variant, see + :func:`resize_with_pad_torch`. ``pad_value`` is keyword-only on purpose: callers + historically used different values (0, -1) and must state their choice explicitly. + """ + if img.ndim != 4: + raise ValueError(f"(b,c,h,w) expected, but got {img.shape}") + + current_height, current_width = img.shape[2:] + if current_height == height and current_width == width: + return img + + ratio = max(current_width / width, current_height / height) + resized_height = int(current_height / ratio) + resized_width = int(current_width / ratio) + resized_img = F.interpolate( + img, size=(resized_height, resized_width), mode="bilinear", align_corners=False + ) + + pad_height = max(0, height - resized_height) + pad_width = max(0, width - resized_width) + padded_img = F.pad(resized_img, (pad_width, 0, pad_height, 0), value=pad_value) + return padded_img diff --git a/tests/policies/common/test_flow_matching.py b/tests/policies/common/test_flow_matching.py new file mode 100644 index 000000000..065883ebc --- /dev/null +++ b/tests/policies/common/test_flow_matching.py @@ -0,0 +1,193 @@ +#!/usr/bin/env python + +# 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. + +"""Behavior-pinning tests for the shared flow-matching sampling primitives. + +``euler_integrate`` is compared against a verbatim copy of the historical pi0/pi05/ +smolvla sampling loop (including its RTC hook semantics): any divergence from that +reference is a behavior change for released checkpoints. +""" + +import torch + +from lerobot.policies.common.flow_matching import ( + euler_integrate, + sample_beta, + sample_noise, + sample_time_beta, +) + + +def test_sample_beta_range_dtype_and_reproducibility(): + torch.manual_seed(0) + s1 = sample_beta(1.5, 1.0, 4096, "cpu") + torch.manual_seed(0) + s2 = sample_beta(1.5, 1.0, 4096, "cpu") + assert torch.equal(s1, s2) + assert s1.shape == (4096,) and s1.dtype == torch.float32 + assert s1.min() >= 0.0 and s1.max() <= 1.0 + # Beta(1.5, 1.0) mean is 1.5/2.5 = 0.6. + assert abs(s1.mean().item() - 0.6) < 0.02 + + +def test_sample_time_beta_openpi_convention(): + torch.manual_seed(1) + time = sample_time_beta(4096, "cpu", alpha=1.5, beta=1.0, scale=0.999, offset=0.001) + assert time.dtype == torch.float32 + assert time.min() >= 0.001 and time.max() <= 1.0 + # Exact composition: Beta sample * scale + offset, same RNG stream. + torch.manual_seed(1) + expected = sample_beta(1.5, 1.0, 4096, "cpu") * 0.999 + 0.001 + torch.testing.assert_close(time, expected, rtol=0, atol=0) + + +def test_sample_noise_seeded(): + torch.manual_seed(2) + n1 = sample_noise((2, 8, 4), "cpu") + torch.manual_seed(2) + n2 = sample_noise((2, 8, 4), "cpu") + assert torch.equal(n1, n2) + assert n1.dtype == torch.float32 and n1.shape == (2, 8, 4) + + +def test_euler_integrate_constant_velocity_is_exact(): + # With v_t == c constant, x_0 = x_1 + sum(dt * c) = x_1 - c exactly (num_steps * dt = -1). + noise = torch.randn(3, 5, 2) + c = torch.randn(3, 5, 2) + out = euler_integrate(lambda x_t, time: c, noise, num_steps=10) + torch.testing.assert_close(out, noise - c, rtol=0, atol=1e-6) + + +def _reference_pi0_loop(denoise_fn, noise, num_steps, rtc_enabled, rtc_processor, kw): + """Verbatim structure of the historical pi0/pi05/smolvla sample_actions loop.""" + bsize = noise.shape[0] + device = noise.device + dt = -1.0 / num_steps + x_t = noise + for step in range(num_steps): + time = 1.0 + step * dt + time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize) + + def denoise_step_partial_call(input_x_t, current_timestep=time_tensor): + return denoise_fn(input_x_t, current_timestep) + + if rtc_enabled: + v_t = rtc_processor.denoise_step( + x_t=x_t, + prev_chunk_left_over=kw.get("prev_chunk_left_over"), + inference_delay=kw.get("inference_delay"), + time=time, + original_denoise_step_partial=denoise_step_partial_call, + execution_horizon=kw.get("execution_horizon"), + ) + else: + v_t = denoise_step_partial_call(x_t) + x_t = x_t + dt * v_t + if rtc_processor is not None and rtc_processor.is_debug_enabled(): + rtc_processor.track(time=time, x_t=x_t, v_t=v_t) + return x_t + + +class _StubRTCProcessor: + def __init__(self, debug_enabled: bool): + self._debug = debug_enabled + self.tracked = [] + self.guidance_calls = [] + + def is_debug_enabled(self): + return self._debug + + def denoise_step( + self, + x_t, + prev_chunk_left_over, + inference_delay, + time, + original_denoise_step_partial, + execution_horizon, + ): + self.guidance_calls.append( + { + "time": time, + "inference_delay": inference_delay, + "execution_horizon": execution_horizon, + "x_t": x_t.clone(), + } + ) + return original_denoise_step_partial(x_t) * 0.5 + + def track(self, time, x_t, v_t): + self.tracked.append({"time": time, "x_t": x_t.clone(), "v_t": v_t.clone()}) + + +def _make_denoise_fn(): + weight = torch.randn(4, 4) * 0.1 + + def denoise_fn(x_t, time_tensor): + return x_t @ weight + time_tensor[:, None, None] + + return denoise_fn + + +def test_euler_integrate_matches_historical_loop(): + torch.manual_seed(3) + denoise_fn = _make_denoise_fn() + noise = torch.randn(2, 6, 4) + ref = _reference_pi0_loop(denoise_fn, noise, 10, rtc_enabled=False, rtc_processor=None, kw={}) + out = euler_integrate(denoise_fn, noise, 10) + assert torch.equal(out, ref) + + +def test_euler_integrate_rtc_guidance_and_kwarg_forwarding(): + torch.manual_seed(4) + denoise_fn = _make_denoise_fn() + noise = torch.randn(2, 6, 4) + leftover = torch.randn(2, 6, 4) + kw = {"inference_delay": 3, "prev_chunk_left_over": leftover, "execution_horizon": 25} + + ref_proc, new_proc = _StubRTCProcessor(False), _StubRTCProcessor(False) + ref = _reference_pi0_loop(denoise_fn, noise, 6, rtc_enabled=True, rtc_processor=ref_proc, kw=kw) + out = euler_integrate( + denoise_fn, + noise, + 6, + rtc_processor=new_proc, + rtc_enabled=True, + inference_delay=3, + prev_chunk_left_over=leftover, + execution_horizon=25, + ) + assert torch.equal(out, ref) + assert len(new_proc.guidance_calls) == 6 + for ref_call, new_call in zip(ref_proc.guidance_calls, new_proc.guidance_calls, strict=True): + assert ref_call["time"] == new_call["time"] + assert new_call["inference_delay"] == 3 and new_call["execution_horizon"] == 25 + # Guidance sees the PRE-update x_t. + assert torch.equal(ref_call["x_t"], new_call["x_t"]) + + +def test_euler_integrate_debug_tracking_fires_even_when_rtc_disabled(): + # Historical behavior: track() fires whenever the processor exists and has debugging + # enabled, independent of whether RTC guidance is active. + torch.manual_seed(5) + denoise_fn = _make_denoise_fn() + noise = torch.randn(2, 6, 4) + proc = _StubRTCProcessor(True) + out = euler_integrate(denoise_fn, noise, 4, rtc_processor=proc, rtc_enabled=False) + assert len(proc.guidance_calls) == 0 + assert len(proc.tracked) == 4 + # track() receives the POST-update x_t; the last one is the returned sample. + assert torch.equal(proc.tracked[-1]["x_t"], out) diff --git a/tests/policies/common/test_vla_utils.py b/tests/policies/common/test_vla_utils.py new file mode 100644 index 000000000..4d57b15ee --- /dev/null +++ b/tests/policies/common/test_vla_utils.py @@ -0,0 +1,195 @@ +#!/usr/bin/env python + +# 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. + +"""Behavior-pinning tests for the shared VLA helpers. + +These helpers are the canonical versions of functions that used to be copy-pasted across +the openpi-derived policies (pi0, pi05, pi0_fast, smolvla, eo1, xvla). The expected +values below encode the historical per-policy behavior exactly; a failure here means a +behavior change that would silently affect released checkpoints. +""" + +import math + +import pytest +import torch + +from lerobot.policies.common.vla_utils import ( + create_sinusoidal_pos_embedding, + make_att_2d_masks, + pad_vector, + prepare_attention_masks_4d, + resize_with_pad, + resize_with_pad_torch, +) +from lerobot.utils.constants import OPENPI_ATTENTION_MASK_VALUE + + +def test_create_sinusoidal_pos_embedding_matches_openpi_formula(): + time = torch.tensor([0.0, 0.25, 1.0]) + dim, min_period, max_period = 8, 4e-3, 4.0 + emb = create_sinusoidal_pos_embedding(time, dim, min_period, max_period, device=torch.device("cpu")) + + assert emb.shape == (3, dim) + # Independent recomputation of the openpi formula in float64. + fraction = torch.linspace(0.0, 1.0, dim // 2, dtype=torch.float64) + period = min_period * (max_period / min_period) ** fraction + scaling = 1.0 / period * 2 * math.pi + sin_input = scaling[None, :] * time.to(torch.float64)[:, None] + expected = torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1) + torch.testing.assert_close(emb, expected, rtol=1e-9, atol=1e-9) + + +def test_create_sinusoidal_pos_embedding_validation(): + with pytest.raises(ValueError, match="divisible by 2"): + create_sinusoidal_pos_embedding(torch.zeros(2), 7, 4e-3, 4.0, device=torch.device("cpu")) + with pytest.raises(ValueError, match="batch_size"): + create_sinusoidal_pos_embedding(torch.zeros(2, 2), 8, 4e-3, 4.0, device=torch.device("cpu")) + + +def test_make_att_2d_masks_docstring_cases(): + # Pure causal attention: [[1 1 1]] + pad = torch.ones(1, 3, dtype=torch.bool) + att = torch.tensor([[1, 1, 1]], dtype=torch.int32) + expected = torch.tensor([[[1, 0, 0], [1, 1, 0], [1, 1, 1]]], dtype=torch.bool) + assert torch.equal(make_att_2d_masks(pad, att), expected) + + # Prefix-LM: [[0 0 1 1]] -> first two tokens attend bidirectionally, rest causal. + att = torch.tensor([[0, 0, 1, 1]], dtype=torch.int32) + pad = torch.ones(1, 4, dtype=torch.bool) + expected = torch.tensor([[[1, 1, 0, 0], [1, 1, 0, 0], [1, 1, 1, 0], [1, 1, 1, 1]]], dtype=torch.bool) + assert torch.equal(make_att_2d_masks(pad, att), expected) + + # Padding removes rows and columns. + pad = torch.tensor([[True, True, False]]) + att = torch.tensor([[0, 1, 1]], dtype=torch.int32) + out = make_att_2d_masks(pad, att) + assert not out[0, :, 2].any() and not out[0, 2, :].any() + + +def test_make_att_2d_masks_validation(): + with pytest.raises(ValueError): + make_att_2d_masks(torch.ones(3, dtype=torch.bool), torch.ones(1, 3, dtype=torch.int32)) + with pytest.raises(ValueError): + make_att_2d_masks(torch.ones(1, 3, dtype=torch.bool), torch.ones(3, dtype=torch.int32)) + + +def test_prepare_attention_masks_4d(): + masks = torch.tensor([[[True, False], [False, True]]]) + out = prepare_attention_masks_4d(masks) + assert out.shape == (1, 1, 2, 2) + expected = torch.tensor([[[[0.0, OPENPI_ATTENTION_MASK_VALUE], [OPENPI_ATTENTION_MASK_VALUE, 0.0]]]]) + assert torch.equal(out, expected) + + out_bf16 = prepare_attention_masks_4d(masks, dtype=torch.bfloat16) + assert out_bf16.dtype == torch.bfloat16 + assert torch.equal(out_bf16, expected.to(torch.bfloat16)) + + +def test_pad_vector_openpi_semantics(): + v = torch.arange(6.0).reshape(2, 3) + padded = pad_vector(v, 5) + assert padded.shape == (2, 5) + assert torch.equal(padded[:, :3], v) and not padded[:, 3:].any() + # Already large enough (>=): returned unchanged, same object. + assert pad_vector(v, 3) is v + assert pad_vector(v, 2) is v + # 3D input. + v3 = torch.ones(2, 4, 3) + assert pad_vector(v3, 7).shape == (2, 4, 7) + + +def test_pad_vector_truncate_semantics(): + v = torch.arange(6.0).reshape(2, 3) + out = pad_vector(v, 2, truncate=True) + assert out.shape == (2, 2) and torch.equal(out, v[:, :2]) + out = pad_vector(v, 5, truncate=True) + assert out.shape == (2, 5) and torch.equal(out[:, :3], v) and not out[:, 3:].any() + assert pad_vector(v, 0, truncate=True).shape == (2, 0) + assert pad_vector(v, 3, truncate=True) is v + + +@pytest.mark.parametrize("channels_last", [True, False]) +def test_resize_with_pad_torch_centered(channels_last): + img = torch.rand(2, 3, 30, 60) if not channels_last else torch.rand(2, 30, 60, 3) + out = resize_with_pad_torch(img, 64, 64) + if channels_last: + assert out.shape == (2, 64, 64, 3) + # Aspect ratio preserved: 30x60 -> 32x64, padded 16 top and 16 bottom (centered). + assert not out[:, :16].any() and not out[:, -16:].any() + assert out[:, 16:48].abs().sum() > 0 + else: + assert out.shape == (2, 3, 64, 64) + assert not out[:, :, :16].any() and not out[:, :, -16:].any() + + +def test_resize_with_pad_torch_uint8_roundtrip(): + img = (torch.rand(1, 3, 20, 20) * 255).to(torch.uint8) + out = resize_with_pad_torch(img, 40, 40) + assert out.dtype == torch.uint8 and out.shape == (1, 3, 40, 40) + with pytest.raises(ValueError, match="Unsupported image dtype"): + resize_with_pad_torch(torch.rand(1, 3, 8, 8, dtype=torch.float64), 16, 16) + + +def test_resize_with_pad_top_left(): + img = torch.rand(2, 3, 30, 60) + out = resize_with_pad(img, 64, 64, pad_value=-1.0) + assert out.shape == (2, 3, 64, 64) + # 30x60 -> 32x64; this variant pads on the TOP only (32 rows of pad_value). + assert torch.equal(out[:, :, :32], torch.full((2, 3, 32, 64), -1.0)) + assert out[:, :, 32:].min() >= 0 + # No-op fast path returns the same object. + assert resize_with_pad(img, 30, 60, pad_value=0.0) is img + with pytest.raises(ValueError, match="expected"): + resize_with_pad(torch.rand(3, 8, 8), 16, 16, pad_value=0.0) + + +def test_clone_past_key_values(): + pytest.importorskip("transformers") + from transformers import DynamicCache + + from lerobot.policies.common.vla_utils import clone_past_key_values + + cache = DynamicCache() + keys, values = torch.rand(1, 2, 4, 8), torch.rand(1, 2, 4, 8) + cache.update(keys, values, 0) + cloned = clone_past_key_values(cache) + (ck, cv, _), (ok, ov, _) = next(iter(cloned)), next(iter(cache)) + assert torch.equal(ck, ok) and torch.equal(cv, ov) + # Deep copy: mutating the clone must not touch the original. + ck.zero_() + assert not torch.equal(ck, ok) + + +def test_clone_past_key_values_is_fullgraph_compilable(): + pytest.importorskip("transformers") + from transformers import DynamicCache + + from lerobot.policies.common.vla_utils import clone_past_key_values + + cache = DynamicCache() + keys, values = torch.rand(1, 2, 4, 8), torch.rand(1, 2, 4, 8) + cache.update(keys, values, 0) + + compiled_clone = torch.compile(clone_past_key_values, backend="eager", fullgraph=True) + cloned = compiled_clone(cache) + + (cloned_keys, cloned_values, _), (original_keys, original_values, _) = ( + next(iter(cloned)), + next(iter(cache)), + ) + assert torch.equal(cloned_keys, original_keys) + assert torch.equal(cloned_values, original_values)