PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2610.19883v1cs.LG

TL;DR

Reward-free interaction data motivates occupancy models, which predict discounted future-state distributions for later reward-based policy extraction. The paper studies occupancy targets, generative models, and extraction methods, then introduces drifting with direct policy extraction. Across several domains, TD-Drift methods outperform the flow-based pipeline while remaining simpler and more computationally efficient.

  • Problem

    Occupancy-based RL seeks to use reward-free data for policy extraction after rewards appear, but existing expressive generative pipelines can require costly sampling and multi-step extraction.

    Source: §1 · p. 1

  • Method

    The paper compares conditional and marginal occupancy targets, introduces TD-Drift for one-step drifting occupancy models, and extracts policies directly through the occupancy model.

    Source: §1 · p. 2, §4.1 · p. 4

  • Results

    Across several domains, TD-Drift methods outperform the existing flow-based pipeline while being conceptually simpler and computationally more efficient.

    Source: §6 · p. 9

  • Takeaways & Limitations

    Direct policy extraction from one-step drifting occupancies provides a simpler alternative to critic distillation, and intention conditioning can help but is not strictly necessary.

    Source: §1 · p. 2, §5.3 · p. 8

    Drifting requires a longer offline pre-training phase because its sample-based distributional distances are generally more expensive to compute than flow vector-field matching.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

Occupancy modeling offers a way to reuse reward-free trajectories: predict discounted future-state distributions, then use revealed rewards to select a policy. The paper analyzes target, model-class, and extraction choices and combines drifting with direct policy extraction into a simpler recipe.

1 INTRODUCTION

  • Reward-free trajectories can support later policy extraction because occupancy models predict discounted future-state distributions that make values for revealed rewards straightforward to compute.The paper frames this as a way to select or extract a policy once a reward becomes available.Source: §1 · p. 1
  • Prior flow-based extraction uses critic distillation because differentiating through its ODE sampler is not straightforward, while inferred intentions add complexity and can fail when modes are poorly inferred.The paper notes that the practical importance of these added components has been only partially explored.Source: §1 · pp. 1–2
  • The paper revisits three pipeline choices: conditional versus marginal occupancy targets, flow matching versus one-step drifting models, and direct policy extraction versus critic distillation.It introduces TD-Drift, a temporal-difference objective for drifting occupancy models.Source: §1 · p. 2
  • Across four DeepMind Control Suite domains with reward-free ExORL data, drifting is competitive with flow at lower sampling cost, intention conditioning helps but is unnecessary, and direct extraction outperforms critic-based extraction.The combined recipe models occupancies with drifting, optionally conditioned on intentions, and extracts policies directly.Source: §1 · p. 2
  • The paper contributes a systematic comparison of pipeline components, a temporal-difference objective for drifting occupancies, and a practical recipe combining drifting with direct extraction.Figure 1 organizes the choices as occupancy target, generative model, and policy extraction.Source: §1 · p. 2, Figure 1 · p. 2

2 RELATED WORK

The paper situates occupancy modeling among generative approaches to RL and reward-free offline pre-training. It emphasizes that occupancy models predict future-state distributions directly, while one-step drifting moves iterative transport from inference into training.

Generative modeling in RL

  • Occupancy models predict discounted future-state distributions directly, avoiding autoregressive rollout error accumulation but losing information about when each state is reached.They are the focus of this work and are discussed as an alternative to generative transition models unrolled for planning.Source: §2 · p. 3
  • Flow models can make differentiation through iterative sampling costly and unstable, so policy optimization has used workarounds such as distillation.The related-work passage cites Park et al. (2025) and Zheng et al. (2026) for such approaches.Source: §2 · p. 3

One-step generative models

  • Drifting models are one-step by construction: they move iterative transport from inference into training by minimizing a kernel-based discrepancy between generated and data samples.This contrasts with approaches that reduce network evaluations by changing probability paths or using consistency, shortcut, or velocity-based models.Source: §2 · p. 3

Unsupervised pre-training and occupancy models

  • Reward-free offline data supports pre-training that can later be adapted when rewards appear, including skill-discovery and behavior-foundation-model approaches.The related-work passage identifies the successor representation and occupancy as central objects in this setting.Source: §2 · p. 3

3 BACKGROUND

