PaperlayerView exact revision on arXiv

Source-linked AI summary

Intention-Conditioned Flow Occupancy Models

Chongyi Zheng, Seohong Park, Sergey Levine, Benjamin Eysenbach

arXiv:2506.08902v4cs.LGcs.AI

TL;DR

RL pre-training remains challenging because agents must reason about long-term action effects, while offline data may combine users pursuing different tasks. InFOM combines latent intentions with flow-matched occupancy prediction; across 36 state-based and 4 image-based benchmark tasks, it reports 1.8× median improvement in returns and 36% improvement in success rates over alternative pre-training methods.

  • Problem

    RL pre-training must capture long-term temporal dependencies, and offline datasets can combine distinct users and tasks whose intentions are unobserved.

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

  • Method

    InFOM uses a latent-variable model to capture intentions and flow matching to learn occupancy models that predict temporally distant future states.

    Source: §1 · p. 2, §4.3 · p. 5

  • Results

    Across 36 state-based and 4 image-based benchmark tasks, InFOM reports 1.8× median improvement in returns and 36% improvement in success rates over alternative pre-training methods.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors report that modeling intentions and long-term behaviors from unstructured data supports effective and efficient fine-tuning across state- and image-based domains.

    Source: §6 · p. 10

    Reducing trajectories to consecutive state-action pairs may not accurately capture their original intentions, although the approach was sufficient for strong benchmark performance.

    Source: §6 · p. 10

Abstract

from arXiv · show

Large-scale pre-training has fundamentally changed how machine learning research is done today: large foundation models are trained once, and then can be used by anyone in the community (including those without data or compute resources to train a model from scratch) to adapt and fine-tune to specific tasks. Applying this same framework to reinforcement learning (RL) is appealing because it offers compelling avenues for addressing core challenges in RL, including sample efficiency and robustness. However, there remains a fundamental challenge to pre-train large models in the context of RL: actions have long-term dependencies, so training a foundation model that reasons across time is important. Recent advances in generative AI have provided new tools for modeling highly complex distributions. In this paper, we build a probabilistic model to predict which states an agent will visit in the temporally distant future (i.e., an occupancy measure) using flow matching. As large datasets are often constructed by many distinct users performing distinct tasks, we include in our model a latent variable capturing the user intention. This intention increases the expressivity of our model, and enables adaptation with generalized policy improvement. We call our proposed method intention-conditioned flow occupancy models (InFOM). Comparing with alternative methods for pre-training, our experiments on $36$ state-based and $4$ image-based benchmark tasks demonstrate that the proposed method achieves $1.8 \times$ median improvement in returns and increases success rates by $36\%$. Website: https://chongyi-zheng.github.io/infom Code: https://github.com/chongyi-zheng/infom

1 INTRODUCTION

The introduction frames RL pre-training as difficult because agents must reason about long-term effects and datasets may combine users pursuing different tasks. InFOM models temporally distant future states while capturing latent intentions, using flow matching to support pre-training and fine-tuning.

1 INTRODUCTION

  • Pre-training and fine-tuning are established in other fields, but applying this paradigm to RL remains open because agents must reason across time and data may combine users performing distinct tasks.The passage contrasts this with foundation-model successes in vision and language.Source: §1 · p. 1
  • Prior RL foundation-model attempts often predict pre-training actions, while occupancy models can predict future events but are typically hard to train and ignore user intentions.The introduction identifies flow matching and diffusion as tools for modeling complex distributions with varied inputs.Source: §1 · p. 2
  • Across 36 state-based and 4 image-based benchmark tasks, InFOM reports 1.8× median improvement in returns and 36% improvement in success rates over alternative pre-training methods.The method also reports intention inference and efficient policy extraction in additional experiments.Source: §1 · p. 2
  • InFOM learns a probabilistic model of temporally distant future states and latent intentions, using flow matching to model occupancy measures.The paper describes this as capturing both time and intention in RL pre-training.Source: §1 · p. 2

2 RELATED WORK

Related work spans offline unsupervised RL, representation learning, generative models, and successor representations. InFOM is positioned as combining latent intention learning and occupancy modeling, with implicit generalized policy improvement for policy extraction.

