Occupancy Modeling via Drifting: One-Step Generation Enables Direct Policy Extraction

Robin Uhrich; Marco Bagatella; Núria Armengol Urpí; Pavel Kolev; Georg Martius

cs.LGOct 14, 2026

Occupancy Modeling via Drifting: One-Step Generation Enables Direct Policy Extraction

Authors: Robin Uhrich, Marco Bagatella, Núria Armengol Urpí, Pavel Kolev, Georg Martius

Organizations: University of Tübingen · ETH Zürich · Max Planck Institute for Intelligent Systems

Abstract

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

\mathcal{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.

Panel of Figure 2 — Figure 2: Attracting TD-Drift field $\mathbf{V^{+}_{q_{\bar{\theta}}^{\pi}}}$ . Each sample from the pushforward distribution (blue) is attracted with weight $1-\gamma$ (light red) towards the next state $s^{\prime}$ , while the attraction towards the bootstrap samples $\bar{s}^{+}$ (bold red) is weighted with $\gamma$ .

Figure 3: Drifting vs. flow matching. Return of InFOM and InDOM per domain (mean

\pm

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

\pm

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

\pm

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

\pm

s.e.m.)

Figure 7: Wall-clock time for policy extraction . Extrapolated:

10

K to

500

K steps. (mean

\pm

s.e.m.).

Environment	BC		FB		InFOM		DOM-Direct		InDOM-Direct
cheetah	$142$ $\pm$	$2$	$631$ $\pm$	$16$	$305$ $\pm$	$13$	$266$ $\pm$	$14$	$\mathbf{331}$ $\pm$	$14$
walker	$232$ $\pm$	$2$	$780$ $\pm$	$12$	$364$ $\pm$	$7$	$390$ $\pm$	$15$	$\mathbf{474}$ $\pm$	$13$
quadruped	$501$ $\pm$	$3$	$601$ $\pm$	$18$	$458$ $\pm$	$37$	$\mathbf{682}$ $\pm$	$6$	$592$ $\pm$	$6$
jaco	$3$ $\pm$	$1$	$2$ $\pm$	$1$	$175$ $\pm$	$3$	$161$ $\pm$	$3$	$\mathbf{201}$ $\pm$	$5$
Mean	$220$		$503$		$326$		$375$		$\mathbf{400}$

Table 1: Comparison with existing methods. Four-task mean return per domain (mean

\pm

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$ $\pm$	$9$	$892$ $\pm$	$32$	$\mathbf{498}$ $\pm$	$14$	$478$ $\pm$	$23$	$432$ $\pm$	$25$	$395$ $\pm$	$15$	$438$ $\pm$	$26$	$362$ $\pm$	$21$
Walk Backward	$256$ $\pm$	$6$	$913$ $\pm$	$35$	$540$ $\pm$	$45$	$527$ $\pm$	$50$	$609$ $\pm$	$30$	$637$ $\pm$	$38$	$\mathbf{702}$ $\pm$	$59$	$576$ $\pm$	$64$
Run	$45$ $\pm$	$3$	$469$ $\pm$	$8$	$96$ $\pm$	$3$	$\mathbf{97}$ $\pm$	$3$	$\mathbf{97}$ $\pm$	$8$	$62$ $\pm$	$2$	$81$ $\pm$	$6$	$56$ $\pm$	$2$
Run Backward	$54$ $\pm$	$2$	$252$ $\pm$	$30$	$84$ $\pm$	$3$	$81$ $\pm$	$3$	$101$ $\pm$	$12$	$84$ $\pm$	$4$	$\mathbf{105}$ $\pm$	$11$	$68$ $\pm$	$4$
Mean	$142$ $\pm$	$2$	$631$ $\pm$	$16$	$305$ $\pm$	$13$	$296$ $\pm$	$15$	$310$ $\pm$	$13$	$295$ $\pm$	$9$	$\mathbf{331}$ $\pm$	$14$	$266$ $\pm$	$14$

Appendix

Table 2: Per-task results, all variants. Undiscounted episode return per task (mean

\pm

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\times 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\times 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	$\gamma$	0.98	0.98	0.99	0.99	0.99	0.99
	$\lambda_{\mathrm{KL}}$	0.05	–	10	–	10	–
	$\alpha$	0.3	0.3	0.3	0.3	0.3	0.3
	$\kappa$	–	–	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 .

Panel of Figure 8 — Figure 8: Positive feedback loop from noise coupling (iii) shared $\epsilon$ 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

Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.

See all 1,727 papers

Unsupervised Behavioral Compression: Learning Low-Dimensional Policy Manifolds through State-Occupancy Matching
Mar 27, 2026Andrea Fraschini, Davide Tenedini, Riccardo Zamboni +2RL Control Distribution Matching
Politecnico di Milano · Technion
Generative OOD-regularized Model-based Policy Optimization
May 23, 2026Aysin Tumay, Jiahe Huang, Elise Jortberg +1Offline RL OOD Detection
University of California, San Diego · Abiomed
Drift, Then Spread: Learning One-Step Generative Policies for Online Reinforcement Learning
Oct 11, 2026Chunyu Luo, Hongyu Zhao, Guangyi Ni +2Reinforcement Learning Maximum Entropy RL
The University of Hong Kong · Beihang University · Harbin Institute of Technology, Shenzhen +2

Mar 27, 2026J/K move · Enter open · S save

Unsupervised Behavioral Compression: Learning Low-Dimensional Policy Manifolds through State-Occupancy Matching

Andrea Fraschini, Davide Tenedini, Riccardo Zamboni, Mirco Mutti+1

Politecnico di Milano · Technion

RL Control Distribution Matching Sample-Efficient RL

Open paper AskPDF arXiv

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

\mathcal Z

using a learned generative mapping

g:\mathcal Z \to Θ

. 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.

Abstract

Reinforcement Learning

Explore similar work

Unsupervised Behavioral Compression: Learning Low-Dimensional Policy Manifolds through State-Occupancy Matching

Generative OOD-regularized Model-based Policy Optimization

Drift, Then Spread: Learning One-Step Generative Policies for Online Reinforcement Learning