Making the flow-matching action head a scorable stochastic policy

5.1 Why the deterministic sampler is a problem

A policy-gradient method needs $\log \pi(a \mid s)$ and its gradient with respect to the parameters. The Euler ODE sampler in infer_action maps an initial noise $x_0$ to a chunk $A$ through a deterministic, invertible-in-principle but practically intractable map. The density of $A$ exists but computing it would require the Jacobian determinant of 20 stacked transformer evaluations, which is out of the question.

The standard solution, introduced for diffusion policies by DPPO and adapted to flow matching by ReinFlow, πRL (Flow-SDE), and Flow-GRPO, is to stop treating the denoising process as a black box. Instead, each denoising step is treated as a transition in a small inner MDP whose transition distribution is Gaussian, and the log-probability of the final action is the sum of the per-step Gaussian log-densities. To make the transitions Gaussian you have to inject noise at each step, which turns the ODE into an SDE. The crucial property is that you can do this while keeping the marginal distribution at each noise level the same as the ODE's, so the model's training objective is not violated.

5.2 The Flow-SDE step in Fast-WAM's variables

Fast-WAM's forward process is $x_\sigma = (1-\sigma )x_0 + \sigma \varepsilon$, and the network predicts $v = \varepsilon - x_0$. From these two relations the score function (the gradient of the log marginal density at noise level σ) is

$$\nabla_x \log p_\sigma(x) = -\big( x + (1 - \sigma) \cdot v_\varphi(x, \sigma) \big) / \sigma$$

(Appendix A derives this and checks it against the form πRL states.) The probability-flow ODE $dx = v d\sigma$ can be converted into a family of SDEs

$$dx = \big[\, v - (g(\sigma)^2/2) \cdot \nabla \log p_\sigma(x) \,\big]\, d\sigma + g(\sigma)\, d\bar{W}$$

that all share the same marginals $p_\sigma$, for any non-negative noise function $g(\sigma)$ (this is the reverse-SDE construction of Song et al., 2021). Both πRL and Flow-GRPO make the specific choice

$$g(\sigma) = a \cdot \sqrt{\sigma / (1 - \sigma)}$$

with a single scalar $a$ controlling the overall noise level. Substituting the score, the drift coefficient $g^2/(2\sigma)$ simplifies to $a^2/(2(1-\sigma))$. Discretising with Fast-WAM's Euler scheme and step $\Delta _k = \sigma _{k+1} - \sigma _k$ (negative, so time runs from noise to data) gives the update the rollout worker will actually execute:

$$\begin{aligned} v_k &= v_\varphi(x_k, \sigma_k) \\ g_k &= a \cdot \sqrt{\sigma_k / (1 - \sigma_k)} \\ \mathrm{drift}_k &= v_k + \big( g_k^2 / (2 \sigma_k) \big) \cdot \big( x_k + (1 - \sigma_k) \cdot v_k \big) \\ &= v_k + \big( a^2 / (2 (1 - \sigma_k)) \big) \cdot \big( x_k + (1 - \sigma_k) \cdot v_k \big) \\ \mu_k &= x_k + \mathrm{drift}_k \cdot \Delta_k \\ \mathrm{std}_k &= g_k \cdot \sqrt{|\Delta_k|} \\ x_{k+1} &= \mu_k + \mathrm{std}_k \cdot \xi_k, \qquad \xi_k \sim \mathcal{N}(0, I) \end{aligned}$$