Offline unsupervised RL

  • InFOM is conceptually related to offline unsupervised skill learning that predicts intentions, but it does not learn multiple skills during pre-training.The authors describe their work as complementary to behavioral-cloning pre-training, with additional gains from modeling intentions and occupancy measures together.Source: §2 · p. 2

Unsupervised representation learning for RL

  • Representation-learning methods provide inputs for policy and value networks, but whether those representations facilitate policy adaptation is often unclear.The passage notes an exception to this uncertainty.Source: §2 · pp. 2–3

RL with generative models

  • Generative models have been used to model trajectories, dynamics, skills, policies, and values, while InFOM uses flow matching for discounted state occupancy measures.The passage identifies the flow-matching objective as state of the art.Source: §2 · p. 3

Successor representations and successor features

  • Unlike prior flow-based occupancy work, InFOM learns intentions with a latent-variable model and uses an expectile loss for implicit generalized policy improvement.The paper reports that these choices lead to higher returns and success rates.Source: §2 · p. 3

3 PRELIMINARIES

The preliminaries define discounted state occupancy as a summary of policy visitation and relate it to value estimation. They also introduce flow matching and its temporal-difference SARSA variant for learning occupancy models.

Occupancy measures

  • A discounted state occupancy measure quantifies how frequently different states are visited under a policy, with discounting.Prior work gives it a Bellman recursion that backs up probability density at current and future time steps.Source: §3 · p. 3
  • The Q-function can be written as the expected reward under the discounted state occupancy measure, scaled by 1/(1 − γ).This relation permits policy evaluation by sampling occupancy states and regressing Q toward their average rewards.Source: §3 · pp. 3–4
  • The occupancy recursion used here starts from the current time step rather than the next time step, unlike some prior approaches.The passage cites prior approaches that start from the next time step.Source: §3 · p. 3

Flow matching and TD flows

  • Flow matching models distributions with ordinary differential equations, whose deterministic nature gives more stable learning objectives and faster inference than denoising diffusion models.The passage contrasts ODE-based flow matching with diffusion models that solve stochastic differential equations.Source: §3 · p. 4
  • TD flows incorporate the occupancy Bellman equation into conditional flow matching, and the SARSA variant uses the behavioral policy as the target policy.The paper uses this SARSA variant to learn its generative occupancy models.Source: §3 · p. 4

4 INTENTION-CONDITIONED FLOW OCCUPANCY MODELS

InFOM pre-trains on offline reward-free data by inferring latent user intentions and modeling intention-conditioned future-state occupancy, then adapts to downstream rewards through generalized policy improvement.

4 INTENTION-CONDITIONED FLOW OCCUPANCY MODELS

  • InFOM combines latent intention inference and flow-based long-horizon occupancy prediction, then uses generalized policy improvement to adapt policies to downstream rewards.The method has separate pre-training and fine-tuning stages.Source: §4 · p. 4, §4.1 · p. 4

4.1 PROBLEM SETTING

  • Pre-training uses unlabeled transitions from a mixture of policies, with latent intentions representing distinct tasks and assumed consistent across consecutive transitions.The downstream task supplies a separate reward-labeled dataset.Source: §4.1 · pp. 4–5

4.2 VARIATIONAL INTENTION INFERENCE

  • The variational encoder infers intention from the next transition while occupancy models predict future states, jointly trained with flow matching and KL regularization.The intention prior is a standard Gaussian, and the encoder and occupancy model are optimized jointly.Source: §4.2 · p. 5

4.3 PREDICTING THE FUTURE VIA SARSA FLOWS

  • Flow occupancy models target multimodal discounted future-state distributions and stitch shared trajectory segments using a SARSA temporal-difference flow loss.Intention conditioning permits a simpler, more stable SARSA bootstrap than a Q-learning-style bootstrap and supports generalization across intentions while avoiding counterfactual errors.Source: §4.3 · pp. 5–6

4.4 GENERATIVE VALUE ESTIMATION AND IMPLICIT GENERALIZED POLICY IMPROVEMENT

  • Implicit GPI distills intention-conditioned Q functions with an upper expectile, then optimizes a behavior-regularized policy instead of maximizing over sampled intentions.Behavioral-cloning regularization reduces errors from out-of-distribution actions and overestimated Q values; the approach avoids instability from differentiating through the ODE solver.Source: §4.4 · pp. 6–7
  • For downstream tasks, InFOM estimates intention-conditioned Q values by averaging rewards over future states sampled from the pretrained occupancy model.The method samples intentions from the prior; the number of future states affects estimate accuracy and variance.Source: §4.4 · p. 6

