feat(flow-matching): support forward-Euler integration convention

Add a `forward_euler` flag to `euler_integrate` so the shared loop can serve
the groot/evo1/wall_x policies (integrate t: 0 -> 1, dt = +1/num_steps) in
addition to the openpi backward convention (t: 1 -> 0), which remains the
default. RTC hook and debug tracking are unchanged in both directions.

Adds forward-convention equivalence tests against a reference loop.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Martino Russi
2026-07-19 12:22:01 +02:00
parent a9879e69ed
commit ddcb61f7dd
2 changed files with 49 additions and 7 deletions
+19 -7
View File
@@ -63,24 +63,35 @@ def euler_integrate(
noise: Tensor,
num_steps: int,
*,
forward_euler: bool = False,
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).
"""Euler integration of a velocity field between the noise and action endpoints.
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.
Two integration conventions are supported via ``forward_euler``:
* Backward (default, openpi: pi0, pi05, eo1, smolvla): integrates from t=1 (noise) to
t=0 (actions) with ``dt = -1/num_steps`` and ``time = 1.0 + step*dt``.
* Forward (groot, evo1, wall_x): integrates from t=0 (noise) to t=1 (actions) with
``dt = +1/num_steps`` and ``time = step*dt``.
In both cases the update is ``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, ...)``.
noise: Initial sample of shape ``(batch_size, ...)``. This is ``x_1`` for the
backward convention and ``x_0`` for the forward convention.
num_steps: Number of Euler steps.
forward_euler: If ``True`` use the forward convention (start at t=0); otherwise
use the backward openpi convention (start at t=1).
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).
@@ -93,10 +104,11 @@ def euler_integrate(
bsize = noise.shape[0]
device = noise.device
dt = -1.0 / num_steps
dt = 1.0 / num_steps if forward_euler else -1.0 / num_steps
t_start = 0.0 if forward_euler else 1.0
x_t = noise
for step in range(num_steps):
time = 1.0 + step * dt
time = t_start + 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):
@@ -71,6 +71,36 @@ def test_euler_integrate_constant_velocity_is_exact():
torch.testing.assert_close(out, noise - c, rtol=0, atol=1e-6)
def test_euler_integrate_forward_constant_velocity_is_exact():
# Forward convention: dt = +1/num_steps, so x_1 = x_0 + sum(dt * c) = x_0 + c exactly.
noise = torch.randn(3, 5, 2)
c = torch.randn(3, 5, 2)
out = euler_integrate(lambda x_t, time: c, noise, num_steps=10, forward_euler=True)
torch.testing.assert_close(out, noise + c, rtol=0, atol=1e-6)
def _reference_forward_loop(denoise_fn, noise, num_steps):
"""Verbatim structure of the groot/evo1/wall_x forward-Euler loop (t: 0 -> 1)."""
bsize = noise.shape[0]
device = noise.device
dt = 1.0 / num_steps
x_t = noise
for step in range(num_steps):
time = 0.0 + step * dt
time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
x_t = x_t + dt * denoise_fn(x_t, time_tensor)
return x_t
def test_euler_integrate_forward_matches_reference_loop():
torch.manual_seed(7)
denoise_fn = _make_denoise_fn()
noise = torch.randn(2, 6, 4)
ref = _reference_forward_loop(denoise_fn, noise, 10)
out = euler_integrate(denoise_fn, noise, 10, forward_euler=True)
assert torch.equal(out, ref)
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]