Value Forcing: Aligning Return Distributions at All Noise Levels

TL;DR

Value Forcing learns return distributions with flow matching. It enforces Bellman consistency at every noise level to learn more accurate return distributions. Across 26 offline RL tasks, Value Forcing improves policy performance, with particularly large gains when data is limited.

Learning return distributions from limited data

Value learning lies at the heart of reinforcement learning (RL), yet accurately estimating the value can be hard: each transition provides only an immediate reward, but no ground-truth returns. Thus, errors from temporal-difference (TD) learning are hard to correct when data is limited or rewards are sparse.

How can we get richer learning signals from each transition?

Distributional RL models the full return distribution and can provide richer learning signals when the whole return distribution is used as the TD target. Recent work has begun to use Flow Matching to represent the return distribution. These methods map Gaussian noise to the return distribution through an iterative denoising process, creating an opportunity to enforce Bellman consistency at every noise level and extract more learning signals from the same data. Prior flow-based methods, however, supervise only some noise levels or use incorrect targets, and can learn inaccurate return distributions that collapse toward a single value (see the rightmost figure above).

Prior methods do not enforce correct Bellman consistency

A flow-based critic turns Gaussian noise into a return sample. Ideally, the Bellman equation should hold at each of the intermediate noise levels. Prior methods differ in where along the flow path they enforce this consistency, and each misses or misapplies it at some noise levels, resulting in sparse or incorrect supervision.

Endpoint-only
supervision

Floq, FlowCritic, FlowIQN

Build Bellman targets from clean returns ($\tau = 1$). They train the flow at other noise levels, but do not enforce a Bellman equation there.

Incorrect supervision
at the initial noise

Value Flows

Uses the clean-return Bellman equation at every noise level. At $\tau = 0$, this conflicts with the initial Gaussian noise.

Noise-sampled
supervision

PCBF

Pairs flow paths with the same initial noise, rather than enforcing each intermediate marginal distribution to satisfy a Bellman equation.

In our experiments, FlowCritic, Value Flows, and PCBF can predict inaccurate distributions that concentrate near a single value:

Return flow paths on humanoidmaze-giant. VF spreads Gaussian noise over a broad return distribution, while FlowCritic, Value Flows and PCBF predict narrower distributions. Gray shows observed returns.

See Fig. 7 in the paper for complete results.

Our method: Value Forcing

Value Forcing uses the return distribution at the next state to bootstrap the return distribution at the current state using the same noise level. We first define a new flow Bellman equation, then turn it into flow-matching losses and two offline RL algorithms: VF and IVF.

A flow Bellman equation at every noise level

The standard distributional Bellman equation relates the return random variables at the current state $ Z^\pi(s,a) $ and at the next state $ Z^\pi(S', A') $:

Distributional Bellman equation $$Z^\pi(s,a) \overset{d}{=} \mathcal{T}^\pi Z^\pi(s,a) \overset{d}{=} r(s,a) + \gamma Z^\pi(S', A')$$

Our flow Bellman equation extends this relation to every noise level $\tau \in [0,1]$. It scales the reward by $\tau$ and adds an independent Gaussian correction $N \sim \mathcal{N}(0,1)$:

Flow Bellman equation $$\begin{aligned} Z^\pi(s,a,\tau) &\overset{d}{=} \mathcal{T}^\pi_\tau Z^\pi(s,a,\tau) \overset{d}{=} \tau\, r(s,a) + \gamma Z^\pi(S', A', \tau) + \underbrace{(1-\tau)\sqrt{1-\gamma^2}\, N}_{\text{Gaussian correction}} \end{aligned}$$

The Gaussian correction accounts for the variance in the discounted return at the next time step. We can understand this by examining the equation at the two ends of the flow path.

At $\tau = 0$ initial noise

$$\underbrace{Z^\pi(s,a,0)}_{\text{standard Gaussian}} \overset{d}{=} \gamma \underbrace{Z^\pi(S',A',0)}_{\text{standard Gaussian}} + \sqrt{1-\gamma^2}\, N$$

Discounting reduces the initial noise variance to $\gamma^2$. The independent correction adds $1-\gamma^2$, keeping both sides standard Gaussian.

At $\tau = 1$ clean returns

$$Z^\pi(s,a,1) \overset{d}{=} r(s,a) + \gamma Z^\pi(S',A',1)$$

The correction disappears, leaving the standard distributional Bellman equation.

In fact, the flow Bellman operator admits Bellman consistency at all intermediate noise levels: we have a unique fixed point at every noise level.

Learning V- and Q-return velocities

We learn two return flows separately: a V-flow conditioned on the state, and a Q-flow conditioned on the state and action. Their noisy returns satisfy:

