lerobot/tests/policies/common/test_vla_utils.py

#!/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
import torch.nn.functional as F  # noqa: N812

from lerobot.policies.common.vla_utils import (
    create_sinusoidal_pos_embedding,
    fuse_action_time_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)


def _make_action_time_layers(action_dim=7, width=16, seed=0):
    torch.manual_seed(seed)
    return (
        torch.nn.Linear(action_dim, width),
        torch.nn.Linear(2 * width, width),
        torch.nn.Linear(width, width),
    )


def test_fuse_action_time_embedding_matches_pi0_inline_loop():
    # Verbatim pi0 block: time_emb cast to timestep.dtype, no gradient checkpointing.
    in_proj, mlp_in, mlp_out = _make_action_time_layers()
    min_period, max_period = 4e-3, 4.0
    noisy_actions = torch.randn(3, 5, 7)
    timestep = torch.rand(3)

    time_emb = create_sinusoidal_pos_embedding(
        timestep, in_proj.out_features, min_period=min_period, max_period=max_period, device=timestep.device
    )
    time_emb = time_emb.type(dtype=timestep.dtype)
    action_emb = in_proj(noisy_actions)
    time_emb = time_emb[:, None, :].expand_as(action_emb)
    action_time_emb = torch.cat([action_emb, time_emb], dim=2)
    ref = mlp_out(F.silu(mlp_in(action_time_emb)))

    out = fuse_action_time_embedding(
        noisy_actions,
        timestep,
        action_in_proj=in_proj,
        action_time_mlp_in=mlp_in,
        action_time_mlp_out=mlp_out,
        min_period=min_period,
        max_period=max_period,
        time_emb_dtype=timestep.dtype,
    )
    assert torch.equal(out, ref)


def test_fuse_action_time_embedding_matches_smolvla_inline_loop():
    # Verbatim smolvla block: time_emb cast to action_emb.dtype (default), no checkpointing.
    in_proj, mlp_in, mlp_out = _make_action_time_layers(seed=1)
    min_period, max_period = 4e-3, 4.0
    noisy_actions = torch.randn(2, 4, 7)
    timestep = torch.rand(2)

    action_emb = in_proj(noisy_actions)
    dtype = action_emb.dtype
    time_emb = create_sinusoidal_pos_embedding(
        timestep, in_proj.out_features, min_period, max_period, device=action_emb.device
    )
    time_emb = time_emb.type(dtype=dtype)
    time_emb = time_emb[:, None, :].expand_as(action_emb)
    action_time_emb = torch.cat([action_emb, time_emb], dim=2)
    action_time_emb = mlp_in(action_time_emb)
    action_time_emb = F.silu(action_time_emb)
    ref = mlp_out(action_time_emb)

    out = fuse_action_time_embedding(
        noisy_actions,
        timestep,
        action_in_proj=in_proj,
        action_time_mlp_in=mlp_in,
        action_time_mlp_out=mlp_out,
        min_period=min_period,
        max_period=max_period,
    )
    assert torch.equal(out, ref)


def test_fuse_action_time_embedding_apply_hook_is_invoked():
    # The apply hook (gradient-checkpoint wrapper in pi0/eo1) must wrap both sub-steps.
    in_proj, mlp_in, mlp_out = _make_action_time_layers(seed=2)
    calls = []

    def apply(fn, arg):
        calls.append(fn)
        return fn(arg)

    out = fuse_action_time_embedding(
        torch.randn(2, 4, 7),
        torch.rand(2),
        action_in_proj=in_proj,
        action_time_mlp_in=mlp_in,
        action_time_mlp_out=mlp_out,
        min_period=4e-3,
        max_period=4.0,
        apply=apply,
    )
    assert len(calls) == 2
    assert out.shape == (2, 4, 16)