feat(flow-matching): cover per-policy sampling divergences

Extend the shared samplers to represent the non-openpi conventions Martino
catalogued, without changing existing openpi behavior:

- sample_noise: add distribution="uniform" for evo1's rand*2-1 noise.
- sample_time_beta: add complement flag (groot/wall_x forward t=(1-beta)*s)
  and optional clamp_min/clamp_max (evo1 Beta(2,2) clamped to [0.02, 0.98]);
  scale/offset now default to the identity mapping.
- sample_beta: build the Beta distribution via a cached, CPU-side helper
  (groot convention) so it is constructed once per (alpha, beta).

Adds tests for uniform noise, complement/clamp timesteps, and caching.

Co-authored-by: Cursor <cursoragent@cursor.com>

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)