The background defines the reward-free MDP and occupancy framework, then describes flow matching and drifting as alternative ways to model distributions. Occupancy evaluation links a learned future-state distribution to values for arbitrary reward tasks.

Markov decision processes

  • A reward-free MDP specifies states, actions, transitions, discounting, and an initial-state distribution; a reward function turns it into a task-specific MDP whose objective is expected discounted cumulative reward.A policy is a conditional distribution over actions given a state.Source: §3 · p. 3

Occupancy models

  • An occupancy is the discounted distribution over future states after an initial state-action pair and subsequent actions under a policy.The occupancy satisfies a Bellman equation relating it to the immediate next-state distribution and later occupancy.Source: §3 · p. 3
  • The action-value function can be expressed using occupancy, so a modeled occupancy supports evaluating a policy for arbitrary reward tasks.This connects discounted state visitation to downstream policy evaluation.Source: §3 · p. 3

Flow matching and drifting

  • Flow matching learns a velocity field and samples by integrating an ODE, requiring many network evaluations; drifting instead trains a one-step generator toward a target distribution.Drifting moves generated samples through a field that attracts them to positive data samples and repels them from generated negative samples; the generator is trained toward a self-consistent fixed point.Source: §3 · p. 4

4 METHODS

The methods pipeline pre-trains occupancy models on reward-free transitions, then extracts a reward-maximizing policy using a reward model and frozen occupancy model. It compares conditionality, generative model class, and policy-extraction strategy.

4 METHODS

  • The pipeline fits occupancy models on reward-free transitions, then uses a reward model and the frozen occupancy model to learn a reward-maximizing policy.The paper examines three choices: whether to model occupancy modes, which generative model to use, and how to extract a policy.Source: §4 · p. 4

4.1 CONDITIONAL OR UNCONDITIONAL MODELING

  • Conditional models infer intentions and learn a family of occupancies, while unconditional models fit the marginal behavior-policy occupancy without an encoder.Jointly learning the occupancy model and intention latent adds complexity, and benefits may be offset by poor intention inference.Source: §4.1 · p. 4

4.2 MODEL CLASS

  • TD-Drift is a Bellman-consistent one-step sampler proposed to avoid flow matching’s inference cost, which affects both bootstrap targets and Monte Carlo value estimates.TD-Flow instead uses next-state and bootstrap regression terms along probability paths; TD-Drift generates bootstrap samples in one forward pass.Source: §4.2 · p. 5

4.3 POLICY EXTRACTION

  • For conditional occupancies, generalized policy improvement searches over actions and intentions; direct extraction differentiates through the drifting sampler instead of distilling a critic.Unconditional extraction needs no intention search, and direct extraction replaces the critic action gradient in the DDPG-BC objective.Source: §4.3 · p. 6

5 EXPERIMENTS

The experiments compare occupancy-model variants across four domains, testing model class, intention conditioning, and policy extraction. InDOM-Direct is strongest overall, while DOM-Direct averages 375 against 326 for InFOM.

5 EXPERIMENTS

  • The six variants cross intention conditioning, flow versus drift, and direct extraction; direct extraction is not tested for flow models because it would require differentiating through the ODE solver.Source: §5 · p. 6
  • Evaluation covers four DeepMind Control Suite domains, with reward-free pretraining data and reward-labeled data for policy extraction.Each domain has four tasks; occupancy models are pretrained on one million transitions and policies extracted using 500,000 reward-labeled transitions.Source: §5 · p. 7

5.1 IS DRIFTING AN EFFECTIVE SOLUTION TO OCCUPANCY MODELING?

  • Figure 3 reports per-domain returns for InFOM and InDOM, with percentage differences measured relative to InFOM.Source: Figure 3 · p. 7

5.2 IS CONDITIONAL OCCUPANCY MODELING NECESSARY?

5.3 IS DIRECT POLICY EXTRACTION COMPETITIVE?

  • Figure 6 compares direct with distilled-critic extraction for InDOM and DOM, reporting return changes relative to critic-based extraction.Source: Figure 6 · p. 8

5.4 IS POLICY EXTRACTION VIA DRIFT OCCUPANCIES MORE EFFICIENT?

  • Figure 7 reports policy-extraction wall-clock times extrapolated from 10K to 500K steps.Source: Figure 7 · p. 8

