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

arXiv:2610.19883v1 Announce Type: new 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.

Related

• Reward-guided Fine-Tuning of One-Step Generative Models via Wasserstein Gradient Flow
• Analytical Correction for Subsampling Bias in Drifting Models
• Score-Calibrated Flow for Sampling from Unnormalized Densities with Applications to Generative Online Reinforcement Learning
• dFlowGRPO: Rate-Aware Policy Optimization for Discrete Flow Models

Source: arXiv cs.LG | 2026-10-15