$$\begin{aligned} Z_V^\pi(s,\tau) &\overset{d}{=} Z_Q^\pi(s, A, \tau), \\ Z_Q^\pi(s,a,\tau) &\overset{d}{=} \tau\, r(s,a) + \gamma Z_V^\pi(S', \tau) + (1-\tau)\sqrt{1-\gamma^2}\, N. \end{aligned}$$

Each flow has a velocity network that moves noisy samples toward clean returns. The V-velocity learns from the Q-velocity at the current state. The Q-velocity learns from the V-velocity at the next state, using the flow Bellman equation at the same noise level.

Value Forcing losses $$\begin{align*} \mathcal{L}(v_V) &= \mathbb{E}_{\substack{ a\sim\pi(a\mid s),\,z^\tau\sim p_Q(z^\tau\mid s,a,\tau) }}\Big[\big(v_V(z^\tau\mid s,\tau)-v_Q(z^\tau\mid s,a,\tau)\big)^2\Big] \tag{1} \\ \mathcal{L}(v_Q) &= \mathbb{E}_{\substack{ s'\sim p(s'\mid s,a),\,\epsilon\sim\mathcal{N}(0,1),\,z^\tau\sim p_V(z^\tau\mid s',\tau) }}\Big[\Big(v_Q\!\Big(\tau(1-\gamma)r+\gamma z^\tau+(1-\tau)\sqrt{1-\gamma^2}\,\epsilon\,\Big|\,s,a,\tau\Big)\\ &\hspace{14.5em} -\Big((1-\gamma)r+\gamma v_V(z^\tau\mid s',\tau)-\sqrt{1-\gamma^2}\,\epsilon\Big)\Big)^2\Big] \tag{2} \end{align*}$$