5.5 WHAT MATTERS FOR OCCUPANCY MODELING AND DOWNSTREAM POLICY EXTRACTION?

6 DISCUSSION

The discussion concludes that TD-Drift-based methods outperform the flow-based pipeline across several domains while remaining simpler and more computationally efficient. It also identifies training-time costs, evaluation boundaries, and directions for extending the pipeline.

6 DISCUSSION

  • Across several domains, TD-Drift methods outperform the flow-based pipeline while remaining conceptually simpler and computationally more efficient.The authors describe occupancy modeling with downstream policy extraction as a promising unsupervised RL technique.Source: §6 · p. 9
  • Drifting trades cheaper inference for more training compute, requiring longer offline pre-training whose cost is amortized over downstream tasks.The authors say the method is more efficient at inference time.Source: §6 · p. 9
  • Evaluation used self-supervised exploration data, not expert demonstrations, and did not inspect what the intention manifold encodes.The reason for the performance gap associated with intentions therefore remains open.Source: §6 · p. 9
  • The authors identify combining first-order and zero-order policy-optimization approaches, and scaling drifting occupancy models with latent-space kernels, as future directions.They also highlight learning the latent space for this purpose.Source: §6 · p. 9

A EXTENDED RELATED WORKS

  • Drifting had been used for policies and world models, but the authors report developing a recursive temporal-difference variant for occupancy models.They state that prior drifting applications targeted data-given quantities such as actions or next states.Source: A EXTENDED RELATED WORKS · p. 15
  • This work belongs to the generative successor-measure family and is closest to InFOM, whose pipeline it decomposes, updates, and streamlines.Prior generative approaches include γ-models, Bellman diffusion, TD-Flow, flow successor representations, and InFOM.Source: A EXTENDED RELATED WORKS · p. 15
  • Prior generative occupancy models distilled an explicit critic because iterative-sampler gradients are costly and poorly conditioned; this work studies direct critic-free extraction.The related-work discussion also notes that gradient-based extraction tends to outperform weighted regression (Park et al., 2024a).Source: A EXTENDED RELATED WORKS · p. 15

B ADDITIONAL RESULTS

The additional-results appendix supplies per-task return results and tables describing shared and domain-specific hyperparameters.

B ADDITIONAL RESULTS

  • Table 2 reports undiscounted episode return per task as mean ± s.e.m. over random seeds, grouped by domain, with four-task domain averages.It uses a fixed configuration per domain and method; bold marks the best method per row.Source: Table 2 · p. 15
  • Table 3 lists common hyperparameters for the method and baselines, with direct (In)DOM extraction inheriting relevant settings except the IQL expectile.Source: Table 3 · p. 16

C HYPERPARAMETERS AND TUNING PROTOCOL

The tuning protocol inherits InFOM settings where applicable and tunes parameters introduced or changed for the drift model. Training-step choices also differ between the pipelines.

C HYPERPARAMETERS AND TUNING PROTOCOL

  • The protocol inherits InFOM’s τ = 0.9 expectile setting and N = 16 flow/drift samples, without retuning its parameters to avoid comparing different tuning budgets.Source: C HYPERPARAMETERS AND TUNING PROTOCOL · p. 16
  • The drift model tunes γ, kernel temperature κ, and intention-encoder KL weight λKL, while using 10^-4 learning and target-update rates and longer pre-training than InFOM.Pre-training uses 1.5M steps for cheetah and 2M for other domains; the occupancy model stays frozen during policy extraction.Source: C HYPERPARAMETERS AND TUNING PROTOCOL · p. 16
  • Discount factors are selected per domain using the unconditional pipeline, then reused for intention-conditioned runs while λKL is swept; InFOM’s discount is also tuned per domain.Source: C HYPERPARAMETERS AND TUNING PROTOCOL · p. 16

C.1 EVALUATION PROTOCOL

The evaluation protocol scores configurations using multi-task, multi-seed averages and reports selected configurations with five seeds. It compares behavior cloning, InFOM, and Forward-Backward as baselines.