5 EXPERIMENTS

The experiments compare InFOM with prior pre-training methods, examine whether its latent intentions align with task behaviors, and ablate its policy-extraction strategy. Across the reported comparisons, InFOM performs competitively on diverse domains, while its intention encoder and implicit GPI show benefits in the studied settings.

5.1 COMPARING TO PRIOR PRE-TRAINING AND FINE-TUNING METHODS

  • The experiments compare reward-free pre-training followed by reward-labeled fine-tuning across 36 state-based and 4 image-based tasks, using 8 baselines.The evaluated baselines include policies, transition models, representations, occupancy classifiers, and latent skills.Source: §5.1 · pp. 7–8
  • 36% higher success rate over the best performing baseline on challenging state-based OGBench manipulation tasks, and 31% improvement over the best baseline using RGB images.InFOM also obtains 20× improvement on jaco, where baselines make only trivial progress.Source: §5.1 · p. 8

5.2 VISUALIZING LATENT INTENTIONS

  • InFOM infers clearly clustered latent intentions for pick and place behaviors, while HILP + FOM and FB + FOM infer mixed intentions.The visualization uses t-SNE on cube double task 1; the paper also reports that InFOM outperforms these two variants on 3 of 4 tasks in additional experiments.Source: §5.2 · pp. 8–9, Figure 4 · p. 9

5.3 IMPORTANCE OF THE IMPLICIT GENERALIZED POLICY IMPROVEMENT

  • InFOM is 44% more performant with 8× smaller variance than its standard-GPI variant, and achieves higher returns and success rates than one-step PI.These policy-extraction ablations were evaluated on quadruped jump and scene task 1.Source: §5.3 · pp. 9–10, Figure 5 · p. 9
  • Additional ablations report faster policy learning, generally better performance from continuous than vector-quantized discrete latents, and the importance of sufficient pre-training and fine-tuning data.Other appendix experiments examine dataset diversity, suboptimal fine-tuning datasets, and hyperparameters.Source: §5.3 · p. 10
  • The paper omits new comparisons with alternative generative occupancy models because prior work found flow-based occupancy models outperforming diffusion, VAE, and GAN approaches.The authors state that these comparisons do not distinguish their contribution.Source: §5.3 · p. 10

6 CONCLUSION

The conclusion presents InFOM as a way to capture diverse intentions and long-term behaviors from unstructured data, with reported benefits for fine-tuning across state- and image-based domains. It also notes that reducing trajectories to consecutive state-action pairs may not fully preserve the original intentions.

6 CONCLUSION

  • InFOM captures diverse intentions and long-term behaviors from unstructured data using flow models, and outperforms prior unsupervised pre-training approaches across state- and image-based domains.The authors report that these learned intentions enable effective and efficient fine-tuning.Source: §6 · p. 10
  • Reducing trajectories to consecutive state-action pairs may not accurately capture their original intentions, though the approach achieves strong performance on the benchmark tasks.The authors leave alternative trajectory encodings and data-collection strategies for future work.Source: §6 · p. 10

REPRODUCIBILITY STATEMENT

The appendix describes the implementation and supporting RL machinery, from actor-critic updates and flow matching to temporal-difference occupancy learning. It also explains how the consistency assumption lets an intention remain associated with its behavioral policy across transitions.

REPRODUCIBILITY STATEMENT

  • All methods share a JAX codebase; experiments ran on one NVIDIA A6000 GPU, taking up to 4 hours for state-based tasks and 12 hours for image-based tasks.The authors provide open-source implementations.Source: REPRODUCIBILITY STATEMENT · p. 10

A.1 VALUE FUNCTIONS AND THE ACTOR-CRITIC FRAMEWORK

  • The actor-critic framework alternates critic evaluation and policy improvement, while the fine-tuning algorithm updates the flow model, intention encoder, reward predictor, critic, and policy.The actor-critic passage describes behavioral policy data as a replay buffer in online RL or a fixed dataset in offline RL.Source: §A.1 · pp. 20–21