Here $z^\tau\sim p_Q(z^\tau\mid s,a,\tau)$ and $z^\tau\sim p_V(z^\tau\mid s',\tau)$ denote samples from the noisy Q- and V-return distributions.

We alternate between updating V- and Q-velocity networks.

Reward scaling. We multiply rewards by $(1-\gamma)$ to normalize the return scale. This substantially improves performance (see Ablations).

Value Forcing Algorithms

We instantiate two offline RL algorithms: (1) evaluate the return distribution of an explicit policy and then perform policy improvement: VF (explicit method), and (2) learn return distributions through an expectile-weighted objective without an explicit actor: IVF (implicit method). Both algorithms use Eq. (2) to learn the Q-velocity but differ in how they learn the V-velocity.

Explicit actor-critic

VF

We define the explicit policy as performing rejection sampling on $K$ actions proposed by a flow BC policy $\pi_\beta$:

$$\pi(a \mid s) \overset{d}{=} \operatorname*{arg\,max}_{a_1,\,\dots,\,a_K:\ a_i \sim \pi_\beta(a \mid s)} Q(s, a_i),$$ where the Q-estimate is the expected return predicted by the Q-flow: $$Q(s,a) = \mathbb{E}_{z^0\sim\mathcal{N}(0,1)}\left[\int_0^1 v_Q(z^\tau\mid s,a,\tau)\,d\tau\right].$$ In practice, we integrate over $ M $ independent initial noises to approximate this expectation. VF uses this policy to train the V-velocity in Eq. (1).

Implicit value learning

IVF

Following prior work (e.g., Sec. 10.3.1 of Fahrmeir et al. (2022)), we define the expectile weight with coefficient $\mu \in [0.5,1)$ as:

$$\begin{aligned} w^\mu(s,a) &\triangleq \mu\,\mathbf{1}\big(Q(s,a) \ge V(s)\big) + (1-\mu)\,\mathbf{1}\big(Q(s,a) < V(s)\big), \end{aligned}$$ where $ V(s) \triangleq \mathbb{E}_{z^0 \sim \mathcal{N}(0, 1)} \left[ \int_0^1 v_V(z^\tau \mid s, \tau) d\tau \right]$ is the V-estimate. We then use the expectile weight to define a regression objective using actions in the dataset: $$\begin{aligned} \mathcal{L}^\mu(v_V) &= \mathbb{E}_{\textcolor{#2f8a4f}{\boldsymbol{a \sim D}}}\Big[w^\mu(s,a) \cdot \left( v_V(z^\tau \mid s,\tau) - v_Q(z^\tau \mid s,a,\tau) \right)^2 \Big]. \end{aligned}$$

This loss extends the scalar expectile regression loss (Kostrikov et al., 2022) to return flows.

ShowHide Algorithm 1

Algorithm 1 Value Forcing (VF) and Implicit Value Forcing (IVF)

Input: Dataset $D$, discount $\gamma$, expectile $\mu$, number of action candidates $K$, number of MC return samples $M$.

Initialize velocities $v_V^\theta, v_Q^\psi, v_\beta^\xi$ and initialize target velocity $\bar\psi \leftarrow \psi$.
for each training iteration do
$(s,a,r,s') \sim D,\ \tau \sim \mathrm{Unif}([0,1])$.
$a_\pi \sim$ rejection sampling, $z^\tau \sim p_Q^{\bar\psi}(z^\tau \mid s,a_\pi,\tau)$.
$\Delta\theta \leftarrow \nabla_\theta \mathcal{L}(v_V^\theta)$. Learn V-velocity for VF
Compute $w^\mu(s,a)$ using the expectile weights, $z^\tau \sim p_Q^{\bar\psi}(z^\tau \mid s,a,\tau)$.
$\Delta\theta \leftarrow \nabla_\theta \mathcal{L}^\mu(v_V^\theta)$. Learn V-velocity for IVF
$z^\tau \sim p_V^\theta(z^\tau \mid s',\tau),\ \epsilon \sim \mathcal{N}(0,1)$.
$\Delta\psi \leftarrow \nabla_\psi \mathcal{L}(v_Q^\psi)$. Learn Q-velocity
$\Delta\xi \leftarrow \nabla_\xi \mathcal{L}(v_\beta^\xi)$. Learn flow BC policy
$\bar\psi \leftarrow$ the exponential moving average of $\psi$. Update the target parameters
end for

Algorithm 1: Lines 4–5 are specific to VF and lines 6–7 are specific to IVF.

Experiments

We evaluate VF and IVF on 26 long-horizon, sparse-reward tasks: 20 from OGBench and 6 from OCBench. OGBench's play-style data comes from scripted policies pursuing random goals. OCBench uses near-expert demonstrations designed to emulate human-like demonstrations.

OGBench

cube-triple

cube-quadruple

puzzle-4x4

humanoidmaze-giant

OCBench

block-lite-single

block-lite-double

chamber-lite-easy

Value Forcing improves return estimation accuracy

We evaluate a fixed policy on one sampled state from each OGBench domain, plotting the return histograms predicted by different methods and their $1$-Wasserstein distances to the ground truth.

For each of four domains: predicted return histograms of VF, FlowCritic, Value Flows and PCBF against the ground truth, and bar charts of the mean 1-Wasserstein distance. VF has a lower distance than FlowCritic, Value Flows and PCBF in every domain. — **VF learns more accurate return distributions than prior flow-based return estimators.** *(Left)* We collect datasets to conduct policy evaluation experiments at a single state from each of the four OGBench domains. *(Center)* VF captures the ground-truth return distributions more accurately, while prior methods are biased and collapse their predictions into narrow peaks. *(Right)* VF achieves lower mean 1-Wasserstein distances than all three prior methods in all tasks. Tap to enlarge.

IVF learns more accurate return distributions than implicit distributional RL baselines. (Left) We conduct policy evaluation experiments on the same set of states from the four OGBench domains. (Center) IVF better captures the ground-truth return distributions, while Implicit C51, Implicit IQN, and DIVL produce shifted or overly concentrated predictions. (Right) IVF achieves lower mean 1-Wasserstein distances than all three baselines at the evaluated state in each of the 4 domains. Tap to enlarge.

Value Forcing outperforms prior offline RL methods on smaller datasets

We compare against prior explicit and implicit value-learning methods in an offline RL setting, experimenting with subsets of OGBench and OCBench datasets.

Bar chart of success rate versus OGBench dataset size (10%, 25%, 50% and 100%) for VF, IVF and 11 baselines, grouped into explicit and implicit methods; VF and IVF have the highest success in their groups at 10% and 25% of the OGBench data. A separate OCBench group is shown on the right. — **Value Forcing improves offline RL performance across dataset sizes.** Success rates are averaged over 20 OGBench tasks *(Left)* and 6 OCBench tasks *(Right)* on each dataset size. On OGBench, VF and IVF outperform their explicit and implicit baselines with 10% and 25% of the data and remain competitive with more data. On OCBench, VF outperforms selected explicit methods on the 25% dataset, while IVF leads among selected implicit methods on the full dataset. Hatched bars indicate distributional RL methods without flow matching.

ShowHide OGBench per-task results

Success rate with 10% of the dataset, mean ± std over 4 seeds.

Task	Floq	FlowCritic	Value Flows	PCBF	C51	IQN	QC	VF (Ours)	DIVL	Implicit C51	Implicit IQN	Implicit QC	IVF (Ours)
	Explicit methods								Implicit methods

ShowHide OCBench per-task results

Success rate with 25% of the dataset, mean ± std over 4 seeds.

Task	Value Flows	C51	QC	VF (Ours)	Implicit C51	Implicit QC	IVF (Ours)
	Explicit methods				Implicit methods

Key design decisions in Value Forcing

Policy extraction and expectile weighting

Figure 4. Success rates of alternative instantiations. Explicit methods: learning both v_Q and v_V (VF) and learning only v_Q perform similarly, while a one-step actor is much lower. Implicit methods: Q, V expectile weighting (IVF) is much higher than v_Q, v_V expectile or quantile weights. — For VF, the results suggest that learning a separate V-velocity is optional, but replacing rejection sampling with a one-step actor reduces success rates. For IVF, expectile weights based on scalar advantages clearly outperform expectile regression on velocities or quantile-based weights.

Reward scaling

BibTeX

bibtex

@article{zheng2026value,
  title={Value Forcing: Aligning Return Distributions at All Noise Levels},
  author={Zheng, Chongyi and Zhou, Zhiyuan and Li, Qiyang and Eysenbach, Benjamin and Levine, Sergey},
  journal={arXiv preprint arXiv:2611.02160},
  year={2026},
  url={https://arxiv.org/abs/2611.02160},
}