refactor(vla): extract shared model components (#4054)

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

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

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)

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)