A.2 FLOW MATCHING

  • Flow matching transforms Gaussian noise into a target distribution by integrating a time-dependent vector field, with Euler steps used to solve the flow ODE.The algorithm uses 10 Euler steps by default.Source: §A.1 · p. 20, §A.2 · p. 21

A.3 TEMPORAL DIFFERENCE FLOWS

  • The TD flow loss reconstructs current states and bootstraps from noisy target future states; choosing the behavioral policy as the target gives the SARSA variant used to learn occupancy models.The target future state is generated by applying the target vector field at the next state-action pair.Source: §A.3 · p. 22

B.1 THE CONSISTENCY ASSUMPTION ON INTENTIONS

  • The consistency assumption treats each latent intention as indexing a behavioral policy whose occupancy measure follows its own Bellman equations, motivating propagation of the same latent across transitions.The authors connect this propagation to using a TD loss to estimate the behavioral occupancy measure.Source: §B.1 · p. 22

B.2 DISTINCTIONS FROM META RL AND MULTI-TASK RL

InFOM differs from offline meta-RL and multi-task RL by inferring task structure from heterogeneous unlabeled data without task descriptions or task-specific reward supervision during pre-training. Its latent-variable objective predicts future states while regularizing the intention encoder, and can also be interpreted through an information bottleneck.

B.2 DISTINCTIONS FROM META RL AND MULTI-TASK RL

  • Unlike offline meta-RL, InFOM infers task structure from heterogeneous data without explicit task descriptions or reward-labeled data during pre-training; unlike multi-task RL, it uses neither task descriptions nor task-specific reward supervision.Instead, pre-training learns a generative multi-step transition model to facilitate downstream value estimation.Source: §B.2 · p. 22

C.1 THE EVIDENCE LOWER BOUND AND ITS CONNECTION WITH AN INFORMATION BOTTLENECK

  • The ELBO maximizes prediction of future states from inferred intentions while penalizing the encoder’s KL divergence from the prior, and admits an information-bottleneck interpretation.The stated bottleneck trades information about current states, actions, and future states against information about the next state-action pair retained in the intention.Source: §C.1 · pp. 23–24

C.2 INTUTIONS AND DISCUSSIONS ABOUT THE IMPLICIT GENERALIZED POLICY IMPROVEMENT

Implicit GPI aggregates intention-conditioned values, with expectile distillation approximating a maximum over intentions. By contrast, one-step policy improvement uses a behavioral value function and does not guarantee convergence to the optimal policy.

C.2 INTUTIONS AND DISCUSSIONS ABOUT THE IMPLICIT GENERALIZED POLICY IMPROVEMENT

  • As the sample count tends to infinity and the expectile tends to 1, expectile distillation converges to the maximum over intentions and Q approaches the greedy value function.The scalar Q summarizes average rewards at future states sampled from the flow occupancy models.Source: §C.2 · p. 24
  • When the flow occupancy model matches the behavioral policy’s conditional future-state distribution, the optimal Q corresponds to a greedy value function under that behavioral policy.Source: §C.2 · p. 24
  • Implicit GPI produces a policy no worse than the policies corresponding to each intention-conditioned Q function, and can converge to the downstream optimal policy if its task-specific intention was captured during pre-training.The optimality claim is conditional on capturing the task-specific intention.Source: §C.2 · p. 24

C.3 ONE-STEP POLICY IMPROVEMENT WITH FLOW OCCUPANCY MODELS

  • One-step policy improvement fits a behavioral Q function and maximizes it with a policy, but its behavioral value function does not guarantee convergence to the optimal policy.Its occupancy model is not conditioned on latent intentions and represents the average behavioral policy’s discounted state occupancy.Source: §C.3 · p. 24

D.1 TASKS AND DATASETS

The evaluation uses ExORL and OGBench control tasks, with reward-free pre-training followed by reward-labeled fine-tuning. Eight baselines span several pre-training strategies, and performance is evaluated with task-specific metrics and late fine-tuning checkpoints.

