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

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)