C.1 EVALUATION PROTOCOL

  • Configurations are scored by each domain’s four-task mean over three pre-training seeds, and only complete 3 × 4 grids are scored.Reported tables and method-comparison figures use each method’s best-performing configuration with five seeds.Source: §C.1 · p. 16
  • The evaluation compares behavior cloning, rerun InFOM, and Forward-Backward, with BC serving as a lower baseline for whether policy improvement beats behavior imitation.Forward-Backward is described as a powerful zero-shot RL baseline aligned with unlabeled-data pre-training and later reward adaptation.Source: §C.1 · p. 16
  • Table 4 reports domain-specific hyperparameters for the method and baselines.Source: Table 4 · p. 17

D IMPLEMENTATION DETAILS

The implementation estimates TD-Drift forces with Monte Carlo samples and supports both critic-based and direct policy extraction. It uses differentiable bounded state predictions and shared normalization across pre-training and extraction.

D IMPLEMENTATION DETAILS

  • The TD-Drift attraction estimator uses observed transitions and dataset next actions, while bootstrap samples estimate positive forces and independent samples estimate repulsion.For deterministic transitions, the estimator is consistent as the number of bootstrap samples grows; stochastic transitions require aggregation over next-state samples.Source: D IMPLEMENTATION DETAILS · p. 17
  • A final tanh activation keeps generated states in the normalized [−1, 1]^d box while retaining differentiability for direct policy gradients.The dataset is max/min normalized; the paper notes this is a stronger assumption than mean/standard-deviation normalization.Source: D IMPLEMENTATION DETAILS · p. 17
  • Normalization statistics are computed once from the pre-training dataset and reused during policy extraction, keeping the frozen occupancy model valid across both phases.The paper reports using this normalization for all drift results.Source: D IMPLEMENTATION DETAILS · p. 18
  • Reusing the same noise for predictions and bootstrap targets destabilizes TD learning through a positive feedback loop, so the implementation uses independent noise samples.The loop begins with identical prediction and target networks, and the EMA target feeds the resulting bias back into learning.Source: D IMPLEMENTATION DETAILS · p. 18, Figure 8 · p. 18
  • Pre-training warms up the occupancy model, optionally infers a policy intention, updates an EMA target, and freezes the learned model before extraction.The intention-conditioned algorithm includes a KL term, while the unconditioned variant sets the latent and KL term to zero.Source: D IMPLEMENTATION DETAILS · p. 19
  • Direct extraction evaluates candidate latent intentions, selects the highest estimated value, and updates the policy with a behavior-regularized DDPG-BC objective.The critic-based route instead uses a Monte Carlo value estimate or a distilled critic for the action gradient.Source: D IMPLEMENTATION DETAILS · pp. 18–19

E ABLATIONS

An initial warm-up spread is crucial for TD-Drift: without it, DOM and DOM-Direct collapse to behavior-cloning baseline performance on the evaluated quadruped tasks.

E.1 AN INITIAL SPREAD MAKES TD-DRIFT WORK

  • Without the 10K-step warm-up, DOM and DOM-Direct collapse to behavior-cloning baseline performance across four quadruped tasks.The warm-up spreads the initial drifting distribution toward dataset coverage, making next-state supervision informative; the comparison is shown in Figure 9.Source: §E.1 · p. 20

E.2 MEDIAN TEMPERATURE RESCALING

Median-distance rescaling addresses variation in pairwise distance scales across datasets, making a shared temperature choice more meaningful. The tested temperature field is a reasonable default, but not a consistent winner.

E.2 MEDIAN TEMPERATURE RESCALING

  • Mapping states to [−1, 1]^d fixes coordinate ranges but not pairwise-distance scales, which still vary with domain dimensionality and geometry.After warm-up, the inter-sample distance distribution is fixed but initially unknown, motivating dataset-specific rescaling.Source: §E.2 · p. 20
  • The temperature field (0.05, 0.1, 0.2) is best in only two of 16 domain × task pairs, though it is usually close to the best tested temperature.The authors therefore describe it as a reasonable default rather than a clear win.Source: §E.2 · p. 20
  • Dividing distances by each dataset’s median expected inter-sample distance unifies their scales and makes one default temperature choice meaningful.Figure 10 contrasts per-domain min-max normalization of distances with division by the median pairwise distance.Source: §E.2 · p. 20, Figure 10 · p. 20

Related papers

Loading 2610.19883v1…