D.1 TASKS AND DATASETS

  • The benchmark suite contains 36 state-based and 4 image-based tasks from ExORL and OGBench, with manipulation tasks including long-horizon ordering and pixel-based reasoning.OGBench uses semi-sparse fine-tuning rewards; ExORL uses dense rewards except for sparse-reward jaco tasks.Source: §D.1 · pp. 24–25
  • ExORL uses 5M unlabeled transitions and 500K reward-labeled transitions, while OGBench uses 1M unlabeled and 500K reward-labeled transitions from the same data-collection policy.OGBench pre-training data comes from a non-Markovian expert with temporally correlated noise; ExORL pre-training data comes from exploration methods.Source: §D.1 · p. 25

D.2 BASELINES

  • The eight baselines cover offline actor-critic, representation learning, one-step model-based learning, contrastive occupancy prediction, and latent-intention methods.Image-based comparisons selectively include four baselines: ReBRAC, CRL + IS, DINO + ReBRAC, and FB + IQL.Source: §D.2 · pp. 25–26

D.3 EVALUATION PROTOCOLS

  • ExORL performance is measured by episode return and OGBench by success rate, with results averaged over the final fine-tuning checkpoints rather than the best checkpoint.For OGBench training, algorithms still use semi-sparse rewards; evaluation averages the last three specified checkpoints for each modality.Source: §D.3 · p. 26

D.4 IMPLEMENTATIONS AND HYPERPARAMETERS

The implementation uses shared hyperparameters where possible, while tuning settings by domain or task. InFOM combines an intention encoder and flow occupancy models for pre-training with reward prediction, critic, and policy learning during fine-tuning.

D.4 IMPLEMENTATIONS AND HYPERPARAMETERS

  • Hyperparameters are shared across methods where possible, tuned on selected tasks per domain, and reused for that domain’s tasks; image-task settings are tuned individually.Asynchronous policy training updates the policy four times less frequently than other models during fine-tuning.Source: §D.4 · pp. 26–27
  • InFOM’s intention encoder maps the next state-action pair to a Gaussian latent intention, while a flow-matching vector field predicts future states conditioned on the current state-action pair and intention.The latent dimension was swept over 64, 128, 256, and 512; 512 was sufficient for most ExORL tasks and 128 was generally sufficient for OGBench.Source: §D.4 · p. 27
  • Fine-tuning updates the intention encoder and flow occupancy models alongside a reward predictor, an unconditioned double-Q critic, and a Gaussian policy.The critic uses expectile distillation over 16 future states sampled from the flow occupancy models.Source: §D.4 · pp. 27–28
  • IQL and ReBRAC fine-tune initialized behavioral-cloning policies, while DINO + ReBRAC learns and fine-tunes state representations alongside the actor and critic.Source: §D.4 · p. 28
  • The baselines use distinct implementations: MBPO adds imaginary rollouts, contrastive occupancy methods use importance sampling to estimate Q, and FB/HILP infer a task representation before IQL fine-tuning.For FB and HILP, the task-specific representation is inferred from 10K fine-tuning transitions.Source: §D.4 · pp. 28–29

E ADDITIONAL VISUALIZATIONS OF LATENT INTENTIONS

F.1 EVALUATION ON ROBOTICS BENCHMARKS

The robotics and ablation evaluations compare InFOM with prior methods, test intention encoding and policy-learning components, and examine settings that affect performance. Results include gains on realistic robotics data and faster fine-tuning, alongside challenges from sparse rewards and discrete intentions.

F.1 EVALUATION ON ROBOTICS BENCHMARKS

  • 34%: InFOM outperforms the best baseline on the robotics evaluation when trained on embeddings from a fixed image encoder.The evaluation uses the RT-1 dataset for reward-free pre-training and reward-labeled fine-tuning, and tests the Google robot pick-coke-can task in simulation.Source: §F.1 · pp. 29, 31, Figure 7 · p. 31

F.2 VARIATIONAL INTENTION INFERENCE IS SIMPLE AND PERFORMANT

  • On 3 of 4 fine-tuning tasks, InFOM outperforms prior intention-encoding methods while using a simpler variational encoder.InFOM also achieves lower future-state prediction errors than HILP + FOM and FB + FOM on two ExORL tasks.Source: §F.2 · pp. 31–32, Figure 8 · p. 31, Figure 9 · p. 32

