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.
Figures & tables
Figure 1: A generative pipeline for unsupervised RL. Starting from an unlabeled dataset D , the occupancy can be modeled in a conditional or unconditional fashion (i). Independently of this choice, the occupancy can be learned through flow matching or drifting (ii). In both cases, the policy can be extracted in a two-step process, but drifting also allows for direct policy extraction (iii). We study all combinations of these choices, and recommend those marked by a bold arrow.
Figure 2: Attracting TD-Drift field Vqθˉπ+ . Each sample from the pushforward distribution (blue) is attracted with weight 1−γ (light red) towards the next state s′ , while the attraction towards the bootstrap samples sˉ+ (bold red) is weighted with γ .
Figure 3: Drifting vs. flow matching. Return of InFOM and InDOM per domain (mean ± s.e.m.). The percentage is the performance difference of InDOM relative to InFOM.
Figure 4: KL weight of the intention encoder. Return of InDOM over the KL weight. Each sweep uses a fixed configuration per domain. (mean ± s.e.m. over three random seeds).
Figure 5: Effect of intention conditioning. Return with and without intention conditioning for the drift (DOM, InDOM) and the flow (FOM, InFOM) occupancy models. The percentage is the change in performance due to conditional measure modeling. (mean ± s.e.m.)
Figure 6: Critic-based vs. direct extraction. Return of InDOM and DOM with policy extraction through a distilled critic and direct policy extraction. The percentage is the performance change of direct policy extraction over extraction from a distilled critic. (mean ± s.e.m.)
Figure 7: Wall-clock time for policy extraction . Extrapolated: 10 K to 500 K steps. (mean ± s.e.m.).
Environment
BC
FB
InFOM
DOM-Direct
InDOM-Direct
cheetah
142±
2
631±
16
305±
13
266±
14
331±
14
walker
232±
2
780±
12
364±
7
390±
15
474±
13
quadruped
501±
3
601±
18
458±
37
682±
6
592±
6
jaco
3±
1
2±
1
175±
3
161±
3
201±
5
Mean
220
503
326
375
400
Table 1: Comparison with existing methods. Four-task mean return per domain (mean ± s.e.m. over five pre-training seeds), each method at its best discount and configuration per domain. FB are rerun using the codebase of Bagatella et al. (2025) . Per-task numbers and the remaining variants are in Table 2 .
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Task
BC
FB
InFOM
FOM
InDOM
DOM
InDOM-Direct
DOM-Direct
cheetah
Walk
212±
9
892±
32
498±
14
478±
23
432±
25
395±
15
438±
26
362±
21
Walk Backward
256±
6
913±
35
540±
45
527±
50
609±
30
637±
38
702±
59
576±
64
Run
45±
3
469±
8
96±
3
97±
3
97±
8
62±
2
81±
6
56±
2
Run Backward
54±
2
252±
30
84±
3
81±
3
101±
12
84±
4
105±
11
68±
4
Mean
142±
2
631±
16
305±
13
296±
15
310±
13
295±
9
331±
14
266±
14
Appendix
Table 2: Per-task results, all variants. Undiscounted episode return per task (mean ± s.e.m. over random seeds), grouped by domain with fixed configuration per domain and method; Mean rows give the four-task domain average. FB numbers are rerun using the codebase from Bagatella et al. (2025) . InFOM and FOM numbers are rerun using the codebase from Zheng et al. (2026) . Bold marks the best method per row.
Hyperparameter
(In)DOM
(In)FOM
learning rate
10−4
3×10−4
optimizer
Adam
Adam
pre-training steps
1.5 M ( cheetah ), 2 M (else)
1 M
policy extraction / finetuning steps
500 K
500 K
batch size
256
256
target network update rate
10−4
5×10−4
Appendix
Table 3: Common hyperparameters for our method and baselines . For direct policy extraction via (In)DOM, we inherit all relevant hyperparameters except the IQL expectile.
Environment
Parameter
InFOM
FOM
InDOM
DOM
InDOM-Direct
DOM-Direct
cheetah
γ
0.98
0.98
0.99
0.99
0.99
0.99
λKL
0.05
–
10
–
10
–
α
0.3
0.3
0.3
0.3
0.3
0.3
κ
–
–
0.2
(0.05,0.1,0.2)
0.1
0.1
M
–
–
–
–
10
–
dim (z)
128
–
128
–
128
–
Appendix
Table 4: Domain-specific hyperparameters for our method and the baselines .
Figure 8: Positive feedback loop from noise coupling (iii) shared ϵ biases the kernel towards self-attraction; (iv) the EMA update feeds this bias into the target, closing the loop.
Figure 9: Effect of warming up initial spread . Policy return on quadruped for DOM and DOM-Direct with (w) and without (w/o) warm-up phase to generate an initial spread.
Figure 10: Rescaling kernel distances per dataset . Two histograms of inter-sample distances. The left one is based on min-max normalization per data domain. Right one: each distance is divided by the median pairwise distance (dashed line in left histogram)
Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.
Deep Reinforcement Learning (DRL) is widely recognized as sample-inefficient, a limitation attributable in part to the high dimensionality and substantial functional redundancy inherent to the policy parameter space. A recent framework, which we refer to as Action-based Policy Compression (APC), mitigates this issue by compressing the parameter space Θ into a low-dimensional latent manifold Z using a learned generative mapping g:Z→Θ. However, its performance is severely constrained by relying on immediate action-matching as a reconstruction loss, a myopic proxy for behavioral similarity that suffers from compounding errors across sequential decisions. To overcome this bottleneck, we introduce Occupancy-based Policy Compression (OPC), which enhances APC by shifting behavior representation from immediate action-matching to long-horizon state-space coverage. Specifically, we propose two principal improvements: (1) we curate the dataset generation with an information-theoretic uniqueness metric that delivers a diverse population of policies; and (2) we propose a fully differentiable compression objective that directly minimizes the divergence between the true and reconstructed mixture occupancy distributions. These modifications force the generative model to organize the latent space around true functional similarity, promoting a latent representation that generalizes over a broad spectrum of behaviors while retaining most of the original parameter space's expressivity. Finally, we empirically validate the advantages of our contributions across multiple continuous control benchmarks.
Andrea Fraschini, Davide Tenedini, Riccardo Zamboni +2
We study sequential decision-making with offline reinforcement learning (RL). Traditional offline RL policies may result in out-of-distribution (OOD) actions when training relies only on sparse offline representations. To ensure safe offline policies in a sparse state-action space, we explore how density estimation models can be integrated into model-based RL methods to avoid the OOD regions. Generative models are capable of explicitly modeling the density in sparse state-action spaces. Building on this, we introduce Generative OOD-regularized Model-based Policy Optimization (GORMPO), a density-regularized offline RL algorithm that uses generative density modeling to restrict policy updates to high-density areas of the dataset. Furthermore, we examine whether better OOD detection corresponds to better model-based offline policies. We compare (1) the OOD detection capabilities of various density estimators and (2) their performance within the GORMPO framework on a real-world medical dataset and sparse offline RL datasets. We theoretically guarantee GORMPO's performance under mild assumptions. Empirically, GORMPO outperforms state-of-the-art baselines by 17% on a real-world medical dataset and enhances the base model on the offline RL datasets. Our empirical findings show that better OOD detection generally results in improved policies in environments with stable dynamics, while conservative penalties with poor density estimation are favored when dynamics are uncertain.
Generative policies offer expressive action distributions for online reinforcement learning (RL), but many diffusion- and flow-based policies require costly iterative sampling. We propose DriftSpread, a one-step generative policy based on Drifting for online maximum-entropy (MaxEnt) RL. To incorporate entropy without a tractable policy density, we build on an energy-based view of Drifting and derive critic-guided Drift and entropy-induced Spread from the negative MaxEnt objective. Our method learns these updates sequentially in the same actor, using action regression for value improvement and distribution matching to capture the effect of entropy. The actor updates require neither direct samples from the critic-defined Boltzmann target nor policy-density or score estimation. Under exact fitting with a fixed critic, we prove convergence to the Boltzmann target in the small-step limit followed by the long-time limit. On MuJoCo, DriftSpread achieves the highest aggregate normalized return among all evaluated methods, while outperforming the diffusion- and flow-based baselines in both training and inference efficiency.
Chunyu Luo, Hongyu Zhao, Guangyi Ni +2
The University of Hong Kong · Beihang University · Harbin Institute of Technology, Shenzhen +2