Generative Value Classifier

A distributional critic predicts the return distribution \(p(z\mid s,a)\) over a fixed set of return bins \(z_1,\dots,z_K\). The Generative Value Classifier (GVC) recovers it from the likelihood of the action under each bin, with Bayes’ rule:

The critic has two parts. The action model \(p_\theta(a\mid s,z)\) is a flow-matching policy conditioned on a return bin, and it can share the actor’s architecture. The return prior \(p_\psi(z\mid s)\) predicts the return distribution from the state alone. Flows do not give exact likelihoods, so we use the conditional flow-matching error of \(a\) under each bin as a surrogate for \(-\log p_\theta(a\mid s,z_i)\) and put tuned exponents on the two factors. A shared trunk encodes the state, the noised action, and the flow time, and each bin adds a learned embedding before a shared head, so one forward pass scores all \(K\) bins.

Training. Both parts of the critic learn from the same TD target. Let \(\widehat p\) be the distribution of \(r+\gamma Z'\) projected onto the return bins, where \(Z'\) is the target critic’s return distribution at the next state and action. The flow-matching loss on the dataset action is weighted across bins by \(\widehat p\), and the return prior is fit to \(\widehat p\) with a cross-entropy loss:

Because the flow error is only a surrogate, with unknown scale and bin-dependent offsets, these two losses do not ensure that \(Q(s,a)\) matches the target mean \(y=\sum_i\widehat p_i z_i\). A squared-error value anchor, \(\mathcal L_{\mathrm{anchor}}=\big(Q(s,a)-y\big)^2\), penalizes this error directly. The full objective is \(\mathcal L=\mathcal L_{\mathrm{actor}}+\lambda_F\mathcal L_{\mathrm{flow}}+\lambda_R\mathcal L_{\mathrm{return}}+\lambda_V\mathcal L_{\mathrm{anchor}}\).

Acting. The actor is a separate flow policy trained by behavior cloning on dataset or replay actions. At each step it proposes \(N\) candidate actions, GVC scores them, and we execute the one with the highest \(Q(s,a)\). The next action \(a'\) in the TD target is chosen the same way.

We follow the training recipe of QC on four RoboMimic manipulation tasks with sparse task-completion rewards and change only the critic. Lift, Can, and Square are eventually solved by GVC, QC, and IDQL, with small sample-efficiency gains for GVC. Tool Hang, a long-horizon task in which the robot inserts a hook into a base and then hangs a wrench from it, is where the critics diverge: GVC reaches about 85% success within 0.5M steps, whereas QC and RLPD only reach about 50% at 1M steps.

We evaluate on 35 tasks across seven OGBench domains, with 8 seeds per task. The closest comparison is BFN (from QC), which uses the same flow actor and best-of-\(N\) selection as GVC but a regression critic. On the two largest puzzles, GVC reaches 92% and 48% on the first task of puzzle-4x5 and puzzle-4x6, against 83% and 22% for BFN. It leads BFN by 12 points on the sparse scene tasks, mostly on task5 (97% vs. 45%), and the two are nearly tied on cube-double (65% vs. 66%). Compared with Value Flows, the strongest distributional baseline, GVC is ahead on puzzle-4x4, puzzle-4x5, and puzzle-4x6 by 8 points on average and behind on cube-double and cube-triple by 4.

We ablate the generative loss, which is what distinguishes GVC from a discriminative critic, and the noise sampling used to score actions. All runs are on Tool Hang with 8 seeds per variant.

GVC spends more compute per scored action than a regression critic. Scoring one candidate takes eight forward passes of the action model (\(T=4\) flow times, each with an antithetic noise pair), and each pass returns all \(K\) bins. On puzzle-4x4 on a single H100, a GVC update takes 15.6 ms, comparable to Value Flows (13.5 ms) and IQN (13.1 ms), the two distributional critics that also evaluate many return bins or quantiles per transition, and 5–9× the 1.8–3.4 ms of the scalar regression critics. Splitting the critic into two models and discretizing the return space also add implementation complexity.