F.3 FLOW OCCUPANCY MODELS ENABLE FASTER POLICY LEARNING

  • InFOM converges faster than methods pre-training behavioral-cloning policies or self-supervised state representations, and faster than methods using one-step transition models or future-state classifiers.Its adaptation speed is similar to FB + IQL.Source: §F.3 · p. 32, Figure 10 · p. 32

F.4 LEARNING WITH DISCRETE INTENTIONS

  • Discrete intentions decrease InFOM performance by 11% on ExORL tasks and mean success rate by 78% on OGBench tasks.The results favor a continuous latent space in these experiments.Source: Figure 11 · p. 33, §F.4 · p. 33

F.5 FINE-TUNING WITH POSTERIOR INTENTIONS

  • Sampling latents from the posterior does not significantly change InFOM performance (+7%), supporting the use of prior samples to estimate Q_z.The authors interpret this as robustness to unseen latents for different (s, a) pairs.Source: Figure 12 · p. 33, §F.5 · p. 33

F.7 IMPORTANCE OF THE BEHAVIORAL CLONING REGULARIZATION

  • Behavioral-cloning regularization with α > 0 is a key component of InFOM’s policy loss.The ablation compares the full algorithm with a variant using α = 0.Source: Figure 14 · p. 34, §F.7 · p. 34

F.6 LEARNING WITH SPARSE REWARDS IS CHALLENGING

  • Using dense rewards results in a 3.6× smaller performance gap on Jaco tasks, suggesting that the original sparse reward function imposes learning challenges.The comparison fine-tunes ReBRAC with dense rewards and evaluates it in the original environments.Source: Figure 13 · p. 34, §F.6 · p. 34

F.8 DIVERSITY OF THE PRE-TRAINING DATASETS

  • Pre-training dataset size and diversity are positively associated: across two ExORL tasks, dataset diversity has an average correlation coefficient of 0.76 with size.Diversity is measured by the number of unique grid points covered after discretizing the state space.Source: §F.8 · p. 34, Figure 15 · p. 34

F.9 THE EFFECT OF DATASET SIZES

These ablations find that larger pre-training and fine-tuning datasets improve performance, while fine-tuning data quality and the number of sampled future states also matter. The chosen future-state count balances estimation accuracy, variance, and computation.

F.9 THE EFFECT OF DATASET SIZES

  • Larger pre-training datasets yield higher returns, and increasing fine-tuning dataset size consistently improves success rates on the selected OGBench tasks.The authors conjecture that diverse reward-free pre-training data reduces out-of-distribution intention sampling, while fine-tuning dataset size affects reward-prediction accuracy.Source: §F.9 · p. 35, Figure 16 · p. 35, Figure 17 · p. 35

F.10 FINE-TUNING ON SUBOPTIMAL DATASETS

  • 9%: Fine-tuning on a 0.2 reward-quantile dataset achieved 9% of original InFOM performance, while a 0.8-quantile dataset achieved 85%.The datasets were filtered by reward quantile from the ExORL dataset collected by RND.Source: §F.10 · p. 35, Figure 18 · p. 36

F.11 THE SUFFICIENT NUMBER OF FUTURE STATES IN THE Q ESTIMATION

  • Increasing the number of flow future states can improve Q-value accuracy and reduce variance, but raises computation; the experiments use N = 16 as a balance.A smaller N can yield errors from unrealistic future states, while a larger N mitigates high variance in some tasks.Source: §F.11 · pp. 35–36, Figure 19 · p. 36

F.12 ADDITIONAL HYPERPARAMETER ABLATIONS

The hyperparameter ablations identify sensitivity to latent dimension and domain-dependent tuning needs for several coefficients. The Euler-step effect saturates after a threshold.

F.12 ADDITIONAL HYPERPARAMETER ABLATIONS

  • InFOM is sensitive to latent intention dimension d, while the effect of Euler steps T saturates after T = 10.The authors suggest using a common T value across tasks.Source: §F.12 · p. 36, Figure 20 · p. 37
  • The expectile µ affects ExORL performance but has minor effects on OGBench; KL coefficient λ and behavioral-cloning coefficient α are crucial and require domain-specific tuning.The authors generally sweep hyperparameters on one task per domain, then use one set for every task in that domain.Source: §F.12 · p. 36, Figure 20 · p. 37

Related papers

Loading 2506.08902v4…