This is, up to notation, exactly the update printed in Flow-GRPO and πRL ($x_{t+\Delta t} = x_t + [\,v + \sigma_t^2/(2t)(x_t + (1-t)v)\,]\Delta t + \sigma_t \sqrt{\Delta t} \cdot \varepsilon$, with their $t$ equal to Fast-WAM's $\sigma$ and their $\sigma_t$ equal to $g$).

Two things to notice. With $a = 0$ the update collapses exactly to the existing deterministic step $x_{k+1} = x_k + v_k \Delta _k$, which is a useful unit test (Chapter 13). And $g(\sigma)$ diverges as $\sigma \to 1$, so the very first step (where $\sigma _0 = 1$) must never be a stochastic one; the mixed sampler in the next section guarantees this, and a defensive clamp of $\,1 - \sigma \ge 10^{-3}$ costs nothing. A constant $g(\sigma ) = a$ is mathematically just as valid and avoids the divergence, but it is not what the cited papers use, so treat it as an untested variant.

The log-probability of the whole chunk is then

$$\log \pi_\varphi(A \mid s) = \sum_{k \,\in\, \text{stochastic steps}} \log \mathcal{N}\big( x_{k+1} \,;\, \mu_k(x_k; \varphi),\ \mathrm{std}_k^2 I \big)$$

summed over all 32×7 dimensions, with the chain $x_0, x_1, \dots , x_K = A$ stored from the rollout. Note that $\mu_k$ depends on the parameters $\varphi$ through $v_\varphi$, while std_k does not, so the gradient flows through the mean only.

5.3 Three practical choices, following πRL and Flow-GRPO

Fewer denoising steps during RL. Each stochastic step adds a term to the log-probability, and each term requires a forward (and, in the learner, backward) pass through the action expert. πRL discretises into $K$ uniform steps and ablates $K \in \{1,\ 2,\ 4,\ 8\}$; $K = 1$ suffers from ODE-to-SDE discretisation error, larger $K$ costs more, and $K = 4$ is the setting its main results use. Flow-GRPO collects training rollouts with 10 steps and evaluates with the original 40 ("denoising reduction"). Fast-WAM's default is 20 steps with shift 1.0, and the README reports that shifts from 1.0 to 3.0 perform similarly, which suggests the sampler is not very sensitive to the schedule. Use $K = 8$ for RL rollouts and training (a compromise between πRL's 4 and Flow-GRPO's 10 that keeps Fast-WAM's uniform σ grid at multiples of 0.125), and verify before you start that the pretrained checkpoint's deterministic success rate at $K = 8$ is within a point of its rate at $K = 20$. Always evaluate at $K = 20$.

Mixed ODE–SDE. πRL's mixed sampler does not make every step stochastic. For each chunk it randomly samples one denoising index $k*$, runs the SDE step there, and runs every other step as the deterministic ODE. The policy's log-probability then has a single Gaussian term, which keeps the learner cheap and the variance low, and because $k*$ is never the first index the $\sigma = 1$ divergence of $g(\sigma)$ never arises. This is the grounded default. A variant worth testing, but not from the literature, is to make the last K_sde (say 4) steps stochastic, which concentrates exploration where the action is taking shape and gives the learner more log-probability terms per sample at proportionally higher cost.

Noise scale $a$. This is the exploration knob. πRL ablates $a \in \{0.2,\ 0.5,\ 0.8\}$ and finds $a = 0.5$ best: $a = 0.2$ is unstable (a much higher PPO clip fraction), and $a = 0.8$ degrades the rollout policy too much. Flow-GRPO uses $a = 0.7$ for image generation, which is a different domain. Since Fast-WAM's actions are min/max normalised to about $[-1,\ 1]$, the same scale as π₀'s, start at $a = 0.5$. If the policy's success rate collapses in the first iteration, $a$ is too high; if advantages are all zero because every rollout does the same thing, $a$ is too low.

5.4 What to store from each rollout

For each macro-step the worker must save enough to recompute the log-probability exactly in the learner:

• the state (image, proprio, task id; the context is looked up by task id);
• the full denoising chain $x_0 \dots x_K$ ([K+1, 32, 7]) together with the index (or indices) of the stochastic step(s);
• the old log-probability per stochastic step $\log \pi_{\mathrm{old},k}$ (one value with πRL's single-step sampler, K_sde values with the variant), computed by the worker at sampling time;
• the reward inputs (Chapter 6) and the done flag.

The chain is small (9 × 32 × 7 floats). The images are the bulk of the storage.

5.5 The alternative for a cheap first baseline: DSRL

If you want to validate the environment plumbing and reward computation before touching the model's gradients, DSRL (Wagenmaker et al., 2025) offers a way to do RL on a frozen flow policy. It keeps the sampler exactly as it is and instead trains a small actor-critic (SAC) whose "action" is the initial noise $x_0$; the frozen Fast-WAM turns that noise into a chunk. Because nothing inside Fast-WAM is trained, the compiled inference path can be used unchanged. It will not get you WAM-RL's world-model benefits, but it is a few hundred lines and a strong sanity check.