diff --git a/src/lerobot/policies/common/flow_matching.py b/src/lerobot/policies/common/flow_matching.py index f66b0aacf..7e3b0a4ee 100644 --- a/src/lerobot/policies/common/flow_matching.py +++ b/src/lerobot/policies/common/flow_matching.py @@ -23,6 +23,7 @@ stateless; adopting them does not affect checkpoints. """ from collections.abc import Callable +from functools import lru_cache from typing import TYPE_CHECKING import torch @@ -32,29 +33,67 @@ 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 +@lru_cache(maxsize=None) +def _beta_distribution(alpha: float, beta: float) -> "torch.distributions.Beta": + # Beta sampling uses _sample_dirichlet which isn't implemented for MPS, so build on CPU. + # Cached (groot convention) so the distribution object is constructed once per (alpha, beta). 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) + return torch.distributions.Beta(alpha_t, beta_t) -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_beta(alpha: float, beta: float, bsize: int, device) -> Tensor: # see openpi (exact copy) + return _beta_distribution(alpha, beta).sample((bsize,)).to(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).""" +def sample_noise(shape, device, *, distribution: str = "normal") -> Tensor: + """Float32 flow-matching noise sample. + + ``distribution="normal"`` (default, openpi: pi0/pi05/eo1/smolvla/groot/wall_x) draws + standard-normal noise. ``distribution="uniform"`` (evo1) draws uniformly from + ``[-1, 1)`` via ``rand * 2 - 1``. + """ + if distribution == "normal": + return torch.normal( + mean=0.0, + std=1.0, + size=shape, + dtype=torch.float32, + device=device, + ) + if distribution == "uniform": + return torch.rand(shape, dtype=torch.float32, device=device) * 2 - 1 + raise ValueError(f"Unknown noise distribution: {distribution!r} (expected 'normal' or 'uniform')") + + +def sample_time_beta( + bsize: int, + device, + *, + alpha: float, + beta: float, + scale: float = 1.0, + offset: float = 0.0, + complement: bool = False, + clamp_min: float | None = None, + clamp_max: float | None = None, +) -> Tensor: + """Beta-distributed flow-matching timesteps. + + Computes ``t = f(Beta(alpha, beta)) * scale + offset`` where ``f`` is the identity by + default or ``1 - x`` when ``complement=True``, then optionally clamps to + ``[clamp_min, clamp_max]``. This covers the known per-policy conventions: + + * openpi backward (pi0/pi05/eo1/smolvla): ``scale=0.999, offset=0.001``. + * forward (groot/wall_x): ``complement=True, scale=0.999`` giving ``(1 - beta) * 0.999``. + * evo1: ``alpha=beta=2, clamp_min=0.02, clamp_max=0.98``. + """ time_beta = sample_beta(alpha, beta, bsize, device) + if complement: + time_beta = 1.0 - time_beta time = time_beta * scale + offset + if clamp_min is not None or clamp_max is not None: + time = time.clamp(min=clamp_min, max=clamp_max) return time.to(dtype=torch.float32, device=device) diff --git a/tests/policies/common/test_flow_matching.py b/tests/policies/common/test_flow_matching.py index 065883ebc..5edf640a9 100644 --- a/tests/policies/common/test_flow_matching.py +++ b/tests/policies/common/test_flow_matching.py @@ -24,6 +24,7 @@ reference is a behavior change for released checkpoints. import torch from lerobot.policies.common.flow_matching import ( + _beta_distribution, euler_integrate, sample_beta, sample_noise, @@ -63,6 +64,57 @@ def test_sample_noise_seeded(): assert n1.dtype == torch.float32 and n1.shape == (2, 8, 4) +def test_sample_noise_normal_is_default(): + torch.manual_seed(2) + default = sample_noise((2, 8, 4), "cpu") + torch.manual_seed(2) + normal = sample_noise((2, 8, 4), "cpu", distribution="normal") + assert torch.equal(default, normal) + + +def test_sample_noise_uniform_evo1(): + torch.manual_seed(2) + n = sample_noise((4096,), "cpu", distribution="uniform") + assert n.dtype == torch.float32 + assert n.min() >= -1.0 and n.max() < 1.0 + # evo1's rand_like * 2 - 1 has mean ~0 over [-1, 1). + assert abs(n.mean().item()) < 0.05 + # Exact match to the historical evo1 expression on the same RNG stream. + torch.manual_seed(2) + expected = torch.rand((4096,), dtype=torch.float32) * 2 - 1 + torch.testing.assert_close(n, expected, rtol=0, atol=0) + + +def test_sample_noise_invalid_distribution(): + import pytest + + with pytest.raises(ValueError, match="Unknown noise distribution"): + sample_noise((2, 2), "cpu", distribution="bogus") + + +def test_sample_time_beta_forward_complement_convention(): + # groot/wall_x forward convention: t = (1 - beta) * 0.999. + torch.manual_seed(9) + time = sample_time_beta(4096, "cpu", alpha=1.5, beta=1.0, scale=0.999, complement=True) + torch.manual_seed(9) + expected = (1.0 - sample_beta(1.5, 1.0, 4096, "cpu")) * 0.999 + torch.testing.assert_close(time, expected, rtol=0, atol=0) + + +def test_sample_time_beta_evo1_clamp(): + # evo1: Beta(2, 2) clamped to [0.02, 0.98]. + torch.manual_seed(10) + time = sample_time_beta(4096, "cpu", alpha=2.0, beta=2.0, clamp_min=0.02, clamp_max=0.98) + assert time.min() >= 0.02 and time.max() <= 0.98 + + +def test_sample_beta_distribution_is_cached(): + a = _beta_distribution(1.5, 1.0) + b = _beta_distribution(1.5, 1.0) + assert a is b + assert _beta_distribution(2.0, 2.0) is not a + + 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)