Occupancy Modeling via Drifting: One-Step Generation Enables Direct Policy Extraction
This paper tackles unsupervised pre-training for reinforcement learning: how to learn from reward-free trajectories so that a good policy can be extracted once a task reward is revealed. The core idea is to replace costly, iterative flow-based occupancy models with one-step drifting models, which makes it possible to extract policies by differentiating directly through the generator instead of distilling a separate critic. This matters because it cuts inference cost, removes an entire distillation stage, and improves performance on standard continuous-control benchmarks.
The paper decomposes the occupancy-to-policy pipeline into three clear axes—target distribution, model class, and extraction method—and introduces TD-Drift, a Bellman-consistent one-step sampler for occupancy measures. The empirical story is coherent within its scope: drifting matches or beats flow matching at a fraction of the inference cost, and direct policy extraction via backprop through the one-step model is simpler and often better than critic distillation. Still, the evaluation is narrow (four DMC domains with ExORL data) and the gains over the prior flow-based method, while consistent on average, are domain-dependent. It is a solid, well-motivated incremental contribution that streamlines an existing generative pipeline.
The strongest element is the technical derivation of the TD-Drift objective. The authors correctly identify that flow matching requires expensive ODE integration for every bootstrap sample and every Monte Carlo value estimate during critic distillation, then build a one-step alternative by constructing a drifting field from the Bellman mixture of next states and bootstrapped samples (Eq. 12). Direct policy extraction is conceptually elegant: because the drift generator $f_\theta$ maps noise to a future state in a single forward pass, the value gradient $\nabla_a \hat{Q}$ can be computed by chaining the reward model Jacobian with the generator Jacobian (Eq. 15), bypassing the two-step distillation pipeline. The experiments also systematically cross all combinations of intention conditioning, model class, and extraction strategy.
The experimental scope is limited to four DMC domains with ExORL exploration data, with no tests on standard offline RL suites such as D4RL or on expert demonstrations. Intention conditioning—the main feature that distinguishes InFOM from earlier unconditional methods—shows mixed value: while InDOM-Direct achieves the highest mean score, the authors themselves observe that pushing the latent close to its prior invites the hypothesis that it carries almost no information, and the far simpler unconditional DOM-Direct is often competitive. The paper also admits that drifting trades cheaper inference for more training compute, yet pre-training wall-clock time is not reported, so the true cost of the proposed pipeline remains unclear. Finally, the drift model relies on a tanh final activation and min/max state normalization, which is a stronger distributional assumption than the normalization used by the flow baselines.
The evidence supports the central claim that drifting replaces flow matching competitively on the tested benchmarks. The authors report a 14% average return improvement for InDOM over InFOM and show that direct extraction outperforms critic distillation in seven of eight settings (Figure 6). The comparison to InFOM appears fair, since it is rerun on the same data and protocol. However, direct extraction is not attempted for flow models—the authors note it would require differentiating through the ODE solver—so the extraction comparison is not fully crossed across model classes. The forward-backward baseline operates in a fundamentally different non-generative framework and serves as an external reference rather than a direct ablation; notably, it still outperforms all occupancy methods on cheetah and walker.
Reproducibility is relatively strong. Hyperparameters are exhaustively documented in Appendix C (Tables 3 and 4), and implementation details including the sample-based TD-Drift field estimator (Eq. 16) and the choice of tanh activation to maintain differentiability under box constraints are given in Appendix D. The authors state that code will be made available on a website. The main practical blocker for an independent reproduction is that the drift models use different pre-training regimes than the flow baselines: longer training (1.5M–2M vs. 1M updates) and a lower learning rate ($10^{-4}$ vs. $3 \times 10^{-4}$), which means the comparison is not strictly controlled for pre-training compute. Additionally, the drifting objective introduces a kernel temperature $\kappa$ and a number of bootstrap samples $N^+$ that must be tuned per domain.
The growing interest in behavior foundation models, combined with recent advancements in generative modeling, has led to an appealing recipe for reinforcement learning from unlabeled data: pre-train a generative model to model occupancies, i.e., discounted state visitations, on reward-free trajectories, such that a policy can be easily extracted once a small reward-labeled dataset becomes available. We decompose this recipe along three axes: (i) the degree to which the multi-modal structure of occupancies must be captured, (ii) the class of generative model, and (iii) the method used for policy extraction. Along each axis, we consider existing choices, propose simpler alternatives, and validate each decision experimentally. This analysis culminates in a novel, streamlined method that leverages drifting, a modern one-step generative sampler, to model the occupancy. We show that this enables policy extraction via direct differentiation through the occupancy model, and bypasses the need for critic distillation. Overall, we observe that this nimble setup leads to improvements in performance on standard benchmarks.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.