Intention-Conditioned Flow Occupancy Models

## Abstract

Abstract 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’s 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 36 state-based and 4 4 image-based benchmark tasks demonstrate that the proposed method achieves 1.8 × 1.8\times median improvement in returns and increases success rates by 36 % 36\% . Website: https://chongyi-zheng.github.io/infom Code: https://github.com/chongyi-zheng/infom

## 1 Introduction

Figure: Figure 1: InFOM is a latent variable model for pre-training and fine-tuning in reinforcement learning. (Left) The datasets are collected by users performing distinct tasks. (Center) We encode intentions by maximizing an evidence lower bound of data likelihood, (Right) enabling intention-aware future prediction using flow matching. See Sec. [4](https://arxiv.org/html/2506.08902v2#S4) for details.
Refer to caption: x1.png

Many of the recent celebrated successes of machine learning have been enabled by training large foundation models on vast datasets, and then adapting those models to downstream tasks. Examples include today’s chatbots (e.g., Gemini (Team et al., 2023) and ChatGPT (Achiam et al., 2023)) and generalist robotic systems (e.g., $\pi_{0}$ (Black et al., 2024) and Octo (Team et al., 2024)). This pre-training-fine-tuning paradigm has been wildly successful in fields ranging from computer vision to natural language processing (Devlin et al., 2019; Brown et al., 2020; Touvron et al., 2023; Zhai et al., 2023; Radford et al., 2021; He et al., 2022; Ouyang et al., 2022; Lu et al., 2019), yet harnessing it in the context of reinforcement learning (RL) remains an open problem.
What fundamentally makes the RL problem difficult is reasoning about time and intention—an effective RL agent must reason about the long-term effect of actions taken now, and must recognize that the data observed are often collected by distinct users performing multiple tasks.
However, current attempts to build foundation models for RL often neglect these two important bits, focusing on predicting the actions in the pre-training dataset instead (Team et al., 2024; O’Neill et al., 2024; Walke et al., 2023).

The closest attempts to building RL algorithms that capture temporal bits are those based on world models (Ding et al., 2024; Hafner et al., 2023; Mendonca et al., 2021) and those based on occupancy models (Janner et al., 2020; Blier et al., 2021; Zheng et al., 2024; Farebrother et al., 2025).(^1^11We will use “successor representations,” “occupancy measures,” and “occupancy models” interchangeably.)
World models can achieve great performance in sample efficiency (Janner et al., 2019) and generalize to diverse tasks (Hafner et al., 2023; Mendonca et al., 2021), although their capacity to perform long-horizon reasoning remains limited because of compounding errors (Talvitie, 2014; Janner et al., 2019; Lambert et al., 2022). Occupancy models (Dayan, 1993) and variants that enable scaling to high-dimensional tasks can also achieve great performance in predicting future events (Sikchi et al., 2024; Barreto et al., 2018; Zheng et al., 2024; 2025; Farebrother et al., 2025), but are typically hard to train and ignore user intentions. Recent advances in generative AI (e.g., flow-matching (Lipman et al., 2024; 2023; Liu et al., 2023) and diffusion (Ho et al., 2020; Song et al., 2021) models) enable modeling complex distributions taking various inputs, providing new tools for constructing occupancy models that depend on intentions.

In this paper, we propose a framework (Fig. [1](https://arxiv.org/html/2506.08902v2#S1.F1)) for pre-training in RL that simultaneously learns a probabilistic model to capture bits about time and intention. Building upon prior work on variational inference (Kingma & Welling, 2013; Alemi et al., 2017) and successor representations (Janner et al., 2020; Touati & Ollivier, 2021; Barreto et al., 2017; Zheng et al., 2024; Farebrother et al., 2025), we learn latent variable models of temporally distant future states, enabling intention-aware prediction. Building upon prior work on generative modeling, we use an expressive flow matching method (Farebrother et al., 2025) to train occupancy models, enabling highly flexible modeling of occupancy measures. We call the resulting algorithm intention-conditioned flow occupancy models (InFOM). Experiments on $36$ state-based and $4$ image-based benchmark tasks show that InFOM outperforms alternative methods for pre-training and fine-tuning by $1.8\times$ median improvement in returns and $36\%$ improvement in success rates. Additional experiments demonstrate that our latent variable model is capable of
inferring underlying user intentions (Sec. [5.2](https://arxiv.org/html/2506.08902v2#S5.SS2)) and enables efficient policy extraction (Sec. [5.3](https://arxiv.org/html/2506.08902v2#S5.SS3)).

## 2 Related work

Offline unsupervised RL.
The goal of offline unsupervised RL is
to pre-train policies, value functions, or models from an unlabeled (reward-free) dataset
to enable efficient learning of downstream tasks.
Prior work has proposed diverse offline unsupervised RL approaches based on
unsupervised skill learning (Touati & Ollivier, 2021; Frans et al., 2024; Park et al., 2024b; Kim et al., 2024; Hu et al., 2023),
offline goal-conditioned RL (Eysenbach et al., 2019; 2022; Valieva & Banerjee, 2024; Park et al., 2023a; Zheng et al., 2024; Park et al., 2025a),
and model-based RL (Mendonca et al., 2021; Mazzaglia et al., 2022).
Among these categories,
our method is conceptually related to offline unsupervised skill learning approaches (Park et al., 2024b; Touati et al., 2023), which also learns a model that predictions intention. However, our approach differs in that it does not learn multiple skills during pre-training.
Our work is complementary to a large body of prior work on using behavioral cloning for pretraining (O’Neill et al., 2024; Team et al., 2024), demonstrating that there are significant additional gains in performance that can be achieved by modeling intentions and occupancy measures simultaneously.

Unsupervised representation learning for RL.
Another way to leverage an unlabeled offline dataset is to learn representations
that facilitate subsequent downstream task learning.
Some works adapt existing representation learning techniques from computer vision,
such as contrastive learning (He et al., 2020; Parisi et al., 2022; Nair et al., 2023)
and masked autoencoding (He et al., 2022; Xiao et al., 2022). Others design specific methods for RL, including self-predictive representations (Schwarzer et al., 2020; Ni et al., 2024)
and temporal distance learning (Sermanet et al., 2018; Ma et al., 2023; Mazoure et al., 2023).
Those learned representations are typically used as inputs for policy and value networks. The key challenge with these representation learning methods is that it is often (Laskin et al., 2020), though not always (Zhang et al., 2021), unclear whether the learned representations will facilitate policy adaptation. In our experiments, we demonstrate that learning occupancy models enables faster policy learning.

RL with generative models.
Modern generative models have been widely adopted to solve RL problems.
Prior work has employed
autoregressive models (Vaswani et al., 2017),
iterative generative models
(e.g., denoising diffusion (Sohl-Dickstein et al., 2015; Ho et al., 2020)
and flow matching (Liu et al., 2023; Lipman et al., 2023; 2024)),
or autoencoders (Kingma & Welling, 2013)
to model
trajectories (Chen et al., 2021; Janner et al., 2021; 2022; Ajay et al., 2023),
environment dynamics (Ding et al., 2024; Alonso et al., 2024),
skills (Ajay et al., 2021; Pertsch et al., 2021; Frans et al., 2024),
and policies (Wang et al., 2023; Hansen-Estruch et al., 2023; Park et al., 2025b).
We employ a state-of-the-art flow-matching objective (Farebrother et al., 2025) to model discounted state occupancy measures.

Successor representations and successor features.
Prior work has used successor representations (Dayan, 1993) and successor features (Barreto et al., 2017)
for transfer learning (Barreto et al., 2017; 2018; Borsa et al., 2018; Nemecek & Parr, 2021; Kim et al., 2022), unsupervised RL (Machado et al., 2017; Hansen et al., 2019; Ghosh et al., 2023; Touati et al., 2023; Park et al., 2024b; 2023b; Chen et al., 2023; Zheng et al., 2025; Jain et al., 2023; Zhu et al., 2024), and goal-conditioned RL (Eysenbach et al., 2020; 2022; Zheng et al., 2024).
Our method is closely related to prior methods that learn successor representations
with generative models (Janner et al., 2020; Thakoor et al., 2022; Tomar et al., 2024; Farebrother et al., 2025).
In particular, the most closely related to ours is the prior work by Farebrother et al. (2025),
which also uses flow-matching to model the occupancy measures
and partly employs the generalized policy improvement (GPI) for policy extraction.
Unlike Farebrother et al. (2025),
which uses forward-backward representations to capture behavioral intentions and perform GPI over a finite set of intentions,
our method employs a latent variable model to learn intentions (Sec. [4.2](https://arxiv.org/html/2506.08902v2#S4.SS2))
and uses an expectile loss to perform implicit GPI (Sec. [4.4](https://arxiv.org/html/2506.08902v2#S4.SS4)).
We empirically show that these choices lead to higher returns and success rates (Sec. [5.1](https://arxiv.org/html/2506.08902v2#S5.SS1),
Sec. [5.3](https://arxiv.org/html/2506.08902v2#S5.SS3)).

## 3 Preliminaries

We consider a Markov decision process (MDP) (Sutton et al., 1998) defined by
a state space ${\mathcal{S}}$, an action space ${\mathcal{A}}$,
an initial state distribution $\rho\in\Delta({\mathcal{S}})$, a reward function $r:{\mathcal{S}}\to\mathbb{R}$, a discount factor $\gamma\in[0,1)$, and a transition distribution $p:{\mathcal{S}}\times{\mathcal{A}}\to\Delta({\mathcal{S}})$, where $\Delta(\cdot)$ denotes the set of all possible probability distributions over a space. We will use $h$ to denote a time step in the MDP and assume the reward function only depends on the state at the current time step $r_{h}\triangleq r(s_{h})$ without loss of generality (Tomar et al., 2024; Frans et al., 2024; Thakoor et al., 2022). In Appendix [A.1](https://arxiv.org/html/2506.08902v2#A1.SS1), we briefly review the definition of value functions and the actor-critic framework in RL.

Occupancy measures. Alternatively, one can summarize the stochasticity over trajectories into the *discounted state occupancy measure* (Dayan, 1993; Eysenbach et al., 2022; Janner et al., 2020; Touati & Ollivier, 2021; Zheng et al., 2024; Myers et al., 2024; Blier et al., 2021) that quantifies the discounted visitation frequency of different states under the policy $\pi$.
Prior work (Dayan, 1993; Janner et al., 2020; Touati & Ollivier, 2021; Zheng et al., 2024) has shown that the discounted state occupancy measure follows a Bellman equation backing up the probability density at the current time step and the future time steps:

$$ $\displaystyle p^{\pi}_{\gamma}(s_{f}\mid s,a)=(1-\gamma)\delta_{s}(s_{f})+\gamma\mathbb{E}_{\begin{subarray}{c}s^{\prime}\sim p(s^{\prime}\mid s,a),\\ a^{\prime}\sim\pi(a^{\prime}\mid s^{\prime})\end{subarray}}\left[p^{\pi}_{\gamma}(s_{f}\mid s^{\prime},a^{\prime})\right],$ (1) $$

where $\delta_{s}(\cdot)$ denotes the Dirac delta measure centered at $s$.(^2^22The recursive relationship in Eq. [1](https://arxiv.org/html/2506.08902v2#S3.E1) starts from the current time step (Eysenbach et al., 2022; Touati & Ollivier, 2021) instead of the next time step as in some prior approaches (Janner et al., 2020; Zheng et al., 2024; Thakoor et al., 2022).) The discounted state occupancy measure allows us to rewrite the Q-function as a linear function of rewards (Barreto et al., 2017; Touati & Ollivier, 2021; Zheng et al., 2024; Sikchi et al., 2024):

$$ $\displaystyle Q^{\pi}(s,a)=\frac{1}{1-\gamma}\mathbb{E}_{s_{f}\sim p^{\pi}_{\gamma}(s_{f}\mid s,a)}\left[r(s_{f})\right].$ (2) $$

The alternative (dual (Sikchi et al., 2024)) definition of Q-function (Eq. [2](https://arxiv.org/html/2506.08902v2#S3.E2)) allows us to cast the policy evaluation step as first learning a generative model $p_{\gamma}(s_{f}\mid s,a)$ to simulate the discounted state occupancy measure of $\pi^{k}$ and then regressing the estimator $Q$ towards the average reward at states sampled from $p_{\gamma}$ (Toussaint & Storkey, 2006; Tomar et al., 2024; Thakoor et al., 2022; Zheng et al., 2024). See Sec. [4.4](https://arxiv.org/html/2506.08902v2#S4.SS4) for detailed formulation.

Flow matching and TD flows. Flow matching (Lipman et al., 2023; 2024; Liu et al., 2023; Albergo & Vanden-Eijnden, 2023) refers to a family of generative models based on ordinary differential equations (ODEs), which are close cousins of denoising diffusion models (Sohl-Dickstein et al., 2015; Song et al., 2021; Ho et al., 2020), which instead solve a stochastic differential equation (SDE).
The deterministic nature of ODEs equips flow-matching methods with simpler learning objectives and faster inference speed than denoising diffusion models (Lipman et al., 2023; 2024; Park et al., 2025b). In Appendix [A.2](https://arxiv.org/html/2506.08902v2#A1.SS2), we discuss the problem setting and the standard learning objective for flow matching.

In the context of RL, prior work has used flow matching to estimate the discounted state occupancy measure (Farebrother et al., 2025)
by incorporating the Bellman equation (Eq. [1](https://arxiv.org/html/2506.08902v2#S3.E1)) into the conditional flow matching loss (Eq. [10](https://arxiv.org/html/2506.08902v2#A1.E10)), resulting in a temporal difference flow matching procedure (TD flows) (Farebrother et al., 2025). In Appendix [A.3](https://arxiv.org/html/2506.08902v2#A1.SS3), we discuss the detailed formulations of the TD flow objective for a target policy $\pi$. Choosing the target policy $\pi$ to be the same as the behavioral policy $\beta$, we obtain a SARSA (Rummery & Niranjan, 1994) variant of the loss optimizing the SARSA flows. We will use the SARSA variant of the TD flow objective to learn our generative occupancy models in Sec. [4.3](https://arxiv.org/html/2506.08902v2#S4.SS3).

## 4 Intention-conditioned flow occupancy models

In this section, we will introduce our method for pre-training and fine-tuning in RL. After formalizing the problem setting, we will dive into the latent variable model for pre-training an intention encoder and flow occupancy model. After pre-training the occupancy models, our method will extract polices for solving different tasks by invoking a generalized policy improvement procedure (Barreto et al., 2017). We refer to our method as intention-conditioned flow occupancy models (InFOM).

### 4.1 Problem setting

We consider learning with purely *offline* datasets, where an unlabeled (reward-free) dataset of transitions $D=\{(s,a,s^{\prime},a^{\prime})\}$ collected by the behavioral policy $\beta$ is provided for pre-training and a reward-labeled dataset $D_{\text{reward}}=\{(s,a,r)\}$ collected by some other policy $\tilde{\beta}$ on a downstream task is used for fine-tuning. Importantly, the behavioral policy $\beta$ used to collect $D$ can consist of a mixture of policies used by different users to complete distinct tasks. We will call this heterogeneous structure of the unlabeled datasets “intentions,” which are latent vectors $z$s in some latent space ${\mathcal{Z}}$. In practice, these intentions can refer to desired goal images or language instructions that index the behavioral policy $\beta=\{\beta(\cdot\mid\cdot,z):z\in{\mathcal{Z}}\}$. Because these latent intentions are *unobserved* to the pre-training algorithm, we want to infer them as a latent random variable $Z$ from the offline dataset, similar to prior work (Hausman et al., 2017; Li et al., 2017; Henderson et al., 2017). In Appendix [B.2](https://arxiv.org/html/2506.08902v2#A2.SS2), we include discussions distinguishing our problem setting from meta RL and multi-task RL problems.

During pre-training, our method exploits the heterogeneous structure of the unlabeled dataset and extracts actionable
information by *(1)* inferring intentions of the data collection policy and *(2)* learning occupancy models to predict long-horizon future states based on those intentions (Sec. [4.2](https://arxiv.org/html/2506.08902v2#S4.SS2) & [4.3](https://arxiv.org/html/2506.08902v2#S4.SS3)). During fine-tuning, we first recover a set of intention-conditioned Q functions by regressing towards average rewards at future states generated by the occupancy models, and then extract a policy to maximize task-specific discounted cumulative returns (Sec. [4.4](https://arxiv.org/html/2506.08902v2#S4.SS4)). Our method builds upon an assumption regarding the consistency of latent intentions.

###### Assumption 1 (Consistency) .

The unlabeled dataset $D$ for pre-training is obtained by executing a behavioral policy following a mixture of unknown intentions $z\in{\mathcal{Z}}$.
We assume that
consecutive transitions $(s,a)$ and $(s^{\prime},a^{\prime})$ share the same intention.

The consistency of intentions across transitions enables both intention inference using two sets of transitions and dynamic programming over trajectory segments. See Appendix [B.1](https://arxiv.org/html/2506.08902v2#A2.SS1) for justifications of this assumption.

###### Assumption 1 (Consistency) .

### 4.2 Variational intention inference

The goal of our pre-training framework is to learn a latent variable model that captures both long-horizon temporal information and unknown user intentions in the unlabeled datasets.

This part of our method aims to infer the intention $z$ based on consecutive transitions $(s,a,s^{\prime},a^{\prime})$ using the encoder $p_{e}(z\mid s^{\prime},a^{\prime})$ and predict the occupancy measures of a future state $s_{f}$ using the occupancy models $q_{d}(s_{f}\mid s,a,z)$.
We want to maximize the likelihood of observing a future state $s_{f}$ starting from a state-action pair $(s,a)$ (amortized variational inference (Kingma & Welling, 2013; Margossian & Blei, 2024)), both sampled from the unlabeled dataset $D$ following the joint behavioral distribution $p^{\beta}(s,a,s_{f})=p^{\beta}(s,a)p^{\beta}(s_{f}\mid s,a)$:

$$ $\displaystyle\max_{q_{d}}\,\mathbb{E}_{p^{\beta}(s,a,s_{f})}\left[\log q_{d}(s_{f}\mid s,a)\right]$ $\displaystyle\geq\max_{p_{e},q_{d}}\mathbb{E}_{p^{\beta}(s,a,s_{f},s^{\prime},a^{\prime})}\left[\mathbb{E}_{p_{e}(z\mid s^{\prime},a^{\prime})}\left[\log q_{d}(s_{f}\mid s,a,z)\right]-\lambda D_{\mathrm{KL}}(p_{e}(z\mid s^{\prime},a^{\prime})\parallel p(z))\right],$ (3) $$

where $p(z)={\mathcal{N}}(0,I)$ denotes an uninformative standard Gaussian prior over intentions and $\lambda$ denotes the coefficient that controls the strength of the KL divergence regularization. We defer the full derivation of the evidence lower bound (ELBO) $L(p_{e},q_{d})$ to Appendix [C.1](https://arxiv.org/html/2506.08902v2#A3.SS1). Inferring the intention $z$ from the next transition $(s^{\prime},a^{\prime})$ follows from our consistency assumption (Assump. [1](https://arxiv.org/html/2506.08902v2#Thmtheorem1)), and is important for avoiding overfitting (Frans et al., 2024).
Importantly, $p_{e}$ and $q_{d}$ are optimized *jointly* with this objective.
One way of understanding this ELBO is as maximizing an information bottleneck with the chain of random variables $(S^{\prime},A^{\prime})\to Z\to(S,A,S_{f})$.
See Appendix [C.1](https://arxiv.org/html/2506.08902v2#A3.SS1) for the connection.

We use flow matching
to reconstruct the discounted state occupancy measure rather than maximizing the likelihood directly, resulting in minimizing a surrogate objective:

$$ $\displaystyle\min_{p_{e},q_{d}}{\mathcal{L}}_{\text{Flow}}(q_{d},p_{e})+\lambda\mathbb{E}_{p^{\beta}(s^{\prime},a^{\prime})}\left[D_{\mathrm{KL}}(p_{e}(z\mid s^{\prime},a^{\prime})\parallel p(z))\right].$ (4) $$

We use ${\mathcal{L}}_{\text{Flow}}$ to denote a placeholder for the flow matching loss and will instantiate this loss for the flow occupancy models $q_{d}$ next.

### 4.3 Predicting the future via SARSA flows

We now present the objective used to learn the flow occupancy models, where we first introduce some motivations and desiderata and then specify the actual loss. Given an unlabeled dataset $D$ and an intention encoder $p_{e}(z\mid s^{\prime},a^{\prime})$, the goal is to learn a *generative* occupancy model $q_{d}(s_{f}\mid s,a,z)$ that approximates the discounted state occupancy measure of the behavioral policy conditioned on different intentions, i.e., $q_{d}(s_{f}\mid s,a,z)\approx p^{\beta}(s_{f}\mid s,a,z)$. We will use $v_{d}:[0,1]\times{\mathcal{S}}\times{\mathcal{S}}\times{\mathcal{A}}\times{\mathcal{Z}}\to{\mathcal{S}}$ to denote the time-dependent vector field that corresponds to $q_{d}$. There are two desired properties of the learned occupancy models: *(1)* distributing the peak probability density to multiple $s_{f}$, i.e., modeling multimodal structure, and *(2)* stitching together trajectory segments that share some transitions in the dataset, i.e., enabling combinatorial generalization.
The first property motivates us to use an expressive flow-matching model (Lipman et al., 2024), while the second property motivates us to learn those occupancy models using temporal difference approaches (Janner et al., 2020; Tomar et al., 2024; Farebrother et al., 2025).
Prior work (Farebrother et al., 2025) has derived the TD version of the regular (Monte Carlo) flow matching loss (Eq. [10](https://arxiv.org/html/2506.08902v2#A1.E10)) that incorporates the Bellman backup into the flow matching procedure, showing the superiority in sample efficiency and the capability of dynamic programming. We will adopt the same idea and use the SARSA variant of the TD flow loss (Eq. [11](https://arxiv.org/html/2506.08902v2#A1.E11)) to learn our intention-conditioned flow occupancy models:

$$ $\displaystyle{\mathcal{L}}_{\text{SARSA flow}}(v_{d},p_{e})$ $\displaystyle=(1-\gamma){\mathcal{L}}_{\text{SARSA current flow}}(v_{d},p_{e})+\gamma{\mathcal{L}}_{\text{SARSA future flow}}(v_{d},p_{e}),$ (5) $\displaystyle{\mathcal{L}}_{\text{SARSA current flow}}(v_{d},p_{e})$ $\displaystyle=\mathbb{E}_{\begin{subarray}{c}(s,a,s^{\prime},a^{\prime})\sim p^{\beta}(s,a,s^{\prime},a^{\prime}),\\ z\sim p_{e}(z\mid s^{\prime},a^{\prime}),\\ t\sim\textsc{Unif}([0,1]),\epsilon\sim{\mathcal{N}}(0,I)\end{subarray}}\left[\lVert v(t,s^{t},s,a,z)-(s-\epsilon)\rVert_{2}^{2}\right],$ $\displaystyle{\mathcal{L}}_{\text{SARSA future flow}}(v_{d},p_{e})$ $\displaystyle=\mathbb{E}_{\begin{subarray}{c}(s,a,s^{\prime},a^{\prime})\sim p^{\beta}(s,a,s^{\prime},a^{\prime}),\\ z\sim p_{e}(z\mid s^{\prime},a^{\prime}),\\ t\sim\textsc{Unif}([0,1]),\epsilon\sim{\mathcal{N}}(0,I)\end{subarray}}\left[\lVert v_{d}(t,\bar{s}_{f}^{t},s,a,z)-\bar{v}_{d}(t,\bar{s}_{f}^{t},s^{\prime},a^{\prime},z)\rVert_{2}^{2}\right].$ $$

Importantly, incorporating the information from latent intentions into the flow occupancy models allows us to *(1)* use the simpler and more stable SARSA bootstrap instead of the Q-learning style bootstrap (Eq. [11](https://arxiv.org/html/2506.08902v2#A1.E11)) on large datasets, *(2)* generalize over latent intentions, avoiding counterfactual errors. Sec. [5.2](https://arxiv.org/html/2506.08902v2#S5.SS2) visualizes the latent intentions, and Appendix [F.2](https://arxiv.org/html/2506.08902v2#A6.SS2) contains additional experiments.

### 4.4 Generative value estimation and implicit generalized policy improvement

We next discuss the fine-tuning process in our algorithm. Our fine-tuning method builds on the dual perspective of value estimation introduced in the preliminaries (Eq. [2](https://arxiv.org/html/2506.08902v2#S3.E2)). We first estimate a *set* of intention-conditioned Q functions using regression and then use those intention-conditioned Q functions to extract a policy, utilizing generalized policy improvement (GPI) (Barreto et al., 2017). The key idea of GPI is that, in addition to taking the maximum over the actions, we can also take the maximum over the intentions. In our setting, the number of intentions is infinite—one for every choice of continuous $z$. Thus, taking the maximum over the intentions is both nontrivial and susceptible to instability (Sec. [5.3](https://arxiv.org/html/2506.08902v2#S5.SS3)).
We address this issue by replacing the greedy “max” with an upper expectile loss, resulting in an implicit generalized policy improvement procedure.

Generative value estimation. Given a reward-labeled dataset $D_{\text{reward}}$ and the pre-trained flow occupancy models $q_{d}$, we can estimate intention-conditioned Q values for a downstream task. Specifically, for a fixed latent intention $z\in{\mathcal{Z}}$, we first sample a set of $N$ future states from the flow occupancy models, $s_{f}^{(1)},\cdots,s_{f}^{(N)}:s_{f}^{(i)}\sim q_{d}(s_{f}\mid s,a,z)$, and then constructs a Monte Carlo (MC) estimation of the Q function using those generative samples:(^3^33We omit the dependency of $Q_{z}$ on $s_{f}^{(1)},\cdots,s_{f}^{(N)}$ to simplify notations.)

$$ $\displaystyle Q_{z}(s,a)=\frac{1}{(1-\gamma)N}\sum_{i=1}^{N}r\left(s_{f}^{(i)}\right),\,s_{f}^{(i)}\sim q_{d}(s_{f}\mid s,a,z),$ (6) $$

where $r(\cdot)$ is the reward function or a learned reward predictor. Importantly, the choice of the number of future states $N$ affects the accuracy and variance of our Q estimate. Ablation experiments in Appendix [F.7](https://arxiv.org/html/2506.08902v2#A6.SS7) indicate that $N=16$ works effectively in our experiments.

Implicit generalized policy improvement.
We can then use those MC estimation of Q functions to learn a policy by invoking the generalized policy improvement.
The naive GPI requires sampling a finite set of latent intentions from the prior distribution $p(z)$,(^4^44We choose to sample from the prior $p(z)$ because of the KL divergence in $L(p_{e},q_{d})$, resembling drawing random samples from a variational auto-encoder (Kingma & Welling, 2013).) $z^{(1)},\cdots,z^{(M)}:\>z^{(j)}\sim p(z)$ and greedily choose one $Q_{z}$ to update the policy:

$$ $\displaystyle\operatorname*{arg\,max}_{\pi}\,\mathbb{E}_{\begin{subarray}{c}s\sim p^{\tilde{\beta}}(s),\>a\sim\pi(a\mid s)\\ z^{(1)},\cdots,z^{(M)}:\>z^{(j)}\sim p(z)\end{subarray}}\left[\max_{z^{(j)}}Q_{z^{(j)}}(s,a)\right].$ $$

Despite its simplicity, the naive GPI suffers from two main disadvantages.
First, using the maximum Q over a finite set of latent intentions to approximate the maximum Q over an infinite number of intentions results in local optima. Second, when we take gradients of this objective with respect to the policy, the chain rule gives one term involving $\nabla_{a}q_{d}(s_{f}|s,a,z)$. Thus, computing the gradients requires differentiating through the ODE solver (backpropagating through time (Park et al., 2025b)), which is unstable.
We address these challenges by learning an explicit scalar Q function to distill the MC estimation of intention-conditioned Q functions. This approach is appealing because gradients of the Q function no longer backpropagate through the ODE solver. We also replace the “max” over a finite set of intention-conditioned Q functions with an upper expectile loss $L_{2}^{\mu}$ (Kostrikov et al., 2022), resulting in the following critic loss

$$ $\displaystyle{\mathcal{L}}(Q)=\mathbb{E}_{(s,a)\sim p^{\tilde{\beta}}(s,a),\>z\sim p(z)}\left[L_{2}^{\mu}\left(Q_{z}(s,a)-Q(s,a)\right)\right],$ (7) $$

where $L_{2}^{\mu}(x)=\lvert\mu-\mathbbm{1}(x<0)\rvert x^{2}$ and $\mu\in[0.5,1)$. After distilling the intention-conditioned Q functions into a single function, we can extract the policy by selecting actions to maximize $Q$ with a behavioral cloning regularization (Fujimoto & Gu, 2021) using the actor loss

$$ $\displaystyle{\mathcal{L}}(\pi)=-\mathbb{E}_{(s,a)\sim p^{\tilde{\beta}}(s,a),a^{\pi}\sim\pi(a^{\pi}\mid s)}[Q(s,a^{\pi})+\alpha\log\pi(a\mid s)],$ (8) $$

where $\alpha$ controls the regularization strength. We use the behavioral cloning regularization to both reduce errors from sampling out-of-distribution (OOD) actions (Kumar et al., 2020; Fujimoto & Gu, 2021) and mitigate error propagations through overestimated $Q_{z}$ values. Alation experiments in Appendix [F.5](https://arxiv.org/html/2506.08902v2#A6.SS5) & [F.8](https://arxiv.org/html/2506.08902v2#A6.SS8) show that this behavioral cloning regularization is important for extracting the policy.
Taken together, we call the expectile Q distillation step (Eq. [7](https://arxiv.org/html/2506.08902v2#S4.E7)) and the policy optimization step (Eq. [8](https://arxiv.org/html/2506.08902v2#S4.E8)) *implicit generalized policy improvement (implicit GPI)*. In Appendix [C.2](https://arxiv.org/html/2506.08902v2#A3.SS2), we discuss the intuition and theoretical soundness of the implicit GPI.

Algorithm summary. We use neural networks to parameterize the intention encoder $p_{\phi}$, the vector field of the occupancy models $v_{\theta}$, the reward predictor $r_{\eta}$, the critic $Q_{\psi}$, and the policy $\pi_{\omega}$. We consider two stages: pre-training and fine-tuning. In Alg. [1](https://arxiv.org/html/2506.08902v2#alg1), we summarize the pre-training process of InFOM. InFOM pre-trains *(1)* the vector field $v_{\theta}$ using the SARSA flow loss (Eq. [5](https://arxiv.org/html/2506.08902v2#S4.E5)) and *(2)* the intention encoder $p_{\phi}$ using the ELBO (Eq. [3](https://arxiv.org/html/2506.08902v2#S4.E3)). Alg. [2](https://arxiv.org/html/2506.08902v2#alg2) shows the pseudocode of InFOM for fine-tuning. InFOM mainly learns *(1)* the reward predictor $r_{\eta}$ via simple regression, *(2)* the critic $Q_{\psi}$ using expectile distillation (Eq. [7](https://arxiv.org/html/2506.08902v2#S4.E7)), and *(3)* the policy $\pi_{\omega}$ by conservatively maximizing the $Q$ (Eq. [8](https://arxiv.org/html/2506.08902v2#S4.E8)). The open-source implementation is available online.(^5^55[https://github.com/chongyi-zheng/infom](https://github.com/chongyi-zheng/infom))

## 5 Experiments

Figure: Figure 3: Domains for evaluation. (Left) ExORL domains (16 state-based tasks). (Right) OGBench domains (20 state-based tasks and 4 image-based tasks).
Refer to caption: figures/domains/cheetah.png

Our experiments start with comparing InFOM to prior methods that first pre-train on reward-free datasets and then fine-tune on reward-labeled datasets, measuring the performance on downstream tasks. We then study the two main components of our method: the variational intention encoder and the implicit GPI policy extraction strategy. Visualizations of the latent intention inferred by our variational intention encoder show alignment with the underlying ground-truth intentions. Our ablation experiments reveal the effect of the implicit GPI policy extraction strategy.
We also include additional experiments showing InFOM enables faster policy learning during fine-tuning in Appendix [F.3](https://arxiv.org/html/2506.08902v2#A6.SS3). Our algorithm is robust to various choices of hyperparameters (Appendix [F.8](https://arxiv.org/html/2506.08902v2#A6.SS8)). Following prior work (Park et al., 2025b), all experiments report means and standard deviations across 8 random seeds for state-based tasks and 4 random seeds for image-based tasks.

Figure: (a) 16 state-based ExORL tasks from Yarats et al. (2022). We average over 4 tasks for each domain.
Refer to caption: x2.png

### 5.1 Comparing to prior pre-training and fine-tuning methods

Our experiments study whether the proposed method (InFOM), which captures actionable information conditioned on user intentions from unlabeled datasets, enables effective pre-training and fine-tuning. We select 36 state-based and 4 image-based tasks across diverse robotic navigation and manipulation domains and compare against 8 baselines. The models pre-trained by those methods include behavioral cloning policies (Kostrikov et al., 2022; Tarasov et al., 2023a), transition models (Janner et al., 2019), representations (Caron et al., 2021), discriminative classifiers that predict occupancy measures (Eysenbach et al., 2022; Zheng et al., 2024), and latent skills (Touati & Ollivier, 2021; Park et al., 2024b). We defer the detailed discussions about benchmarks and datasets to Appendix [D.1](https://arxiv.org/html/2506.08902v2#A4.SS1) and the rationale for choosing different baselines to Appendix [D.2](https://arxiv.org/html/2506.08902v2#A4.SS2). Whenever possible, we use the same hyperparameters for all methods (Table [1](https://arxiv.org/html/2506.08902v2#A4.T1)). See Appendix [D.3](https://arxiv.org/html/2506.08902v2#A4.SS3) for details of the evaluation protocol and Appendix [D.4](https://arxiv.org/html/2506.08902v2#A4.SS4) for implementations and hyperparameters of each method.

We report results in Fig. [4](https://arxiv.org/html/2506.08902v2#S5.F4), aggregating over four tasks in each domain of ExORL and five tasks in each domain of OGBench, and present the full results in Table [4](https://arxiv.org/html/2506.08902v2#A4.T4). These results show that InFOM matches or surpasses all baselines on six out of eight domains. On ExORL benchmarks, all methods perform similarly on the two easier domains (cheetah and quadruped), while InFOM can obtain $20\times$ improvement on jaco, where baselines only make trivial progress (Table [4](https://arxiv.org/html/2506.08902v2#A4.T4)). We suspect the outsized improvement on the jaco task is because of the high-dimensional state space (twice that of the other ExORL tasks (Yarats et al., 2022)) and because it has sparse rewards; Appendix Fig. [11](https://arxiv.org/html/2506.08902v2#A6.F11) supports this hypothesis by showing that the ReBRAC baseline achieves significantly higher returns when using dense rewards.
On those more challenging state-based manipulation tasks from OGBench, we find a marked difference between baselines and InFOM; our method achieves $36\%$ higher success rate over the best performing baseline. In addition, InFOM is able to outperform the best baseline by $31\%$ using RGB images as input directly (visual tasks). We hypothesize that the baselines fail to solve these more challenging tasks because of the semi-sparse reward functions. In contrast, our method can explore different regions of the state space using the different intentions, thereby addressing the challenge of reward sparsity.
We conjecture that the variance of InFOM across seeds in some experiments (e.g., cheetah, cube single, and puzzle 4x4) reflects stochasticity in the MC Q estimates (Eq. [6](https://arxiv.org/html/2506.08902v2#S4.E6)), which might be mitigated by increasing the number of sampled future states (See Appendix [F.7](https://arxiv.org/html/2506.08902v2#A6.SS7)). In Appendix [F.1](https://arxiv.org/html/2506.08902v2#A6.SS1), we compare InFOM against selective baselines on real robotics datasets, showing $34\%$ improvement.

### 5.2 Visualizing latent intentions

Figure: Figure 5: Visualization of latent intentions. (Top) The optimal policy picks up the blue block from the left and places it on the right. (Bottom) Using t-SNE (Maaten & Hinton, 2008), we visualize the latent intentions inferred by the variational intention encoder in InFOM, comparing against latent representations inferred by HILP and FB for learning FOMs. The predictions from InFOM align with the underlying intentions. See Sec. [5.2](https://arxiv.org/html/2506.08902v2#S5.SS2) for details and Appendix [E](https://arxiv.org/html/2506.08902v2#A5) for more visualizations.
Refer to caption: x4.png

Our next experiment studies the intention encoder in our algorithm. To investigate whether the proposed method discovers distinct user intentions from an unlabeled dataset, we visualize latent intentions inferred by our variational intention encoder. We include comparisons against two alternative intention encoding mechanisms proposed by prior methods. Specifically, we consider replacing the variational intention encoder with either *(1)* a set of Hilbert representations (Park et al., 2024b) or *(2)* a set of forward-backward representations (Touati & Ollivier, 2021), and then pre-training the flow occupancy models (FOM) conditioned on these two sets of representations. We call these two variants HILP + FOM and FB + FOM. Note that FB + FOM is equivalent to TD flows with GPI in Farebrother et al. (2025). Using t-SNE (Maaten & Hinton, 2008), we visualize latent intentions predicted by these three methods on cube double task 1 from the OGBench benchmarks.

Fig. [5](https://arxiv.org/html/2506.08902v2#S5.F5) shows the optimal trajectory, where the manipulator picks the blue block from the left and then places it on the right, and the visualizations. The 2D t-SNE visualizations indicate that both FB + FOM and HILP + FOM infer mixed latent intentions for “pick” and “place” behaviors, while InFOM predicts a sequence of latent intentions with clear clustering. This result suggests that InFOM is capable of inferring latent intentions that align with the underlying ground-truth intentions. See Appendix [E](https://arxiv.org/html/2506.08902v2#A5) for more visualizations.

### 5.3 Importance of the implicit generalized policy improvement

Figure: Figure 6: Comparison to alternative policy extraction strategies. We compare InFOM to alternative policy extraction strategies based on the standard generalized policy improvement or one-step policy improvement. Our method is $44\%$ more performant with $8\times$ smaller variance than the variant using the standard GPI. See Sec. [5.3](https://arxiv.org/html/2506.08902v2#S5.SS3) for details.
Refer to caption: x6.png

Our final experiments study different approaches for policy optimization. We hypothesize that our proposed method is more efficient and robust than other policy extraction strategies. To test this hypothesis, we conduct ablation experiments on one task in the ExORL benchmarks (quadruped jump) and another task taken from the OGBench benchmarks (scene task 1), again following the evaluation protocols in Appendix [D.3](https://arxiv.org/html/2506.08902v2#A4.SS3). We compare two alternative policy learning approaches in the fine-tuning phase. First, we ablate the effect of the upper expectile loss by comparing against the standard GPI, which maximizes Q functions over a finite set of intentions $\{z^{(1)},\cdots,z^{(M)}\}$. We choose $M=32$ latent intentions to balance between performance and compute budget, and call this variant InFOM + GPI. Second, we ablate the effect of the variational intention encoder by removing the intention dependency in the flow occupancy models and extracting the policy via one-step policy improvement (PI) (Wang et al., 2018; Brandfonbrener et al., 2021; Peters & Schaal, 2007; Peters et al., 2010). We call this method FOM + one-step PI and defer the detailed formulation to Appendix [C.3](https://arxiv.org/html/2506.08902v2#A3.SS3).

As shown in Fig. [6](https://arxiv.org/html/2506.08902v2#S5.F6), InFOM achieves significantly higher returns and success rates than its variant based on one-step policy improvement, suggesting the importance of inferring user intentions. Compared with its GPI counterpart, our method is $44\%$ more performant with $8\times$ smaller variance (the error bar indicates one standard deviation), demonstrating that the implicit GPI indeed performs a relaxed maximization over intentions while maintaining robustness.

Additional experiments. In Appendix [F.3](https://arxiv.org/html/2506.08902v2#A6.SS3), we include additional ablations showing that InFOM enables faster policy learning. The dataset size ablations in Appendix [F.6](https://arxiv.org/html/2506.08902v2#A6.SS6) show that using sufficient pre-training and fine-tuning data is important. Our hyperparameter ablations can be found in Appendix [F.8](https://arxiv.org/html/2506.08902v2#A6.SS8).

## 6 Conclusion

In this work, we presented InFOM, a method that captures diverse intentions and their long-term behaviors from an unstructured dataset, leveraging the expressivity of flow models. We empirically showed that the intentions captured in flow occupancy models enable effective and efficient fine-tuning, outperforming prior unsupervised pre-training approaches on diverse state- and image-based domains.

Limitations.
One limitation of InFOM is that our reduction from trajectories to consecutive state-action pairs might not always accurately capture the original intentions in the trajectories. While we empirically showed that this simple approach is sufficient to achieve strong performance on our benchmark tasks, it can be further improved with alternative trajectory encoding techniques and data collection strategies, which we leave for future work.

## Reproducibility statement

We implement InFOM and all baselines in the same codebase using JAX (Bradbury et al., 2018). Our implementations build on top of OGBench’s and FQL’s implementations (Park et al., 2025a; b). We include the common hyperparameters for all the methods in Appendix Table [1](https://arxiv.org/html/2506.08902v2#A4.T1), the hyperparameters for InFOM in Appendix Table [2](https://arxiv.org/html/2506.08902v2#A4.T2) and Appendix Table [3](https://arxiv.org/html/2506.08902v2#A4.T3), and the hyperparameters for baselines in Appendix Table [3](https://arxiv.org/html/2506.08902v2#A4.T3). All the experiments were run on a single NVIDIA A6000 GPU and can be finished in 4 hours for state-based tasks and 12 hours for image-based tasks. We provide open-source implementations of InFOM and all baselines at [https://github.com/chongyi-zheng/infom](https://github.com/chongyi-zheng/infom).

## Acknowledgments

We thank Kevin Frans, Homer Walker, Aravind Venugopal, and Bogdan Mazoure for their helpful discussions. We also thank Tongzhou Wang for sharing code during the early stages of our experiments. This work used the Della computational cluster provided by Princeton Research Computing, as well as the Ionic and Neuronic computing clusters maintained by the Department of Computer Science at Princeton University. This research was supported by ONR N00014-22-1-2773 and was partly supported by the Korea Foundation for Advanced Studies (KFAS). This material is based upon work supported by the National Science Foundation under Award No. 2441665. Any opinions, findings, conclusions, or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the National Science Foundation.

## References

- Achiam et al. (2023)
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al.
Gpt-4 technical report.
*arXiv preprint arXiv:2303.08774*, 2023.
- Ajay et al. (2021)
Anurag Ajay, Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofir Nachum.
OPAL: Offline primitive discovery for accelerating offline reinforcement learning.
In *International Conference on Learning Representations*, 2021.
URL [https://openreview.net/forum?id=V69LGwJ0lIN](https://openreview.net/forum?id=V69LGwJ0lIN).
- Ajay et al. (2023)
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B. Tenenbaum, Tommi S. Jaakkola, and Pulkit Agrawal.
Is conditional generative modeling all you need for decision making?
In *The Eleventh International Conference on Learning Representations*, 2023.
URL [https://openreview.net/forum?id=sP1fo2K9DFG](https://openreview.net/forum?id=sP1fo2K9DFG).
- Albergo & Vanden-Eijnden (2023)
Michael Samuel Albergo and Eric Vanden-Eijnden.
Building normalizing flows with stochastic interpolants.
In *The Eleventh International Conference on Learning Representations*, 2023.
URL [https://openreview.net/forum?id=li7qeBbCR1t](https://openreview.net/forum?id=li7qeBbCR1t).
- Alemi et al. (2017)
Alexander A. Alemi, Ian Fischer, Joshua V. Dillon, and Kevin Murphy.
Deep variational information bottleneck.
In *International Conference on Learning Representations*, 2017.
URL [https://openreview.net/forum?id=HyxQzBceg](https://openreview.net/forum?id=HyxQzBceg).
- Alonso et al. (2024)
Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos J Storkey, Tim Pearce, and François Fleuret.
Diffusion for world modeling: Visual details matter in atari.
*Advances in Neural Information Processing Systems*, 37:58757–58791, 2024.
- Barber & Agakov (2004)
David Barber and Felix Agakov.
The im algorithm: a variational approach to information maximization.
*Advances in neural information processing systems*, 16(320):201, 2004.
- Barreto et al. (2017)
André Barreto, Will Dabney, Rémi Munos, Jonathan J Hunt, Tom Schaul, Hado P van Hasselt, and David Silver.
Successor features for transfer in reinforcement learning.
*Advances in neural information processing systems*, 30, 2017.
- Barreto et al. (2018)
Andre Barreto, Diana Borsa, John Quan, Tom Schaul, David Silver, Matteo Hessel, Daniel Mankowitz, Augustin Zidek, and Remi Munos.
Transfer in deep reinforcement learning using successor features and generalised policy improvement.
In *International Conference on Machine Learning*, pp. 501–510. PMLR, 2018.
- Black et al. (2024)
Kevin Black, Noah Brown, Danny Driess, Adnan Esmail, Michael Equi, Chelsea Finn, Niccolo Fusai, Lachy Groom, Karol Hausman, Brian Ichter, et al.
$\pi_{0}$: A vision-language-action flow model for general robot control.
*arXiv preprint arXiv:2410.24164*, 2024.
- Blier et al. (2021)
Léonard Blier, Corentin Tallec, and Yann Ollivier.
Learning successor states and goal-dependent values: A mathematical viewpoint.
*arXiv preprint arXiv:2101.07123*, 2021.
- Borsa et al. (2018)
Diana Borsa, André Barreto, John Quan, Daniel Mankowitz, Rémi Munos, Hado Van Hasselt, David Silver, and Tom Schaul.
Universal successor features approximators.
*arXiv preprint arXiv:1812.07626*, 2018.
- Bradbury et al. (2018)
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang.
JAX: composable transformations of Python+NumPy programs, 2018.
URL [http://github.com/jax-ml/jax](http://github.com/jax-ml/jax).
- Brandfonbrener et al. (2021)
David Brandfonbrener, Will Whitney, Rajesh Ranganath, and Joan Bruna.
Offline rl without off-policy evaluation.
*Advances in neural information processing systems*, 34:4933–4946, 2021.
- Brohan et al. (2022)
Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, et al.
Rt-1: Robotics transformer for real-world control at scale.
*arXiv preprint arXiv:2212.06817*, 2022.
- Brown et al. (2020)
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.
Language models are few-shot learners.
*Advances in neural information processing systems*, 33:1877–1901, 2020.
- Burda et al. (2019)
Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov.
Exploration by random network distillation.
In *International Conference on Learning Representations*, 2019.
URL [https://openreview.net/forum?id=H1lJJnR5Ym](https://openreview.net/forum?id=H1lJJnR5Ym).
- Campbell et al. (2024)
Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola.
Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design.
In *International Conference on Machine Learning*, pp. 5453–5512. PMLR, 2024.
- Caron et al. (2021)
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin.
Emerging properties in self-supervised vision transformers.
In *Proceedings of the IEEE/CVF international conference on computer vision*, pp. 9650–9660, 2021.
- Chen et al. (2023)
Boyuan Chen, Chuning Zhu, Pulkit Agrawal, Kaiqing Zhang, and Abhishek Gupta.
Self-supervised reinforcement learning that transfers using random features.
*Advances in Neural Information Processing Systems*, 36:56411–56436, 2023.
- Chen et al. (2021)
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch.
Decision transformer: Reinforcement learning via sequence modeling.
*Advances in neural information processing systems*, 34:15084–15097, 2021.
- Dao et al. (2023)
Quan Dao, Hao Phung, Binh Nguyen, and Anh Tran.
Flow matching in latent space.
*arXiv preprint arXiv:2307.08698*, 2023.
- Dayan (1993)
Peter Dayan.
Improving generalization for temporal difference learning: The successor representation.
*Neural computation*, 5(4):613–624, 1993.
- Devlin et al. (2019)
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.
Bert: Pre-training of deep bidirectional transformers for language understanding.
In *Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume 1 (long and short papers)*, pp. 4171–4186, 2019.
- Ding et al. (2024)
Zihan Ding, Amy Zhang, Yuandong Tian, and Qinqing Zheng.
Diffusion world model.
*arXiv e-prints*, pp. arXiv–2402, 2024.
- Duan et al. (2016)
Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel.
Rl^2: Fast reinforcement learning via slow reinforcement learning.
*arXiv preprint arXiv:1611.02779*, 2016.
- Durrett (2019)
Rick Durrett.
*Probability: theory and examples*, volume 49.
Cambridge university press, 2019.
- Espeholt et al. (2018)
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, Iain Dunning, et al.
Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures.
In *International conference on machine learning*, pp. 1407–1416. PMLR, 2018.
- Eysenbach et al. (2019)
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine.
Diversity is all you need: Learning skills without a reward function.
In *International Conference on Learning Representations*, 2019.
URL [https://openreview.net/forum?id=SJx63jRqFm](https://openreview.net/forum?id=SJx63jRqFm).
- Eysenbach et al. (2020)
Benjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine.
C-learning: Learning to achieve goals via recursive classification.
*arXiv preprint arXiv:2011.08909*, 2020.
- Eysenbach et al. (2022)
Benjamin Eysenbach, Tianjun Zhang, Sergey Levine, and Russ R Salakhutdinov.
Contrastive learning as goal-conditioned reinforcement learning.
*Advances in Neural Information Processing Systems*, 35:35603–35620, 2022.
- Farebrother et al. (2025)
Jesse Farebrother, Matteo Pirotta, Andrea Tirinzoni, Remi Munos, Alessandro Lazaric, and Ahmed Touati.
Temporal difference flows.
In *ICLR 2025 Workshop on World Models: Understanding, Modelling and Scaling*, 2025.
URL [https://openreview.net/forum?id=nYL0pn3z3u](https://openreview.net/forum?id=nYL0pn3z3u).
- Frans et al. (2024)
Kevin Frans, Seohong Park, Pieter Abbeel, and Sergey Levine.
Unsupervised zero-shot reinforcement learning via functional reward encodings.
In *International Conference on Machine Learning*, pp. 13927–13942. PMLR, 2024.
- Frans et al. (2025)
Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel.
One step diffusion via shortcut models.
In *The Thirteenth International Conference on Learning Representations*, 2025.
URL [https://openreview.net/forum?id=OlzB6LnXcS](https://openreview.net/forum?id=OlzB6LnXcS).
- Fujimoto & Gu (2021)
Scott Fujimoto and Shixiang Shane Gu.
A minimalist approach to offline reinforcement learning.
*Advances in neural information processing systems*, 34:20132–20145, 2021.
- Fujimoto et al. (2018)
Scott Fujimoto, Herke Hoof, and David Meger.
Addressing function approximation error in actor-critic methods.
In *International conference on machine learning*, pp. 1587–1596. PMLR, 2018.
- Ghosh et al. (2023)
Dibya Ghosh, Chethan Anand Bhateja, and Sergey Levine.
Reinforcement learning from passive data via latent intentions.
In *International Conference on Machine Learning*, pp. 11321–11339. PMLR, 2023.
- Grill et al. (2020)
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al.
Bootstrap your own latent-a new approach to self-supervised learning.
*Advances in neural information processing systems*, 33:21271–21284, 2020.
- Haarnoja et al. (2018)
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine.
Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.
In *International conference on machine learning*, pp. 1861–1870. Pmlr, 2018.
- Hafner et al. (2023)
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap.
Mastering diverse domains through world models.
*arXiv preprint arXiv:2301.04104*, 2023.
- Hansen et al. (2019)
Steven Hansen, Will Dabney, Andre Barreto, Tom Van de Wiele, David Warde-Farley, and Volodymyr Mnih.
Fast task inference with variational intrinsic successor features.
*arXiv preprint arXiv:1906.05030*, 2019.
- Hansen-Estruch et al. (2023)
Philippe Hansen-Estruch, Ilya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine.
Idql: Implicit q-learning as an actor-critic method with diffusion policies.
*arXiv preprint arXiv:2304.10573*, 2023.
- Hausman et al. (2017)
Karol Hausman, Yevgen Chebotar, Stefan Schaal, Gaurav Sukhatme, and Joseph J Lim.
Multi-modal imitation learning from unstructured demonstrations using generative adversarial nets.
*Advances in neural information processing systems*, 30, 2017.
- He et al. (2020)
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick.
Momentum contrast for unsupervised visual representation learning.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pp. 9729–9738, 2020.
- He et al. (2022)
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick.
Masked autoencoders are scalable vision learners.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pp. 16000–16009, 2022.
- Henderson et al. (2017)
Peter Henderson, Wei-Di Chang, Pierre-Luc Bacon, David Meger, Joelle Pineau, and Doina Precup.
Optiongan: Learning joint reward-policy options using generative adversarial inverse reinforcement learning.
*ArXiv*, abs/1709.06683, 2017.
URL [https://api.semanticscholar.org/CorpusID:7079525](https://api.semanticscholar.org/CorpusID:7079525).
- Hendrycks & Gimpel (2016)
Dan Hendrycks and Kevin Gimpel.
Gaussian error linear units (gelus).
*arXiv preprint arXiv:1606.08415*, 2016.
- Higgins et al. (2017a)
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner.
beta-vae: Learning basic visual concepts with a constrained variational framework.
In *International conference on learning representations*, 2017a.
- Higgins et al. (2017b)
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner.
beta-VAE: Learning basic visual concepts with a constrained variational framework.
In *International Conference on Learning Representations*, 2017b.
URL [https://openreview.net/forum?id=Sy2fzU9gl](https://openreview.net/forum?id=Sy2fzU9gl).
- Ho et al. (2020)
Jonathan Ho, Ajay Jain, and Pieter Abbeel.
Denoising diffusion probabilistic models.
*Advances in neural information processing systems*, 33:6840–6851, 2020.
- Hu et al. (2023)
Hao Hu, Yiqin Yang, Jianing Ye, Ziqing Mai, and Chongjie Zhang.
Unsupervised behavior extraction via random intent priors.
*Advances in Neural Information Processing Systems*, 36:51491–51514, 2023.
- Jain et al. (2023)
Arnav Kumar Jain, Lucas Lehnert, Irina Rish, and Glen Berseth.
Maximum state entropy exploration using predecessor and successor representations.
*Advances in Neural Information Processing Systems*, 36:49991–50019, 2023.
- Janner et al. (2019)
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine.
When to trust your model: Model-based policy optimization.
*Advances in neural information processing systems*, 32, 2019.
- Janner et al. (2020)
Michael Janner, Igor Mordatch, and Sergey Levine.
Gamma-models: Generative temporal difference learning for infinite-horizon prediction.
*Advances in neural information processing systems*, 33:1724–1735, 2020.
- Janner et al. (2021)
Michael Janner, Qiyang Li, and Sergey Levine.
Offline reinforcement learning as one big sequence modeling problem.
*Advances in neural information processing systems*, 34:1273–1286, 2021.
- Janner et al. (2022)
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine.
Planning with diffusion for flexible behavior synthesis.
In *International Conference on Machine Learning*, pp. 9902–9915. PMLR, 2022.
- Jeen et al. (2024)
Scott Jeen, Tom Bewley, and Jonathan Cullen.
Zero-shot reinforcement learning from low quality data.
*Advances in Neural Information Processing Systems*, 37:16894–16942, 2024.
- Kim et al. (2022)
Jaekyeom Kim, Seohong Park, and Gunhee Kim.
Constrained gpi for zero-shot transfer in reinforcement learning.
*Advances in Neural Information Processing Systems*, 35:4585–4597, 2022.
- Kim et al. (2024)
Junsu Kim, Seohong Park, and Sergey Levine.
Unsupervised-to-online reinforcement learning.
*arXiv preprint arXiv:2408.14785*, 2024.
- Kingma (2014)
Diederik P Kingma.
Adam: A method for stochastic optimization.
*arXiv preprint arXiv:1412.6980*, 2014.
- Kingma & Welling (2013)
Diederik P Kingma and Max Welling.
Auto-encoding variational bayes.
*arXiv preprint arXiv:1312.6114*, 2013.
- Kostrikov et al. (2022)
Ilya Kostrikov, Ashvin Nair, and Sergey Levine.
Offline reinforcement learning with implicit q-learning.
In *International Conference on Learning Representations*, 2022.
URL [https://openreview.net/forum?id=68n2s9ZJWF8](https://openreview.net/forum?id=68n2s9ZJWF8).
- Kumar et al. (2020)
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine.
Conservative q-learning for offline reinforcement learning.
*Advances in neural information processing systems*, 33:1179–1191, 2020.
- Lambert et al. (2022)
Nathan Lambert, Kristofer Pister, and Roberto Calandra.
Investigating compounding prediction errors in learned dynamics models.
*arXiv preprint arXiv:2203.09637*, 2022.
- Laskin et al. (2020)
Michael Laskin, Aravind Srinivas, and Pieter Abbeel.
Curl: Contrastive unsupervised representations for reinforcement learning.
In *International conference on machine learning*, pp. 5639–5650. PMLR, 2020.
- Li et al. (2024)
Xuanlin Li, Kyle Hsu, Jiayuan Gu, Karl Pertsch, Oier Mees, Homer Rich Walke, Chuyuan Fu, Ishikaa Lunawat, Isabel Sieh, Sean Kirmani, et al.
Evaluating real-world robot manipulation policies in simulation.
*arXiv preprint arXiv:2405.05941*, 2024.
- Li et al. (2017)
Yunzhu Li, Jiaming Song, and Stefano Ermon.
Infogail: Interpretable imitation learning from visual demonstrations.
*Advances in neural information processing systems*, 30, 2017.
- Lipman et al. (2023)
Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le.
Flow matching for generative modeling.
In *The Eleventh International Conference on Learning Representations*, 2023.
URL [https://openreview.net/forum?id=PqvMRDCJT9t](https://openreview.net/forum?id=PqvMRDCJT9t).
- Lipman et al. (2024)
Yaron Lipman, Marton Havasi, Peter Holderrieth, Neta Shaul, Matt Le, Brian Karrer, Ricky TQ Chen, David Lopez-Paz, Heli Ben-Hamu, and Itai Gat.
Flow matching guide and code.
*arXiv preprint arXiv:2412.06264*, 2024.
- Liu et al. (2023)
Xingchao Liu, Chengyue Gong, and qiang liu.
Flow straight and fast: Learning to generate and transfer data with rectified flow.
In *The Eleventh International Conference on Learning Representations*, 2023.
URL [https://openreview.net/forum?id=XVjTT1nw5z](https://openreview.net/forum?id=XVjTT1nw5z).
- Lu et al. (2019)
Jiasen Lu, Dhruv Batra, Devi Parikh, and Stefan Lee.
Vilbert: Pretraining task-agnostic visiolinguistic representations for vision-and-language tasks.
*Advances in neural information processing systems*, 32, 2019.
- Ma et al. (2023)
Yecheng Jason Ma, Shagun Sodhani, Dinesh Jayaraman, Osbert Bastani, Vikash Kumar, and Amy Zhang.
VIP: Towards universal visual reward and representation via value-implicit pre-training.
In *The Eleventh International Conference on Learning Representations*, 2023.
URL [https://openreview.net/forum?id=YJ7o2wetJ2](https://openreview.net/forum?id=YJ7o2wetJ2).
- Maaten & Hinton (2008)
Laurens van der Maaten and Geoffrey Hinton.
Visualizing data using t-sne.
*Journal of machine learning research*, 9(Nov):2579–2605, 2008.
- Machado et al. (2017)
Marlos C Machado, Clemens Rosenbaum, Xiaoxiao Guo, Miao Liu, Gerald Tesauro, and Murray Campbell.
Eigenoption discovery through the deep successor representation.
*arXiv preprint arXiv:1710.11089*, 2017.
- Margossian & Blei (2024)
Charles C Margossian and David M Blei.
Amortized variational inference: When and why?
In *Uncertainty in Artificial Intelligence*, pp. 2434–2449. PMLR, 2024.
- Mazoure et al. (2023)
Bogdan Mazoure, Benjamin Eysenbach, Ofir Nachum, and Jonathan Tompson.
Contrastive value learning: Implicit models for simple offline RL.
In *7th Annual Conference on Robot Learning*, 2023.
URL [https://openreview.net/forum?id=oqOfLP6bJy](https://openreview.net/forum?id=oqOfLP6bJy).
- Mazzaglia et al. (2022)
Pietro Mazzaglia, Tim Verbelen, Bart Dhoedt, Alexandre Lacoste, and Sai Rajeswar.
Choreographer: Learning and adapting skills in imagination.
*arXiv preprint arXiv:2211.13350*, 2022.
- Mendonca et al. (2021)
Russell Mendonca, Oleh Rybkin, Kostas Daniilidis, Danijar Hafner, and Deepak Pathak.
Discovering and achieving goals via world models.
*Advances in Neural Information Processing Systems*, 34:24379–24391, 2021.
- Mnih et al. (2015)
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al.
Human-level control through deep reinforcement learning.
*nature*, 518(7540):529–533, 2015.
- Myers et al. (2024)
Vivek Myers, Chongyi Zheng, Anca Dragan, Sergey Levine, and Benjamin Eysenbach.
Learning temporal distances: Contrastive successor features can provide a metric structure for decision-making.
In *International Conference on Machine Learning*, pp. 37076–37096. PMLR, 2024.
- Nair et al. (2023)
Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta.
R3m: A universal visual representation for robot manipulation.
In *Conference on Robot Learning*, pp. 892–909. PMLR, 2023.
- Nemecek & Parr (2021)
Mark Nemecek and Ronald Parr.
Policy caches with successor features.
In *International Conference on Machine Learning*, pp. 8025–8033. PMLR, 2021.
- Ni et al. (2024)
Tianwei Ni, Benjamin Eysenbach, Erfan Seyedsalehi, Michel Ma, Clement Gehring, Aditya Mahajan, and Pierre-Luc Bacon.
Bridging state and history representations: Understanding self-predictive rl.
*arXiv preprint arXiv:2401.08898*, 2024.
- Ouyang et al. (2022)
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al.
Training language models to follow instructions with human feedback.
*Advances in neural information processing systems*, 35:27730–27744, 2022.
- O’Neill et al. (2024)
Abby O’Neill, Abdul Rehman, Abhiram Maddukuri, Abhishek Gupta, Abhishek Padalkar, Abraham Lee, Acorn Pooley, Agrim Gupta, Ajay Mandlekar, Ajinkya Jain, et al.
Open x-embodiment: Robotic learning datasets and rt-x models: Open x-embodiment collaboration 0.
In *2024 IEEE International Conference on Robotics and Automation (ICRA)*, pp. 6892–6903. IEEE, 2024.
- Parisi et al. (2022)
Simone Parisi, Aravind Rajeswaran, Senthil Purushwalkam, and Abhinav Gupta.
The unsurprising effectiveness of pre-trained vision models for control.
In *international conference on machine learning*, pp. 17359–17371. PMLR, 2022.
- Park et al. (2023a)
Seohong Park, Dibya Ghosh, Benjamin Eysenbach, and Sergey Levine.
Hiql: Offline goal-conditioned rl with latent states as actions.
*Advances in Neural Information Processing Systems*, 36:34866–34891, 2023a.
- Park et al. (2023b)
Seohong Park, Oleh Rybkin, and Sergey Levine.
Metra: Scalable unsupervised rl with metric-aware abstraction.
*arXiv preprint arXiv:2310.08887*, 2023b.
- Park et al. (2024a)
Seohong Park, Kevin Frans, Sergey Levine, and Aviral Kumar.
Is value learning really the main bottleneck in offline rl?
*arXiv preprint arXiv:2406.09329*, 2024a.
- Park et al. (2024b)
Seohong Park, Tobias Kreiman, and Sergey Levine.
Foundation policies with hilbert representations.
In *International Conference on Machine Learning*, pp. 39737–39761. PMLR, 2024b.
- Park et al. (2025a)
Seohong Park, Kevin Frans, Benjamin Eysenbach, and Sergey Levine.
OGBench: Benchmarking offline goal-conditioned RL.
In *The Thirteenth International Conference on Learning Representations*, 2025a.
URL [https://openreview.net/forum?id=M992mjgKzI](https://openreview.net/forum?id=M992mjgKzI).
- Park et al. (2025b)
Seohong Park, Qiyang Li, and Sergey Levine.
Flow q-learning.
*arXiv preprint arXiv:2502.02538*, 2025b.
- Pertsch et al. (2021)
Karl Pertsch, Youngwoon Lee, and Joseph Lim.
Accelerating reinforcement learning with learned skill priors.
In *Conference on robot learning*, pp. 188–204. PMLR, 2021.
- Peters & Schaal (2007)
Jan Peters and Stefan Schaal.
Reinforcement learning by reward-weighted regression for operational space control.
In *Proceedings of the 24th international conference on Machine learning*, pp. 745–750, 2007.
- Peters et al. (2010)
Jan Peters, Katharina Mulling, and Yasemin Altun.
Relative entropy policy search.
In *Proceedings of the AAAI Conference on Artificial Intelligence*, volume 24, pp. 1607–1612, 2010.
- Pong et al. (2022)
Vitchyr H Pong, Ashvin V Nair, Laura M Smith, Catherine Huang, and Sergey Levine.
Offline meta-reinforcement learning with online self-supervision.
In *International Conference on Machine Learning*, pp. 17811–17829. Pmlr, 2022.
- Poole et al. (2019)
Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker.
On variational bounds of mutual information.
In *International conference on machine learning*, pp. 5171–5180. PMLR, 2019.
- Radford et al. (2018)
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al.
Improving language understanding by generative pre-training.
2018.
- Radford et al. (2021)
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.
Learning transferable visual models from natural language supervision.
In *International conference on machine learning*, pp. 8748–8763. PmLR, 2021.
- Rakelly et al. (2019)
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen.
Efficient off-policy meta-reinforcement learning via probabilistic context variables.
In *International conference on machine learning*, pp. 5331–5340. PMLR, 2019.
- Ren et al. (2025)
Allen Z. Ren, Justin Lidard, Lars Lien Ankile, Anthony Simeonov, Pulkit Agrawal, Anirudha Majumdar, Benjamin Burchfiel, Hongkai Dai, and Max Simchowitz.
Diffusion policy policy optimization.
In *The Thirteenth International Conference on Learning Representations*, 2025.
URL [https://openreview.net/forum?id=mEpqHvbD2h](https://openreview.net/forum?id=mEpqHvbD2h).
- Rombach et al. (2022)
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.
High-resolution image synthesis with latent diffusion models.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pp. 10684–10695, 2022.
- Rummery & Niranjan (1994)
Gavin A Rummery and Mahesan Niranjan.
*On-line Q-learning using connectionist systems*, volume 37.
University of Cambridge, Department of Engineering Cambridge, UK, 1994.
- Saxe et al. (2018)
Andrew Michael Saxe, Yamini Bansal, Joel Dapello, Madhu Advani, Artemy Kolchinsky, Brendan Daniel Tracey, and David Daniel Cox.
On the information bottleneck theory of deep learning.
In *International Conference on Learning Representations*, 2018.
URL [https://openreview.net/forum?id=ry_WPG-A-](https://openreview.net/forum?id=ry_WPG-A-).
- Schulman et al. (2015)
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz.
Trust region policy optimization.
In *International conference on machine learning*, pp. 1889–1897. PMLR, 2015.
- Schulman et al. (2017)
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.
Proximal policy optimization algorithms.
*arXiv preprint arXiv:1707.06347*, 2017.
- Schwarzer et al. (2020)
Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman.
Data-efficient reinforcement learning with self-predictive representations.
*arXiv preprint arXiv:2007.05929*, 2020.
- Sermanet et al. (2018)
Pierre Sermanet, Corey Lynch, Yevgen Chebotar, Jasmine Hsu, Eric Jang, Stefan Schaal, Sergey Levine, and Google Brain.
Time-contrastive networks: Self-supervised learning from video.
In *2018 IEEE international conference on robotics and automation (ICRA)*, pp. 1134–1141. IEEE, 2018.
- Sikchi et al. (2024)
Harshit Sikchi, Qinqing Zheng, Amy Zhang, and Scott Niekum.
Dual RL: Unification and new methods for reinforcement and imitation learning.
In *The Twelfth International Conference on Learning Representations*, 2024.
URL [https://openreview.net/forum?id=xt9Bu66rqv](https://openreview.net/forum?id=xt9Bu66rqv).
- Sodhani et al. (2021)
Shagun Sodhani, Amy Zhang, and Joelle Pineau.
Multi-task reinforcement learning with context-based representations.
In *International conference on machine learning*, pp. 9767–9779. PMLR, 2021.
- Sohl-Dickstein et al. (2015)
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.
Deep unsupervised learning using nonequilibrium thermodynamics.
In *International conference on machine learning*, pp. 2256–2265. pmlr, 2015.
- Song et al. (2021)
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.
Score-based generative modeling through stochastic differential equations.
In *International Conference on Learning Representations*, 2021.
URL [https://openreview.net/forum?id=PxTIG12RRHS](https://openreview.net/forum?id=PxTIG12RRHS).
- Sutton et al. (1998)
Richard S Sutton, Andrew G Barto, et al.
*Reinforcement learning: An introduction*, volume 1.
MIT press Cambridge, 1998.
- Talvitie (2014)
Erik Talvitie.
Model regularization for stable sample rollouts.
In *UAI*, pp. 780–789, 2014.
- Tarasov et al. (2023a)
Denis Tarasov, Vladislav Kurenkov, Alexander Nikulin, and Sergey Kolesnikov.
Revisiting the minimalist approach to offline reinforcement learning.
*Advances in Neural Information Processing Systems*, 36:11592–11620, 2023a.
- Tarasov et al. (2023b)
Denis Tarasov, Alexander Nikulin, Dmitry Akimov, Vladislav Kurenkov, and Sergey Kolesnikov.
Corl: Research-oriented deep offline reinforcement learning library.
*Advances in Neural Information Processing Systems*, 36:30997–31020, 2023b.
- Tassa et al. (2018)
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al.
Deepmind control suite.
*arXiv preprint arXiv:1801.00690*, 2018.
- Team et al. (2023)
Gemini Team, Rohan Anil, Sebastian Borgeaud, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, Katie Millican, et al.
Gemini: a family of highly capable multimodal models.
*arXiv preprint arXiv:2312.11805*, 2023.
- Team et al. (2024)
Octo Model Team, Dibya Ghosh, Homer Walke, Karl Pertsch, Kevin Black, Oier Mees, Sudeep Dasari, Joey Hejna, Tobias Kreiman, Charles Xu, et al.
Octo: An open-source generalist robot policy.
*arXiv preprint arXiv:2405.12213*, 2024.
- Thakoor et al. (2022)
Shantanu Thakoor, Mark Rowland, Diana Borsa, Will Dabney, Rémi Munos, and André Barreto.
Generalised policy improvement with geometric policy composition.
In *International Conference on Machine Learning*, pp. 21272–21307. PMLR, 2022.
- Tishby et al. (2000)
Naftali Tishby, Fernando C Pereira, and William Bialek.
The information bottleneck method.
*arXiv preprint physics/0004057*, 2000.
- Tomar et al. (2024)
Manan Tomar, Philippe Hansen-Estruch, Philip Bachman, Alex Lamb, John Langford, Matthew E Taylor, and Sergey Levine.
Video occupancy models.
*arXiv preprint arXiv:2407.09533*, 2024.
- Touati & Ollivier (2021)
Ahmed Touati and Yann Ollivier.
Learning one representation to optimize all rewards.
*Advances in Neural Information Processing Systems*, 34:13–23, 2021.
- Touati et al. (2023)
Ahmed Touati, Jérémy Rapin, and Yann Ollivier.
Does zero-shot reinforcement learning exist?
In *The Eleventh International Conference on Learning Representations*, 2023.
URL [https://openreview.net/forum?id=MYEap_OcQI](https://openreview.net/forum?id=MYEap_OcQI).
- Toussaint & Storkey (2006)
Marc Toussaint and Amos Storkey.
Probabilistic inference for solving discrete and continuous state markov decision processes.
In *Proceedings of the 23rd international conference on Machine learning*, pp. 945–952, 2006.
- Touvron et al. (2023)
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al.
Llama 2: Open foundation and fine-tuned chat models.
*arXiv preprint arXiv:2307.09288*, 2023.
- Valieva & Banerjee (2024)
Khadichabonu Valieva and Bikramjit Banerjee.
Quasimetric value functions with dense rewards.
*arXiv preprint arXiv:2409.08724*, 2024.
- Van Hasselt et al. (2016)
Hado Van Hasselt, Arthur Guez, and David Silver.
Deep reinforcement learning with double q-learning.
In *Proceedings of the AAAI conference on artificial intelligence*, volume 30, 2016.
- Vaswani et al. (2017)
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.
Attention is all you need.
*Advances in neural information processing systems*, 30, 2017.
- Walke et al. (2023)
Homer Rich Walke, Kevin Black, Tony Z Zhao, Quan Vuong, Chongyi Zheng, Philippe Hansen-Estruch, Andre Wang He, Vivek Myers, Moo Jin Kim, Max Du, et al.
Bridgedata v2: A dataset for robot learning at scale.
In *Conference on Robot Learning*, pp. 1723–1736. PMLR, 2023.
- Wang et al. (2018)
Qing Wang, Jiechao Xiong, Lei Han, Han Liu, Tong Zhang, et al.
Exponentially weighted imitation learning for batched historical data.
*Advances in Neural Information Processing Systems*, 31, 2018.
- Wang et al. (2023)
Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou.
Diffusion policies as an expressive policy class for offline reinforcement learning.
In *The Eleventh International Conference on Learning Representations*, 2023.
URL [https://openreview.net/forum?id=AHvFDPi-FA](https://openreview.net/forum?id=AHvFDPi-FA).
- Watkins & Dayan (1992)
Christopher JCH Watkins and Peter Dayan.
Q-learning.
*Machine learning*, 8:279–292, 1992.
- Xiao et al. (2022)
Tete Xiao, Ilija Radosavovic, Trevor Darrell, and Jitendra Malik.
Masked visual pre-training for motor control.
*arXiv preprint arXiv:2203.06173*, 2022.
- Yarats et al. (2022)
Denis Yarats, David Brandfonbrener, Hao Liu, Michael Laskin, Pieter Abbeel, Alessandro Lazaric, and Lerrel Pinto.
Don’t change the algorithm, change the data: Exploratory data for offline reinforcement learning.
In *ICLR 2022 Workshop on Generalizable Policy Learning in Physical World*, 2022.
URL [https://openreview.net/forum?id=Su-zh4a41Z5](https://openreview.net/forum?id=Su-zh4a41Z5).
- Yu et al. (2020)
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn.
Gradient surgery for multi-task learning.
*Advances in neural information processing systems*, 33:5824–5836, 2020.
- Yu et al. (2022)
Tianhe Yu, Aviral Kumar, Yevgen Chebotar, Karol Hausman, Chelsea Finn, and Sergey Levine.
How to leverage unlabeled data in offline reinforcement learning.
In *International Conference on Machine Learning*, pp. 25611–25635. PMLR, 2022.
- Zhai et al. (2023)
Xiaohua Zhai, Basil Mustafa, Alexander Kolesnikov, and Lucas Beyer.
Sigmoid loss for language image pre-training.
In *Proceedings of the IEEE/CVF international conference on computer vision*, pp. 11975–11986, 2023.
- Zhang et al. (2021)
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine.
Learning invariant representations for reinforcement learning without reconstruction.
In *International Conference on Learning Representations*, 2021.
URL [https://openreview.net/forum?id=-2FCwDKRREu](https://openreview.net/forum?id=-2FCwDKRREu).
- Zheng et al. (2024)
Chongyi Zheng, Ruslan Salakhutdinov, and Benjamin Eysenbach.
Contrastive difference predictive coding.
*The Twelfth International Conference on Learning Representations*, 2024.
URL [https://openreview.net/forum?id=0akLDTFR9x](https://openreview.net/forum?id=0akLDTFR9x).
- Zheng et al. (2025)
Chongyi Zheng, Jens Tuyls, Joanne Peng, and Benjamin Eysenbach.
Can a MISL fly? analysis and ingredients for mutual information skill learning.
In *The Thirteenth International Conference on Learning Representations*, 2025.
URL [https://openreview.net/forum?id=xoIeVdFO7U](https://openreview.net/forum?id=xoIeVdFO7U).
- Zhou et al. (2025)
Zhiyuan Zhou, Andy Peng, Qiyang Li, Sergey Levine, and Aviral Kumar.
Efficient online reinforcement learning fine-tuning need not retain offline data.
In *The Thirteenth International Conference on Learning Representations*, 2025.
URL [https://openreview.net/forum?id=HN0CYZbAPw](https://openreview.net/forum?id=HN0CYZbAPw).
- Zhu et al. (2024)
Chuning Zhu, Xinqi Wang, Tyler Han, Simon S Du, and Abhishek Gupta.
Distributional successor features enable zero-shot policy optimization.
*arXiv preprint arXiv:2403.06328*, 2024.

## Appendix A Preliminaries

Figure: Algorithm 3 Euler method for solving the flow ODE (Eq. [9](https://arxiv.org/html/2506.08902v2#A1.E9)).

### A.1 Value functions and the Actor-critic framework

The goal of RL is to learn a policy $\pi:{\mathcal{S}}\to\Delta({\mathcal{A}})$ that maximizes the expected discounted return $J(\pi)=\mathbb{E}_{\tau\sim\pi(\tau)}\left[\sum_{h=0}^{\infty}\gamma^{h}r_{h}\right]$, where $\tau$ is a trajectory sampled by the policy. We will use $\beta:{\mathcal{S}}\to\Delta({\mathcal{A}})$ to denote the behavioral policy. Given a policy $\pi$, we measure the expected discounted return starting from a state-action pair $(s,a)$ and a state $s$ as the (unnormalized) Q-function and the value function, respectively:

$$ $\displaystyle Q^{\pi}(s,a)=\mathbb{E}_{\tau\sim\pi(\tau)}\left[\left.\sum_{h=0}^{\infty}\gamma^{h}r_{h}\right|s_{0}=s,a_{0}=a\right],\>V^{\pi}(s)=\mathbb{E}_{a\sim\pi(a\mid s)}\left[Q^{\pi}(s,a)\right].$ $$

Prior actor-critic methods (Schulman et al., 2015; 2017; Haarnoja et al., 2018; Fujimoto et al., 2018; Kumar et al., 2020; Fujimoto & Gu, 2021) typically maximize the RL objective $J(\pi)$ by *(1)* learning an estimate $Q$ of $Q^{\pi}$ via the temporal difference (TD) loss (policy evaluation) and then *(2)* improving the policy $\pi$ by selecting actions that maximizes $Q$ (policy improvement):

$$ $\displaystyle Q^{k+1}$ $\displaystyle\leftarrow\operatorname*{arg\,max}_{Q}\mathbb{E}_{(s,a,r,s^{\prime})\sim p^{\beta}(s,a,r,s^{\prime}),a^{\prime}\sim\pi^{k}(a^{\prime}\mid s^{\prime})}\left[\left(Q(s,a)-(r+\gamma Q^{k}(s^{\prime},a^{\prime}))\right)^{2}\right]$ $\displaystyle\pi^{k+1}$ $\displaystyle\leftarrow\operatorname*{arg\,max}_{\pi}\mathbb{E}_{s\sim p^{\beta}(s),a\sim\pi(a\mid s)}\left[Q^{k+1}(s,a)\right],$ $$

where $k$ indicates the number of updates and $\beta$ is the behavioral policy representing either a replay buffer (online RL) or a fixed dataset (offline RL).

### A.2 Flow matching

The goal of flow matching methods is to transform a simple noise distribution (e.g., a $d$-dimensional standard Gaussian) into a target distribution $p_{{\mathcal{X}}}$ over some space ${\mathcal{X}}\subset\mathbb{R}^{d}$ that we want to approximate.
Specifically, flow matching uses a time-dependent vector field $v:[0,1]\times\mathbb{R}^{d}\to\mathbb{R}^{d}$ to construct a time-dependent diffeomorphic flow $\phi:[0,1]\times\mathbb{R}^{d}\to\mathbb{R}^{d}$  (Lipman et al., 2023; 2024) that realizes the transformation from a single noise $\epsilon$ to a generative sample $\hat{x}$, following the ODE

$$ $\displaystyle\frac{d}{dt}\phi(t,\epsilon)=v(t,\phi(t,\epsilon)),\;\phi(0,\epsilon)=\epsilon,\;\phi(1,\epsilon)=\hat{x}.$ (9) $$

We will use $t$ to denote a time step for flow matching and sample the noise $\epsilon$ from a standard Gaussian distribution ${\mathcal{N}}(0,I)$ throughout our discussions.(^6^66In theory, the noise can be drawn from any distribution, not necessarily limited to a Gaussian (Liu et al., 2023).) Prior work has proposed various formulations for learning the vector field (Lipman et al., 2023; Campbell et al., 2024; Liu et al., 2023; Albergo & Vanden-Eijnden, 2023) and we adopt the simplest flow matching objective building upon optimal transport (Liu et al., 2023) and conditional flow matching (CFM) (Lipman et al., 2023),

$$ $\displaystyle{\mathcal{L}}_{\text{CFM}}(v)=\mathbb{E}_{\begin{subarray}{c}t\sim\textsc{Unif}([0,1]),\\ x\sim p_{{\mathcal{X}}}(x),\epsilon\sim{\mathcal{N}}(0,I)\end{subarray}}\left[\lVert v(t,x^{t})-(x-\epsilon)\rVert_{2}^{2}\right],$ (10) $$

where $\textsc{Unif}([0,1])$ is the uniform distribution over the unit interval and $x^{t}=tx+(1-t)\epsilon$ is a linear interpolation between the ground-truth sample $x$ and the Gaussian noise $\epsilon$.
Importantly, we can generate a sample from the vector field $v$ by numerically solving the ODE (Eq. [9](https://arxiv.org/html/2506.08902v2#A1.E9)).
We will use the Euler method (Alg. [3](https://arxiv.org/html/2506.08902v2#alg3)) as our ODE solver following prior practice (Liu et al., 2023; Park et al., 2025b; Frans et al., 2025).

### A.3 Temporal difference flows

Given a policy $\pi$, prior work (Farebrother et al., 2025) models the occupancy measure $p^{\pi}_{\gamma}$ by optimizing the vector field $v:[0,1]\times{\mathcal{S}}\times{\mathcal{S}}\times{\mathcal{A}}\to{\mathcal{S}}$ using the following loss:

$$ $\displaystyle{\mathcal{L}}_{\text{TD flow}}(v)$ $\displaystyle=(1-\gamma){\mathcal{L}}_{\text{TD current flow}}(v)+\gamma{\mathcal{L}}_{\text{TD future flow}}(v)$ (11) $\displaystyle{\mathcal{L}}_{\text{TD current flow}}(v)$ $\displaystyle=\mathbb{E}_{\begin{subarray}{c}t\sim\textsc{Unif}([0,1]),\epsilon\sim{\mathcal{N}}(0,I),\\ (s,a)\sim p^{\beta}(s,a)\end{subarray}}\left[\lVert v(t,s^{t},s,a)-(s-\epsilon)\rVert_{2}^{2}\right]$ $\displaystyle{\mathcal{L}}_{\text{TD future flow}}(v)$ $\displaystyle=\mathbb{E}_{\begin{subarray}{c}t\sim\textsc{Unif}([0,1]),\epsilon\sim{\mathcal{N}}(0,I),\\ (s,a,s^{\prime})\sim p^{\beta}(s,a,s^{\prime}),a^{\prime}\sim\pi(a^{\prime}\mid s^{\prime})\\ \end{subarray}}\left[\lVert v(t,\bar{s}_{f}^{t},s,a)-\bar{v}(t,\bar{s}_{f}^{t},s^{\prime},a^{\prime})\rVert_{2}^{2}\right],$ $$

where $p^{\beta}(s,a)$ and $p^{\beta}(s,a,s^{\prime})$ denote the joint distribution of transitions, $s^{t}=ts+(1-t)\epsilon$ is a linear interpolation between the current state $s$ and the noise $\epsilon$, and $\bar{v}$ denotes an exponential moving average of historical $v$ over iterations (a target vector field) (Grill et al., 2020; Mnih et al., 2015; Caron et al., 2021). Of particular note is that we obtain a target future state $\bar{s}_{f}$ by applying the Euler method (Alg. [3](https://arxiv.org/html/2506.08902v2#alg3)) to $\bar{v}$ at the next state-action pair $(s^{\prime},a^{\prime})$, where $a^{\prime}$ is sampled from the target policy $\pi$ of interest, and the noisy future state $\bar{s}_{f}^{t}=t\bar{s}_{f}+(1-t)\epsilon$ is a linear interpolation between this future state $\bar{s}_{f}$ and the noise $\epsilon$. Intuitively, minimizing ${\mathcal{L}}_{\text{TD current flow}}$ reconstructs the distribution of current state $s$, while minimizing ${\mathcal{L}}_{\text{TD future flow}}$ bootstraps the vector field $v$ at a noisy target future state $\bar{s}_{f}^{t}$, similar to Q-learning (Watkins & Dayan, 1992). Choosing the target policy $\pi$ to be the same as the behavioral policy $\beta$, we obtain a SARSA (Rummery & Niranjan, 1994) variant of the loss optimizing the SARSA flows. We call the loss in Eq. [11](https://arxiv.org/html/2506.08902v2#A1.E11) the TD flow loss(^7^77The TD flow loss is called the TD^2-CFM loss in Farebrother et al. (2025) and we rename it for simplicity.) and use the SARSA variant of it to learn generative occupancy models.

## Appendix B Further discussions on the problem setting

### B.1 The consistency assumption on intentions

We now discuss the reason for making the consistency assumption (Assumption [1](https://arxiv.org/html/2506.08902v2#Thmtheorem1)) on latent intentions. Since we use a heterogeneous behavioral policy to collect the unlabeled dataset, each unknown user intention indexed their own behavioral policy $\beta:{\mathcal{S}}\times{\mathcal{Z}}\to\Delta({\mathcal{A}})$. The key observation is that the occupancy measure of each intention-conditioned behavioral policy follows its own Bellman equations (Similar to Eq. [1](https://arxiv.org/html/2506.08902v2#S3.E1)):

$$ $\displaystyle p^{\beta}_{\gamma}(s_{f}\mid s,a,{\color[rgb]{1,.5,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,.5,0}z})=(1-\gamma)\delta_{s}(s_{f})+\gamma\mathbb{E}_{\begin{subarray}{c}s^{\prime}\sim p(s^{\prime}\mid s,a),\\ a^{\prime}\sim\beta(a^{\prime}\mid s^{\prime},{\color[rgb]{1,.5,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,.5,0}z})\end{subarray}}\left[p^{\beta}_{\gamma}(s_{f}\mid s^{\prime},a^{\prime},{\color[rgb]{1,.5,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,.5,0}z})\right],$ $$

suggesting that the same latent $z$ propagates through the transitions with the same underlying user intentions. Importantly, this propagation requires using a TD loss to estimate the behavioral occupancy measure, which aligns with the goal of our SARSA flow-matching losses (Eq. [5](https://arxiv.org/html/2506.08902v2#S4.E5)). We note that prior work (Touati & Ollivier, 2021) also adapts the same formulation of the intention-conditioned occupancy measure for zero-shot RL.

### B.2 Distinctions from Meta RL and Multi-task RL

Our problem setting is conceptually similar to meta RL (Duan et al., 2016; Rakelly et al., 2019; Pong et al., 2022) with two key distinctions. First, offline meta RL methods typically have access to explicit task descriptions (e.g., a one-hot task indicator) together with task-specific datasets. These descriptions and datasets induce a clear clustering of transitions. In contrast, our method must infer this structure from a heterogeneous dataset in an unsupervised manner. Second, offline meta RL trains on reward-labeled data during the meta-training phase, where task-specific rewards provide supervision for policy learning. In contrast, during pre-training, our method learns a generative model that predicts future states from inferred intentions without using any task-specific reward signals.

Similar to the distinctions between our setting and offline meta RL problems, our method does *not* fall into the multi-task RL category (Sodhani et al., 2021; Yu et al., 2020). During pre-training, *(1)* InFOM does not have access to task descriptions or task-specific datasets, and *(2)* it does not use any supervision from task-specific reward signals. Instead, InFOM pre-trains a generative, multi-step transition model that facilitates value estimation for downstream tasks.

## Appendix C Theoretical analyses

### C.1 The evidence lower bound and its connection with an information bottleneck

We first derive the evidence lower bound for optimizing the latent variable model and then show its connection with an information bottleneck. Given the unlabeled dataset $D$, we want to maximize the likelihood of consecutive transitions $(s,a,s^{\prime},a^{\prime})$ and a future state $s_{f}$ sampled from the same trajectory following the behavioral joint distribution $p^{\beta}(s,a,s_{f},s^{\prime},a^{\prime})=p^{\beta}(s)\beta(a\mid s)p^{\beta}_{\gamma}(s_{f}\mid s,a)p(s^{\prime}\mid s,a)\beta(a^{\prime}\mid s^{\prime})$. We use $(s^{\prime},a^{\prime})$ to encode the intention $z$ by the encoder $p_{e}(z\mid s,a)$ and $(s,a,s_{f},z)$ to learn the occupancy models $q_{d}(s_{f}\mid s,a,z)$, employing an ELBO of the likelihood of the prior data:

$$ $\displaystyle\hskip 13.00005pt\mathbb{E}_{p^{\beta}(s,a,s_{f})}\left[\log q_{d}(s_{f}\mid s,a)\right]$ $\displaystyle=\mathbb{E}_{p^{\beta}(s,a,s_{f},s^{\prime},a^{\prime})}\left[\log q_{d}(s_{f}\mid s,a)\right]$ $\displaystyle=\mathbb{E}_{p^{\beta}(s,a,s_{f},s^{\prime},a^{\prime})}\left[\log\mathbb{E}_{p(z)}\left[q_{d}(s_{f}\mid s,a,z)\right]\right]$ $\displaystyle\stackrel{{\scriptstyle(a)}}{{=}}\mathbb{E}_{p^{\beta}(s,a,s_{f},s^{\prime},a^{\prime})}\left[\log\mathbb{E}_{p(z)}\left[q_{d}(s_{f}\mid s,a,z)\frac{p_{e}(z\mid s^{\prime},a^{\prime})}{p_{e}(z\mid s^{\prime},a^{\prime})}\right]\right]$ $\displaystyle\stackrel{{\scriptstyle(b)}}{{\geq}}\mathbb{E}_{p^{\beta}(s,a,s_{f},s^{\prime},a^{\prime})}[\mathbb{E}_{p_{e}(z\mid s^{\prime},a^{\prime})}[\log q_{d}(s_{f}\mid s,a,z)]-D_{\text{KL}}(p_{e}(z\mid s^{\prime},a^{\prime})\parallel p(z)]$ $\displaystyle\stackrel{{\scriptstyle(c)}}{{\geq}}\mathbb{E}_{p^{\beta}(s,a,s_{f},s^{\prime},a^{\prime})}[\mathbb{E}_{p_{e}(z\mid s^{\prime},a^{\prime})}[\log q_{d}(s_{f}\mid s,a,z)]-\lambda D_{\text{KL}}(p_{e}(z\mid s^{\prime},a^{\prime})\parallel p(z)]$ $\displaystyle=L(p_{e},q_{d}).$ $$

where in *(a)* we introduce the amortized variational encoder $p_{e}(z\mid s^{\prime},a^{\prime})$, in *(b)* we apply the Jensen’s inequality (Durrett, 2019), and in *(c)* we introduce a coefficient $\lambda$ to control the strength of the KL divergence regularization as in Higgins et al. (2017b). This ELBO can also be cast as a variational lower bound on an information bottleneck with the chain of random variables $(S^{\prime},A^{\prime})\to Z\to(S,A,S_{f})$ (Tishby et al., 2000; Alemi et al., 2017; Saxe et al., 2018):

$$ $\displaystyle\hskip 13.00005ptI^{\beta}(S,A,S_{f};Z)-\lambda I^{\beta}(S^{\prime},A^{\prime};Z)$ $\displaystyle\stackrel{{\scriptstyle(a)}}{{=}}I^{\beta}(S,A,S_{f};Z)-\lambda\mathbb{E}_{p^{\beta}(s^{\prime},a^{\prime})}\left[D_{\mathrm{KL}}(p_{e}(z\mid s^{\prime},a^{\prime})\parallel p_{e}(z))\right]$ $\displaystyle\stackrel{{\scriptstyle(b)}}{{\geq}}I^{\beta}(S,A,S_{f};Z)-\lambda\mathbb{E}_{p^{\beta}(s^{\prime},a^{\prime})}\left[D_{\mathrm{KL}}(p_{e}(z\mid s^{\prime},a^{\prime})\parallel p(z))\right]$ $\displaystyle\stackrel{{\scriptstyle(c)}}{{\geq}}\mathbb{E}_{\begin{subarray}{c}p^{\beta}(s,a,s_{f},s^{\prime},a^{\prime})\\ p_{e}(z\mid s^{\prime},a^{\prime})\end{subarray}}\left[\log q_{d}(s,a,s_{f}\mid z)\right]-\lambda\mathbb{E}_{p^{\beta}(s^{\prime},a^{\prime})}\left[D_{\mathrm{KL}}(p_{e}(z\mid s^{\prime},a^{\prime})\parallel p(z))\right]+H^{\beta}(S,A,S_{f})$ $\displaystyle\stackrel{{\scriptstyle(d)}}{{\geq}}\mathbb{E}_{\begin{subarray}{c}p^{\beta}(s,a,s_{f},s^{\prime},a^{\prime})\\ p_{e}(z\mid s^{\prime},a^{\prime})\end{subarray}}\left[\log q_{d}(s_{f}\mid s,a,z)\right]-\lambda\mathbb{E}_{p^{\beta}(s^{\prime},a^{\prime})}\left[D_{\mathrm{KL}}(p_{e}(z\mid s^{\prime},a^{\prime})\parallel p(z))\right]+\text{const}$ $\displaystyle=L(p_{e},q_{d})+\text{const},$ $$

where in *(a)* we use the definition of $I^{\beta}(S^{\prime},A^{\prime};Z)$ and $p_{e}(z)$ is the marginal distribution of latent intentions $z$ defined as $p_{e}(z)=\int p^{\beta}(s^{\prime},a^{\prime})p_{e}(z\mid s^{\prime},a^{\prime})ds^{\prime}da^{\prime}$, in *(b)* we apply the non-negative property of the KL divergence $D_{\mathrm{KL}}(p_{e}(z)\parallel p(z))$, in *(c)* we apply the standard variation lower bound of the mutual information (Barber & Agakov, 2004; Poole et al., 2019) to incorporate the decoder (occupancy models) $q_{d}(s,a,s_{f}\mid z)$, and in *(d)* we choose the variational decoder to satisfy $\log q_{d}(s,a,s_{f}\mid z)=\log p^{\beta}(s,a)+\log q_{d}(s_{f}\mid s,a,z)$ and consider the entropy $H^{\beta}(S,A,S_{f})$ as a constant.

### C.2 Intutions and discussions about the implicit generalized policy improvement

The intuition for the expectile distillation loss (Eq. [7](https://arxiv.org/html/2506.08902v2#S4.E7)) is that the scalar Q function $Q(\cdot,\cdot)$ is a *one-step* summary of the average reward at future states sampled from the flow occupancy models, while the expectile loss serves as a "softmax" operator over the entire latent intention space. Theoretically, this expectile loss is guaranteed to converge to the maximum over $p(z)$ under mild assumptions (See Sec. 4.4 in Kostrikov et al. (2022) for details). Therefore, given an infinite amount of samples ($N\to\infty$) and an expectile $\mu\to 1$, the $Q$ converges to the greedy value functions:

$$ $\displaystyle Q^{\star}(s,a)=\max_{z\sim p(z)}\frac{1}{(1-\gamma)}\mathbb{E}_{q_{d}(s_{f}\mid s,a,z)}[r(s_{f})].$ $$

If we further assume that the flow occupancy models are optimal, i.e., $q_{d}^{\star}(s_{f}\mid s,a,z)=p^{\beta}(s_{f}\mid s,a,z)$, then the optimal $Q$ corresponds to a greedy value function under the behavioral policy $\beta$:

$$ $\displaystyle Q^{\star}(s,a)=\max_{z\sim p(z)}Q^{\beta}(s,a,z).$ $$

Unlike Q-learning, which converges to the optimal Q-function sequentially (Watkins & Dayan, 1992; Sutton et al., 1998), the implicit GPI proposes a new policy that is strictly no worse than the set of policies that correspond to each $Q_{z}$ in parallel (See Sec. 4.1 in Barreto et al. (2017) for further discussions). Unlike one-step policy improvement (Wang et al., 2018; Brandfonbrener et al., 2021; Peters & Schaal, 2007; Peters et al., 2010), implicit GPI is able to converge to the optimal policy for a downstream task, assuming that the task-specific intention has been captured during pre-training.

### C.3 One-step policy improvement with flow occupancy models

The FOM + one-step PI variant performs one-step policy improvement using a flow occupancy model $q_{d}(s_{f}\mid s,a)$ that is *not* conditioned on latent intentions. This flow occupancy model captures the discounted state occupancy measure of the (average) behavioral policy. After training the flow occupancy model, FOM + one-step PI fits a Q function and extracts a behavioral-regularized policy:

$$ $\displaystyle Q$ $\displaystyle\leftarrow\operatorname*{arg\,min}_{Q}\frac{1}{1-\gamma}\mathbb{E}_{(s,a)\sim p^{\tilde{\beta}}(s,a),s_{f}\sim q_{d}(s_{f}\mid s,a)}[(Q(s,a)-r(s_{f}))^{2}],$ $\displaystyle\pi$ $\displaystyle\leftarrow\operatorname*{arg\,max}_{\pi}\mathbb{E}_{(s,a)\sim p^{\tilde{\beta}}(s,a),a^{\pi}\sim\pi(a^{\pi}\mid s)}\left[Q(s,a^{\pi})+\alpha\log\pi(a\mid s)\right].$ $$

Intuitively, the first objective fits the behavioral Q function based on the dual definition (Eq. [2](https://arxiv.org/html/2506.08902v2#S3.E2)),
and the second objective trains a policy to maximize this behavioral Q function, invoking one-step policy improvement.
While this simple objective sometimes achieves strong performance on some benchmark tasks (Brandfonbrener et al., 2021; Eysenbach et al., 2022),
it does not guarantee convergence to the optimal policy due to the use of a behavioral value function.

## Appendix D Experimental details

### D.1 Tasks and datasets

Our experiments use a suite of 36 state-based and 4 image-based control tasks taken from ExORL benchmarks Yarats et al. (2022) and OGBench task suite (Park et al., 2025a) (Fig. [3](https://arxiv.org/html/2506.08902v2#S5.F3)).

#### ExORL.

We use 16 state-based tasks from the ExORL (Yarats et al., 2022) benchmarks based on the DeepMind Control Suite (Tassa et al., 2018). These tasks involve controlling four robots (cheetah, walker, quadruped, and jaco) to achieve different locomotion behaviors. For each domain, the specific tasks are: cheetah {run, run backward, walk, walk backward}, walker {walk, run, stand, flip}, quadruped {run, jump, stand, walk}, jaco {reach top left, reach top right, reach bottom left, reach bottom right}. For all tasks in cheetah, walker, and quadruped, both the episode length and the maximum return are 1000. For all tasks in jaco, both the episode length and the maximum return are 250. Following prior work (Park et al., 2024b), we multiply the return of jaco tasks by $4$ to match other ExORL tasks during aggregation.

Following the prior work (Touati et al., 2023; Park et al., 2024b; Kim et al., 2024), we will use 5M unlabeled transitions collected by some exploration methods (e.g., RND (Burda et al., 2019)) for pre-training, and another 500K reward-labeled transitions collected by the same exploratory policy for fine-tuning. The fine‑tuning datasets are labeled with task‑specific dense rewards (Yarats et al., 2022), except in the jaco domains, where the reward signals are sparse.

#### OGBench.

We use 20 state-based manipulation tasks from four domains (cube single, cube double, scene, and puzzle 4x4) in the OGBench task suite Park et al. (2025a), where the goal is to control a simulated robot arm to rearrange various objects. For each domain, the specific tasks are: cube single {task 1 (pick and place cube to left), task 2 (pick and place cube to front), task 3 (pick and place cube to back), task 4 (pick and place cube diagonally), task 5 (pick and place cube off-diagonally)}, cube double {task 1 (pick and place one cube), task 2 (pick and place two cubes to right), task 3 (pick and place two cubes off-diagonally), task 4 (swap cubes), task 5 (stack cubes)}, scene {task 1 (open drawer and window), task 2 (close and lock drawer and window), task 3 (open drawer, close window, and pick and place cube to right), task 4 (put cube in drawer), task 5 (fetch cube from drawer and close window)}, puzzle 4x4 {task 1 (all red to all blue), task 2 (all blue to central red), task 3 (two blue to mix), task 4 (central red to all red), task 5 (mix to all red)}. Note that some of these tasks, e.g., cube double task 5 (stack cubes) and scene task 4 (put cube in drawer), involve interacting with the environment in a specific order and thus require long-horizon temporal reasoning. For all tasks in cube single, cube double, and scene, the maximum episode length is 400. For all tasks in puzzle 4x4, the maximum episode length is 800. We also use 4 image-based tasks in the OGBench task suite. Specifically, we consider visual cube single task 1, visual cube double task 1, visual scene task 1, and visual puzzle 4x4 task 1 from each domain respectively. The observations are $64\times 64\times 3$ RGB images. These tasks are challenging because the agent needs to reason from pixels directly. All the manipulation tasks from OGBench are originally designed for evaluating goal-conditioned RL algorithms (Park et al., 2025a).

For both state-based and image-based tasks from OGBench, we will use 1M unlabeled transitions collected by a non-Markovian expert policy with temporally correlated noise (the play datasets) for pre-training, and another 500K reward-labeled transitions collected by the same noisy expert policy for fine-tuning. Unlike the ExORL benchmarks, the fine-tuning datasets for OGBench tasks are relabeled with *semi-sparse* rewards (Park et al., 2025b), providing less supervision for the algorithm.

### D.2 Baselines

We compare InFOM with eight baselines across five categories of prior methods, focusing on different strategies for pre-training and fine-tuning in RL. First, implicit Q-Learning (IQL) (Kostrikov et al., 2022) and revisited behavior-regularized actor-critic (ReBRAC) (Tarasov et al., 2023a) are state-of-the-art offline RL algorithms based on the standard actor-critic framework (Appendix [A.1](https://arxiv.org/html/2506.08902v2#A1.SS1)).
Second, we compare to a variant of ReBRAC learning on top of representations pre-trained on the unlabeled datasets. We chose
an off-the-shelf self-supervised learning objective in vision tasks
called
self-distillation with no labels (DINO) (Caron et al., 2021) as our representation learning loss and name the resulting baseline DINO + ReBRAC. Third, our next baseline, model-based policy optimization (MBPO) (Janner et al., 2019), pre-trains a one-step model to predict transitions in the environment, similar to the next token prediction in language models (Radford et al., 2018). The one-step model is then used to
augment the datasets for downstream policy optimization. We will again use ReBRAC to extract the policy (MBPO + ReBRAC). Fourth, we also include comparisons against the InfoNCE variant of contrastive RL (Eysenbach et al., 2019) and temporal difference InfoNCE (Zheng et al., 2024), which pre-train the discounted state occupancy measure using Monte Carlo or temporal difference contrastive losses.
While our method fits generative occupancy models,
These two approaches predict the ratio of occupancy measures over some marginal densities serving as the discriminative counterparts. After pre-training the ratio predictors, importance sampling is required to recover the Q function (CRL + IS & TD InfoNCE + IS) (Mazoure et al., 2023; Zheng et al., 2024), enabling policy maximization. Fifth, our final set of baselines are prior unsupervised RL methods that pre-train a set of latent intentions and intention-conditioned policies using forward-backward representations (Touati & Ollivier, 2021) or a Hilbert space (Park et al., 2024b). Given a downstream task, these methods first infer the corresponding intention in a zero-shot manner and then fine-tune the policy using offline RL (Kim et al., 2024), differing from the implicit GPI as in our method.
We will use IQL as the fine-tuning algorithm and call the resulting methods FB + IQL and HILP + IQL. For image-based tasks, we selectively compare to four baselines: ReBRAC, CRL + IS, DINO + ReBRAC, and FB + IQL.

### D.3 Evaluation protocols

We compare the performance of InFOM against the eight baselines (Sec. [5.1](https://arxiv.org/html/2506.08902v2#S5.SS1)) after first pre-training each method for 1M gradient steps (250K gradient steps for image-based tasks) and then fine-tuning for 500K gradient steps (100K gradient steps for image-based tasks). We measure the episode return for tasks from ExORL benchmarks and the success rate for tasks from the OGBench task suite. For OGBench tasks, the algorithms still use the semi-sparse reward instead of the success rate for training. Following prior practice (Park et al., 2025b; Tarasov et al., 2023b), we do *not* report the best performance during fine-tuning and report the evaluation results averaged over 400K, 450K, and 500K gradient steps instead. For image-based tasks, we report the evaluation results averaged over 50K, 75K, and 100K gradient steps during fine-tuning. For evaluating the performance of different methods throughout the entire fine-tuning process, we defer the details to specific figures (e.g., Fig. [10](https://arxiv.org/html/2506.08902v2#A6.F10) & [10](https://arxiv.org/html/2506.08902v2#A6.F10)).

### D.4 Implementations and hyperparameters

**Table 1: Common hyperparameters for our method and the baselines.**
| Hyperparameter | Value |
| --- | --- |
| learning rate | $3\times 10^{-4}$ |
| optimizer | Adam (Kingma, 2014) |
| pre-training gradient steps | $1\times 10^{6}$ for state-based tasks, $2.5\times 10^{5}$ for image-based tasks |
| fine-tuning gradient steps | $5\times 10^{5}$ for state-based tasks, $1\times 10^{5}$ for image-based tasks |
| batch size | $256$ |
| MLP hidden layer sizes | $(512,512,512,512)$ |
| MLP activation function | GELU (Hendrycks & Gimpel, 2016) |
| discount factor $\gamma$ | $0.99$ |
| target network update coefficient | $5\times 10^{-3}$ |
| double Q aggregation | min |
| policy update frequency in fine-tuning | $1/4$ |
| image encoder | small IMPALA encoder (Espeholt et al., 2018; Park et al., 2025b) |
| image augmentation method | random cropping |
| image augmentation probability | $1.0$ for DINO + ReBRAC, $0.5$ for all other methods |
| image frame stack | $3$ |

**Table 2: Hyperparameters for InFOM. See Appendix [D.4](https://arxiv.org/html/2506.08902v2#A4.SS4) for descriptions of each hyperparameter.**
| Hyperparameter | Value |
| --- | --- |
| latent intention dimension $d$ | See Table [3](https://arxiv.org/html/2506.08902v2#A4.T3) |
| number of steps for the Euler method $T$ | $10$ |
| number of future states $N$ | $16$ |
| normalize the Q loss term in ${\mathcal{L}}(\pi)$ (Eq. [8](https://arxiv.org/html/2506.08902v2#S4.E8)) | No |
| expectile $\mu$ | See Table [3](https://arxiv.org/html/2506.08902v2#A4.T3) |
| KL divergence regularization coefficient $\lambda$ | See Table [3](https://arxiv.org/html/2506.08902v2#A4.T3) |
| behavioral cloning regurlaization coefficient $\alpha$ | See Table [3](https://arxiv.org/html/2506.08902v2#A4.T3) |

**Table 3: Domain-specific hyperparameters for our method and the baselines. We individually tune these hyperparameters for each domain and use the same set of hyperparameters for tasks in the same domain. See Appendix [D.4](https://arxiv.org/html/2506.08902v2#A4.SS4) for tasks used to tune these hyperparameters and descriptions of each hyperparameter. “-” indicates that the hyperparameter does not exist.**
|  | InFOM (Ours) | IQL | ReBRAC | DINO + ReBRAC | MBPO + ReBRAC | CRL + IS | TD InfoNCE + IS | FB + IQL | HILP + IQL |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Domain or Task | $d$ | $\mu$ | $\lambda$ | $\alpha$ | $\alpha$ | $\alpha_{\text{actor}}$ | $\alpha_{\text{critic}}$ | $\kappa_{\text{student}}$ | $N_{\text{imaginary}}$ | $H_{\text{imaginary}}$ | $\alpha$ | $\alpha$ | $\alpha_{\text{repr}}$ | $\alpha_{\text{AWR}}$ | $\alpha$ |
| cheetah | $128$ | $0.9$ | $0.05$ | $0.3$ | $1$ | $0.1$ | $0.1$ | $0.1$ | $128$ | $1$ | $0.03$ | $0.003$ | $1$ | $1$ | $1$ |
| walker | $512$ | $0.9$ | $0.1$ | $0.3$ | $1$ | $10$ | $0.1$ | $0.1$ | $128$ | $1$ | $0.03$ | $0.03$ | $1$ | $10$ | $10$ |
| quadruped | $512$ | $0.9$ | $0.005$ | $0.3$ | $10$ | $1$ | $1$ | $0.1$ | $128$ | $1$ | $0.03$ | $0.03$ | $10$ | $1$ | $10$ |
| jaco | $512$ | $0.9$ | $0.2$ | $0.1$ | $0.1$ | $0.1$ | $0.1$ | $0.1$ | $128$ | $1$ | $0.003$ | $0.03$ | $1$ | $1$ | $1$ |
| cube single | $512$ | $0.95$ | $0.05$ | $30$ | $1$ | $1$ | $1$ | $0.04$ | $256$ | $2$ | $30$ | $30$ | $10$ | $1$ | $1$ |
| cube double | $128$ | $0.9$ | $0.025$ | $30$ | $1$ | $1$ | $1$ | $0.04$ | $256$ | $2$ | $30$ | $30$ | $1$ | $10$ | $1$ |
| scene | $128$ | $0.99$ | $0.2$ | $300$ | $1$ | $1$ | $1$ | $0.1$ | $256$ | $2$ | $3$ | $3$ | $10$ | $10$ | $1$ |
| puzzle 4x4 | $128$ | $0.95$ | $0.1$ | $300$ | $10$ | $0.1$ | $0.1$ | $0.1$ | $256$ | $2$ | $3$ | $3$ | $10$ | $10$ | $1$ |
| visual cube single task 1 | $512$ | $0.95$ | $0.025$ | $30$ | - | $1$ | 0 | $0.1$ | - | - | $30$ | - | $10$ | $1$ | - |
| visual cube double task 1 | $128$ | $0.9$ | $0.01$ | $30$ | - | $0.1$ | 0 | $0.1$ | - | - | $30$ | - | $10$ | $1$ | - |
| visual scene task 1 | $128$ | $0.99$ | $0.1$ | $300$ | - | $0.1$ | $0.01$ | $0.1$ | - | - | $3$ | - | $10$ | $10$ | - |
| visual puzzle 4x4 task 1 | $128$ | $0.95$ | $0.1$ | $300$ | - | $0.1$ | $0.01$ | $0.1$ | - | - | $3$ | - | $10$ | $10$ | - |

In this section, we discuss the implementation details and hyperparameters for InFOM and the eight baselines. Whenever possible, we use the same set of hyperparameters for all methods (Table [1](https://arxiv.org/html/2506.08902v2#A4.T1)) across all tasks, including learning rate, network architecture, batch size, image encoder, etc. Of particular note is that we use asynchronous policy training (Zhou et al., 2025), where we update the policy 4 times less frequently than other models during fine-tuning. For specific hyperparameters of each method, we tune them on the following tasks from each domain and use one set of hyperparameters for every task in that domain. For image-based tasks, we tune hyperparameters for each task individually.

- •
cheetah: cheetah run
- •
walker: walker walk
- •
quadruped: quadruped jump
- •
jaco: jaco reach top left
- •
cube single: cube single task 2
- •
cube double: cube double task 2
- •
scene: scene task 2
- •
puzzle 4x4: puzzle 4x4 task 4

#### InFOM.

InFOM consists of two main components for pre-training: the intention encoder and the flow occupancy models. First, we use a Gaussian distribution conditioned on the next state-action pair as the intention encoding distribution. Following prior work (Kingma & Welling, 2013; Alemi et al., 2017), we model the intention encoder as a multilayer perceptron (MLP) that takes the next state-action pair $(s^{\prime},a^{\prime})$ as input and outputs two heads representing the mean and the (log) standard deviation of the Gaussian. We apply layer normalization to the intention encoder to stabilize optimization. We use the reparameterization trick (Kingma & Welling, 2013) to backpropagate the gradients from the flow-matching loss and the KL divergence regularization (Eq. [4](https://arxiv.org/html/2506.08902v2#S4.E4)) into the intention encoder. Our initial experiments suggest that the dimension of the latent intention space $d$ is an important hyperparameter, and we sweep over $\{64,128,256,512\}$ and find that $d=512$ is sufficient for most ExORL tasks and $d=128$ is generally good enough for all OGBench tasks. For the coefficient of the KL divergence regularization $\lambda$, we sweep over $\{2.0,1.0,0.2,0.1,0.05,0.025,0.01,0.005\}$ to find the best $\lambda$ for each domain. Second, we use flow-matching vector fields to model the flow occupancy models. The vector field is an MLP that takes in a noisy future state $s_{f}^{t}$, a state-action pair $(s,a)$, and a latent intention $z$, and outputs the vector field with the same dimension as the state. We apply layer normalization to the vector field to stabilize optimization. As mentioned in Sec. [3](https://arxiv.org/html/2506.08902v2#S3), we use flow-matching objectives based on optimal transport (linear path) and sample the time step $t$ from the uniform distribution over the unit interval. Following prior work (Park et al., 2025b), we use a fixed $T=10$ steps (step size = $0.1$) for the Euler method and do not apply a sinusoidal embedding for the time. To make a fair comparison with other baselines, we also pre-train a behavioral cloning policy that serves as initialization for fine-tuning.

For fine-tuning, InFOM learns three components: the reward predictor, the critic, and the policy, while fine-tuning the intention encoder and the flow occupancy models. The reward predictor is an MLP that predicts the scalar reward of a state trained using mean squared error. We apply layer normalization to the reward predictor to stabilize learning. The critic is an MLP that predicts double Q values (Van Hasselt et al., 2016; Fujimoto et al., 2018) of a state-action pair, without conditioning on the latent intention. We apply layer normalization to the critic to stabilize learning. We train the critic using the expectile distillation loss (Eq. [7](https://arxiv.org/html/2506.08902v2#S4.E7)) and sweep the expectile over $\{0.9,0.95,0.99\}$ to find the best $\mu$ for each domain. We use $N=16$ future states sampled from the flow occupancy models to compute the average reward, which we find to be sufficient. We use the minimum of the double Q predictions to prevent overestimation. The policy is an MLP that outputs a Gaussian distribution with a unit standard deviation. In our initial experiments, we find that the behavioral cloning coefficient $\alpha$ in Eq. [8](https://arxiv.org/html/2506.08902v2#S4.E8) is important, and we sweep over $\{300,30,3,0.3\}$ to find the best $\alpha$ for each domain. Following prior practice (Park et al., 2025b), we do not normalize the Q loss term in the actor loss ${\mathcal{L}}(\pi)$ (Eq. [8](https://arxiv.org/html/2506.08902v2#S4.E8)) as in Fujimoto & Gu (2021). Other choices of the policy network include the diffusion model (Ren et al., 2025; Wang et al., 2023) and the flow-matching model (Park et al., 2025b), and we leave investigating these policy networks to future work.

For image-based tasks, following prior work (Park et al., 2025b), we use a smaller variant of the IMPALA encoder (Espeholt et al., 2018) and apply random cropping augmentation with a probability of 0.5. We also apply frame stacking with three images. Table [2](https://arxiv.org/html/2506.08902v2#A4.T2) and Table [3](https://arxiv.org/html/2506.08902v2#A4.T3) summarize the hyperparameters for InFOM.

#### IQL and ReBRAC.

We reuse the IQL (Kostrikov et al., 2022) implementation and the ReBRAC (Tarasov et al., 2023a) implementation from Park et al. (2025b). Since learning a critic requires reward-labeled datasets or relabeling rewards for unlabeled datasets (Yu et al., 2022), we simply pre-train a behavioral cloning policy. During the fine-tuning, we use the behavioral cloning policy as initialization and train a critic from scratch using the TD error (Kostrikov et al., 2022; Fujimoto & Gu, 2021; Tarasov et al., 2023a). Following prior work (Park et al., 2025b), we use the same expectile value $0.9$ for IQL on all tasks, and sweep over $\{100,10,1,0.1,0.01\}$ to find the best AWR inverse temperature $\alpha$ for each domain. For ReBRAC, we tune the behavioral cloning (BC) regularization coefficients for the actor and the critic separately. We use the range $\{100,10,1,0.1\}$ to search for the best actor BC coefficient $\alpha_{\text{actor}}$ and use the range $\{100,10,1,0.1,0\}$ to search for the best critic BC coefficient $\alpha_{\text{critic}}$. We use the default values for other hyperparameters following the implementation from Park et al. (2025b). See Table [3](https://arxiv.org/html/2506.08902v2#A4.T3) for domain-specific hyperparameters.

#### DINO + ReBRAC.

We implement DINO on top of ReBRAC. DINO (Caron et al., 2021) learns a state encoder using two augmentations of the same state. For state-based tasks, the state encoder is an MLP that outputs representations. We apply two clipped Gaussian noises centered at zero to the same state to obtain those augmentations. The standard deviation of the Gaussian noise is set to $0.2$, and we clip the noise into $[-0.2,0.2]$ on all domains. For image-based tasks, the state encoder is the small IMPALA encoder that also outputs representations. We apply two different random cropings to the same image observation to obtain those augmentations. We sweep over $\{0.01,0.04,0.1,0.4\}$ for the temperature for student representations $\kappa_{\text{student}}$ and use a fixed temperature $0.04$ for teacher representations on all domains. We use a representation space with 512 dimensions. We update the target representation centroid with a fixed ratio $0.1$. During pre-training, we learns the DINO representations along with a behavioral cloning policy. During fine-tuning, we learn the actor and the critic using ReBRAC on top of DINO representations, while continuing to fine-tune those DINO representations. We use the same BC coefficients $\alpha_{\text{actor}}$ and $\alpha_{\text{critic}}$ as in ReBRAC. For image-based tasks, we apply random cropping to the same image twice with a probability of 1.0 and use those two augmentations to compute the teacher and the student representations. See Table [3](https://arxiv.org/html/2506.08902v2#A4.T3) for domain-specific temperatures for student representations.

#### MBPO + ReBRAC.

We implement MBPO (Janner et al., 2019) on top of ReBRAC and only consider this baseline for state-based tasks. MBPO learns a one-step transition MLP to predict the residual between the next state $s^{\prime}$ and the current state $s$ conditioned on the current state-action pair $(s,a)$. We pre-train the one-step model with a behavioral cloning policy. During fine-tuning, we use the model with a learned reward predictor to collect imaginary rollouts. We *only* use these imaginary rollouts to learn the actor and the critic. We sweep over $\{64,128,256\}$ for the number of imaginary rollouts to collect for each gradient step $N_{\text{imaginary}}$ and sweep over $\{1,2,4\}$ for the number of steps in each rollout $H_{\text{imaginary}}$. We use the same BC coefficient as in ReBRAC. See Table [3](https://arxiv.org/html/2506.08902v2#A4.T3) for the domain-specific number of imaginary rollouts and number of steps in each rollout.

#### CRL + IS and TD InfoNCE + IS.

We mostly reuse the CRL (Eysenbach et al., 2022) implementation based on the InfoNCE loss from Park et al. (2025a) and adapt it to our setting by adding the important sampling component. We implement TD InfoNCE by adapting the official implementations (Zheng et al., 2024). For both methods, we pre-train the classifiers that predict the ratio between the occupancy measures and the marginal densities over future states with a behavioral cloning policy. We use the SARSA variant of TD InfoNCE during pre-training. After pre-training the classifiers, we learn a reward predictor and apply importance sampling weights predicted by the classifiers to a set of future states sampled from the fine-tuning datasets to estimate $Q$. This Q estimation then drives policy optimization. We use a single future state from the fine-tuning dataset to construct the importance sampling estimation, which is sufficient. We use 512-dimensional contrastive representations. We sweep over $\{300,30,3,0.3,0.03\}$ for the BC coefficient $\alpha$ (Table [3](https://arxiv.org/html/2506.08902v2#A4.T3)).

#### FB + IQL and HILP + IQL.

We implement FB (Touati & Ollivier, 2021) and HILP (Park et al., 2024b) by adapting the FB implementation from Jeen et al. (2024) and the HILP implementation from Kim et al. (2024). During pre-training, for FB, we pre-train the forward-backward representations and the intention-conditioned policies in an actor-critic manner. We use a coefficient $1$ for the orthonormality regularization of the backward representations. We use 512-dimensional forward-backward representations. We sample the latent intentions for pre-training from either a standard Gaussian distribution (with probability $0.5$) or the backward representations for a batch of states (with probability $0.5$), normalizing those latent intentions to length $\sqrt{512}$. We sweep over $\{100,10,1,0.1\}$ for the BC coefficient $\alpha_{\text{repr}}$. For HILP, we pre-train the Hilbert representations $\phi$ and Hilbert foundation policies using an actor-critic framework as well. We use implicit value learning to learn the Hilbert representations following implementations from Park et al. (2024a; 2025a). We set the expectile to $0.9$ for all domains. We sweep over $\{100,10,1,0.1\}$ to find the best AWR inverse temperature $\alpha$. We also use a 512-dimensional Hilbert representation space. To construct the intrinsic rewards, we first sample the latent intention $z$ from a standard Gaussian, normalizing them to length $\sqrt{512}$, and then use the representation of the next state $\phi(s^{\prime})$ and the representation of the current state $\phi(s)$ to compute the intrinsic reward $(\phi(s^{\prime})-\phi(s))^{\top}z$.

During fine-tuning, we first infer a task-specific backward representation or a Hilbert representation using a small amount of transitions (10K) from the fine-tuning datasets, and then invoke IQL to learn the critic and the actor using downstream rewards conditioned on the inferred representations. For FB, we sweep over $\{100,10,1,0.1\}$ for the AWR inverse temperature $\alpha_{\text{AWR}}$ for IQL. For HILP, we reuse the same AWR inverse temperature in representation learning for IQL. See Table [3](https://arxiv.org/html/2506.08902v2#A4.T3) for domain-specific BC coefficients and AWR inverse temperatures.

**Table 4: Evaluation on ExORL and OGBench benchmarks. Following OGBench (Park et al., 2025a), we bold values at and above $95\%$ of the best performance for each task.**
| Task | InFOM (Ours) | IQL | ReBRAC | DINO + ReBRAC | MBPO + ReBRAC | CRL + IS | TD InfoNCE + IS | FB + IQL | HILP + IQL |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| cheetah run | $97.6\pm 7.8$ | $80.0\pm 8.4$ | $97.2\pm 12.9$ | $87.2\pm 8.6$ | $\mathbf{104.7}\pm{2.4}$ | $73.3\pm 6.7$ | $68.2\pm 8.9$ | $83.3\pm 10.9$ | $90.3\pm 1.9$ |
| cheetah run backward | $\mathbf{104.7}\pm{7.3}$ | $77.0\pm 12.6$ | $84.9\pm 3.7$ | $67.1\pm 6.4$ | $87.0\pm 4.8$ | $74.7\pm 8.1$ | $74.3\pm 17.1$ | $67.3\pm 7.0$ | $64.4\pm 6.4$ |
| cheetah walk | $254.8\pm 158.6$ | $357.9\pm 16.4$ | $\mathbf{443.4}\pm{15.3}$ | $383.5\pm 10.3$ | $\mathbf{447.4}\pm{12.7}$ | $327.4\pm 38.7$ | $336.7\pm 22.1$ | $346.5\pm 24.3$ | $366.8\pm 6.9$ |
| cheetah walk backward | $251.8\pm 116.9$ | $303.7\pm 12.6$ | $\mathbf{403.0}\pm{16.1}$ | $318.4\pm 23.0$ | $\mathbf{398.6}\pm{16.0}$ | $330.2\pm 8.5$ | $326.3\pm 45.1$ | $298.0\pm 22.8$ | $318.1\pm 11.4$ |
| walker walk | $\mathbf{467.3}\pm{82.1}$ | $208.6\pm 3.7$ | $208.1\pm 5.8$ | $228.0\pm 3.7$ | $327.6\pm 4.5$ | $213.3\pm 7.8$ | $212.2\pm 13.2$ | $225.3\pm 6.7$ | $225.4\pm 3.7$ |
| walker run | $\mathbf{116.3}\pm{15.3}$ | $92.4\pm 0.6$ | $97.8\pm 1.2$ | $98.5\pm 1.0$ | $107.6\pm 1.2$ | $91.5\pm 3.2$ | $91.0\pm 3.7$ | $97.4\pm 1.2$ | $97.4\pm 2.2$ |
| walker stand | $\mathbf{581.2}\pm{72.1}$ | $409.1\pm 2.3$ | $460.6\pm 1.1$ | $453.0\pm 3.1$ | $458.1\pm 2.5$ | $409.0\pm 7.5$ | $397.2\pm 6.0$ | $446.8\pm 7.1$ | $443.3\pm 3.8$ |
| walker flip | $\mathbf{358.8}\pm{10.3}$ | $260.3\pm 2.8$ | $\textbf{344.6}\pm 2.7$ | $320.3\pm 4.3$ | $\textbf{341.8}\pm 3.7$ | $255.0\pm 8.0$ | $231.6\pm 6.9$ | $287.0\pm 3.1$ | $280.7\pm 5.4$ |
| quadruped run | $341.8\pm 41.2$ | $358.0\pm 6.2$ | $343.0\pm 2.6$ | $344.7\pm 2.9$ | $\mathbf{395.1}\pm{2.6}$ | $323.4\pm 2.9$ | $222.1\pm 39.7$ | $367.0\pm 3.8$ | $371.1\pm 11.5$ |
| quadruped jump | $626.0\pm 6.8$ | $628.5\pm 7.8$ | $605.2\pm 7.8$ | $573.0\pm 9.6$ | $\mathbf{666.9}\pm{3.4}$ | $576.7\pm 13.7$ | $421.4\pm 93.4$ | $\mathbf{639.4}\pm{8.9}$ | $626.5\pm 14.5$ |
| quadruped stand | $\mathbf{718.3}\pm{18.7}$ | $\mathbf{714.2}\pm{9.8}$ | $688.6\pm 5.0$ | $663.2\pm 8.3$ | $\mathbf{703.7}\pm{3.6}$ | $653.1\pm 8.4$ | $457.1\pm 47.7$ | $\mathbf{728.9}\pm{11.5}$ | $\mathbf{715.6}\pm{13.9}$ |
| quadruped walk | $360.7\pm 7.9$ | $\mathbf{375.1}\pm{3.7}$ | $343.5\pm 7.1$ | $\mathbf{391.4}\pm{7.2}$ | $\mathbf{390.0}\pm{5.7}$ | $309.6\pm 9.6$ | $243.1\pm 29.2$ | $\mathbf{388.9}\pm{7.0}$ | $\mathbf{393.4}\pm{3.4}$ |
| jaco reach top left | $\mathbf{742.5}\pm{43.7}$ | $74.7\pm 19.6$ | $59.0\pm 4.9$ | $17.5\pm 3.8$ | $60.1\pm 6.2$ | $29.1\pm 4.7$ | $31.5\pm 3.0$ | $25.0\pm 11.4$ | $40.4\pm 11.5$ |
| jaco reach top right | $\mathbf{687.5}\pm{46.7}$ | $40.6\pm 14.0$ | $38.0\pm 13.1$ | $11.0\pm 4.1$ | $52.5\pm 10.8$ | $21.4\pm 6.5$ | $25.5\pm 10.3$ | $16.2\pm 3.2$ | $25.1\pm 9.6$ |
| jaco reach bottom left | $\mathbf{746.7}\pm{12.6}$ | $77.1\pm 12.5$ | $44.5\pm 4.0$ | $13.7\pm 2.8$ | $43.4\pm 4.6$ | $19.8\pm 8.8$ | $26.6\pm 5.9$ | $19.8\pm 4.0$ | $27.8\pm 4.6$ |
| jaco reach bottom right | $\mathbf{733.0}\pm{19.6}$ | $78.7\pm 19.1$ | $41.4\pm 5.0$ | $8.3\pm 2.8$ | $34.0\pm 6.0$ | $19.6\pm 2.0$ | $25.4\pm 5.7$ | $12.4\pm 2.7$ | $24.7\pm 3.9$ |
| cube single task 1 | $\mathbf{92.5}\pm{4.0}$ | $53.0\pm 8.7$ | $67.3\pm 14.2$ | $1.8\pm 1.0$ | $77.8\pm 11.7$ | $10.1\pm 2.7$ | $13.8\pm 3.8$ | $17.7\pm 8.8$ | $32.9\pm 9.2$ |
| cube single task 2 | $78.4\pm 12.3$ | $51.7\pm 15.1$ | $\mathbf{93.7}\pm{3.5}$ | $1.2\pm 0.6$ | $\mathbf{94.2}\pm{2.0}$ | $3.7\pm 2.8$ | $8.5\pm 5.6$ | $16.7\pm 8.6$ | $26.5\pm 15.4$ |
| cube single task 3 | $56.4\pm 36.9$ | $41.5\pm 5.3$ | $\mathbf{94.8}\pm{0.8}$ | $1.5\pm 1.4$ | $\mathbf{93.1}\pm{4.7}$ | $12.5\pm 3.2$ | $11.7\pm 7.4$ | $16.0\pm 12.2$ | $35.5\pm 14.7$ |
| cube single task 4 | $\mathbf{91.5}\pm{14.2}$ | $42.2\pm 8.3$ | $\mathbf{89.5}\pm{3.6}$ | $0.5\pm 1.0$ | $\mathbf{88.7}\pm{4.7}$ | $1.7\pm 1.7$ | $3.3\pm 3.0$ | $18.7\pm 9.9$ | $36.4\pm 14.9$ |
| cube single task 5 | $70.0\pm 39.1$ | $33.7\pm 12.9$ | $83.3\pm 6.8$ | $0.5\pm 0.6$ | $\mathbf{87.8}\pm{2.7}$ | $4.3\pm 2.2$ | $4.0\pm 3.2$ | $14.2\pm 12.0$ | $18.5\pm 5.6$ |
| cube double task 1 | $\mathbf{29.3}\pm{10.5}$ | $17.8\pm 9.6$ | $2.2\pm 1.7$ | $0.0\pm 0.0$ | $2.7\pm 1.1$ | $4.1\pm 1.9$ | $6.7\pm 2.7$ | $0.2\pm 0.3$ | $0.7\pm 1.1$ |
| cube double task 2 | $\mathbf{12.5}\pm{10.7}$ | $1.3\pm 1.2$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ |
| cube double task 3 | $\mathbf{11.6}\pm{8.3}$ | $0.3\pm 0.4$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ |
| cube double task 4 | $\mathbf{0.3}\pm{0.4}$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ |
| cube double task 5 | $\mathbf{2.8}\pm{4.6}$ | $1.5\pm 1.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.3\pm 0.7$ | $0.2\pm 0.3$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ |
| scene task 1 | $\mathbf{97.8}\pm{1.0}$ | $66.5\pm 13.1$ | $47.7\pm 7.2$ | $26.7\pm 4.3$ | $35.3\pm 7.7$ | $17.5\pm 5.1$ | $21.0\pm 4.3$ | $12.3\pm 11.3$ | $8.8\pm 3.0$ |
| scene task 2 | $\mathbf{15.6}\pm{3.4}$ | $2.5\pm 1.5$ | $7.8\pm 4.9$ | $1.3\pm 0.0$ | $5.6\pm 5.6$ | $2.3\pm 0.7$ | $1.7\pm 1.3$ | $1.5\pm 1.8$ | $1.2\pm 1.7$ |
| scene task 3 | $\mathbf{43.5}\pm{2.8}$ | $0.7\pm 0.5$ | $1.7\pm 1.1$ | $0.2\pm 0.3$ | $2.4\pm 0.8$ | $0.8\pm 0.3$ | $0.5\pm 1.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ |
| scene task 4 | $1.0\pm 0.7$ | $0.2\pm 0.3$ | $\mathbf{2.8}\pm{0.8}$ | $0.2\pm 0.3$ | $2.0\pm 1.3$ | $1.2\pm 1.4$ | $0.7\pm 1.3$ | $0.2\pm 0.3$ | $0.0\pm 0.0$ |
| scene task 5 | $\mathbf{0.3}\pm{0.4}$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $\mathbf{0.3}\pm{0.1}$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ |
| puzzle 4x4 task 1 | $\mathbf{24.2}\pm{14.4}$ | $2.3\pm 2.3$ | $12.8\pm 3.1$ | $0.3\pm 0.7$ | $16.9\pm 1.4$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.2\pm 0.3$ | $0.3\pm 0.6$ |
| puzzle 4x4 task 2 | $\mathbf{14.5}\pm{9.4}$ | $0.5\pm 0.6$ | $0.5\pm 0.6$ | $0.0\pm 0.0$ | $0.2\pm 0.4$ | $0.3\pm 0.4$ | $0.0\pm 0.0$ | $0.2\pm 0.3$ | $0.4\pm 0.6$ |
| puzzle 4x4 task 3 | $\mathbf{26.3}\pm{13.4}$ | $1.0\pm 0.9$ | $5.0\pm 2.7$ | $0.0\pm 0.0$ | $5.1\pm 2.8$ | $0.3\pm 0.4$ | $0.0\pm 0.0$ | $0.2\pm 0.3$ | $0.1\pm 0.3$ |
| puzzle 4x4 task 4 | $\mathbf{12.0}\pm{7.1}$ | $0.3\pm 0.7$ | $0.8\pm 0.8$ | $0.0\pm 0.0$ | $0.4\pm 0.4$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.2\pm 0.3$ | $0.1\pm 0.3$ |
| puzzle 4x4 task 5 | $\mathbf{12.3}\pm{6.2}$ | $0.7\pm 0.8$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.5\pm 0.6$ | $0.0\pm 0.0$ | $0.0\pm 0.0$ | $0.1\pm 0.3$ |
| visual cube single task 1 | $\mathbf{52.1}\pm{20.8}$ | - | $10.6\pm 7.2$ | $15.3\pm 14.6$ | - | $12.0\pm 5.6$ | - | $31.0\pm 15.0$ | - |
| visual cube double task 1 | $\mathbf{11.2}\pm{9.2}$ | - | $0.0\pm 0.0$ | $5.0\pm 2.0$ | - | $5.0\pm 3.6$ | - | $1.3\pm 1.5$ | - |
| visual scene task 1 | $\mathbf{72.4}\pm{17.7}$ | - | $32.0\pm 13.0$ | $26.0\pm 17.2$ | - | $9.0\pm 6.6$ | - | $\mathbf{74.7}\pm{22.2}$ | - |
| visual puzzle 4x4 task 1 | $\mathbf{6.0}\pm{3.2}$ | - | $0.0\pm 0.0$ | $0.0\pm 0.0$ | - | $0.0\pm 0.0$ | - | $0.0\pm 0.0$ | - |

## Appendix E Additional visualizations of latent intentions

Figure: Figure 7: Visualization of latent intentions on quadruped-jump.
Refer to caption: x7.png

We include additional visualization of latent intentions on quadruped-jump in Fig. [7](https://arxiv.org/html/2506.08902v2#A5.F7).

## Appendix F Additional Experiments

Figure: Figure 8: Evaluation on robotics datasets. InFOM outperforms the best baseline by $34\%$ when trained on top of embeddings from a fixed image encoder. See Appendix [F.1](https://arxiv.org/html/2506.08902v2#A6.SS1) for details.
Refer to caption: x9.png

### F.1 Evaluation on robotics benchmarks

To further study the pre-training and fine-tuning effects of our method on realistic datasets. Specifically, we choose the RT-1 dataset (Brohan et al., 2022), which contains $73499$ episodes of transitions. This dataset was collected by commanding a Google robot to pick, place, and move 17 objects in the Google micro-kitchens, covering a diverse set of intentions. Since collecting distinct robotics datasets for pre-training and fine-tuning is difficult, we use the entire dataset as both the reward-free pre-training dataset and the reward-labeled fine-tuning dataset. For the evaluation task, we use google robot pick coke can from the SimplerEnv (Li et al., 2024), which contains a suite of simulation tasks that efficiently and informatively complement real-world evaluations of the Google robot.

We compare against two baselines from our experiments (ReBRAC and DINO + ReBRAC) due to computational constraints, and also include a behavioral cloning (BC) baseline for reference. Our initial experiments indicate that all the algorithms (except DINO + ReBRAC) perform poorly when trained end-to-end from pixels directly. Following prior practice in latent flow matching (Rombach et al., 2022; Dao et al., 2023), we therefore pre-train a $\beta$-VAE (Higgins et al., 2017a) to encode images into a latent embedding space and then learn algorithms on top of those embeddings. For DINO + ReBRAC, we directly use the image representations learned by DINO to train the actor and the critic. We report means and standard deviations of success rates over 4 random seeds.

Results in Fig. [8](https://arxiv.org/html/2506.08902v2#A6.F8) suggest that InFOM outperforms the best baseline by $34\%$ when trained on top of embeddings from a fixed image encoder, indicating that our method can effectively fine-tune on challenging, realistic datasets with overlapping intentions.

### F.2 Variational intention inference is simple and performant

Figure: Figure 9: Comparison to prior intention encoding mechanisms. We compare InFOM to prior intention encoding mechanisms based on unsupervised skill discovery (HILP (Park et al., 2024b)) or successor feature learning (FB (Touati & Ollivier, 2021)). FB + FOM is equivalent to TD flows with GPI in Farebrother et al. (2025). We observe that InFOM outperforms prior methods on 3 out of the 4 tasks. See Appendix [F.2](https://arxiv.org/html/2506.08902v2#A6.SS2) for details.
Refer to caption: x10.png

We now conduct experiments ablating a key component in our method: the variational intention encoder. To investigate whether this framework induces a simple and performant way to infer diverse user intentions from an unlabeled dataset, we compare it to various intention encoding mechanisms proposed by prior methods. Specifically, we consider replacing the variational intention encoder with either *(1)* a set of Hilbert representations and Hilbert foundation policies (Park et al., 2024b) (HILP + FOM) or *(2)* a set of forward-backward representations and representation-conditioned policies (Touati & Ollivier, 2021) (FB + FOM), and then pre-training the flow occupancy models conditioned on these two sets of representations. Note that FB + FOM is equivalent to TD flows with GPI in Farebrother et al. (2025). We choose two tasks in the ExORL benchmarks (walker flip and quadruped jump) and another two tasks taken from the OGBench benchmarks (cube double task 1 and scene task 1), following the same evaluation protocols as in Appendix [D.3](https://arxiv.org/html/2506.08902v2#A4.SS3).

Results in Fig. [10](https://arxiv.org/html/2506.08902v2#A6.F10) indicate that InFOM can outperform prior intention encoding methods on 3 of 4 tasks, while being simpler.
Both HILP and FB capture intentions with full unsupervised RL objectives based on an actor-critic backbone. In contrast, we capture intentions by simply training an intention encoder based on a latent variable model over adjacent transitions, without relying on a potentially complicated offline RL procedure (Tarasov et al., 2023b; Park et al., 2024a).

Figure: Figure 11: Reward function structure can impose challenges. The baseline ReBRAC achieves $3.6\times$ higher performance on variants of jaco tasks with a dense reward function.
Refer to caption: x12.png

### F.3 Flow occupancy models enable faster policy learning

We then investigate whether the proposed method leads to faster policy learning on downstream tasks. We answer this question by an ablation study with a high evaluation frequency, analyzing the performance of various methods throughout the entire fine-tuning phase every 2K gradient steps. We compare InFOM to prior methods on two ExORL tasks (cheetah run and quadruped jump), including ReBRAC, CRL + IS, DINO + ReBRAC, MBPO + ReBRAC, and FB + IQL (See Appendix [D.2](https://arxiv.org/html/2506.08902v2#A4.SS2) for details of these baselines). We choose these baselines because they perform similarly to our method, helping to prevent counterfactual errors derived from the performance deviation when comparing convergence speed.

We compare different algorithms by plotting the returns at each evaluation step, with the shaded regions indicating one standard deviation. As shown in Fig. [10](https://arxiv.org/html/2506.08902v2#A6.F10), InFOM converges faster than prior methods that only pre-train behavioral cloning policies (ReBRAC) or self-supervised state representations (DINO + ReBRAC), demonstrating the effectiveness of extracting temporal information. The observation that methods utilizing a one-step transition model (MBPO + ReBRAC) or a future state classifier (CRL + IS) learn more slowly than our method highlights the importance of predicting long-horizon future events using expressive generative models. Additionally, our flow occupancy models extract rich latent intentions from the unlabeled datasets, resulting in adaptation speed similar to the prior zero-short RL method (FB + IQL).

### F.4 Learning with sparse rewards is challenging

We hypothesize that the sparse reward function on jaco tasks explains the performance gap between InFOM and baselines. To test this hypothesis, we conduct ablation experiments on jaco reach top left and jaco reach bottom right, studying whether using *dense* rewards will mitigate the performance gap. Specifically, the dense reward function is defined as $r(s,g)=-\lVert s-g\rVert_{2}$ with $g$ as the target position. To make a fair comparison, we fine-tune the ReBRAC baseline on variants of those two jaco tasks with dense reward functions, measuring the performance in the original environments. We report returns across 8 random seeds.

Results in Fig. [11](https://arxiv.org/html/2506.08902v2#A6.F11) highlight that using a dense reward function results in $3.6\times$ smaller performance gap, suggesting that the original sparse reward function imposes challenges for learning on jaco tasks. We note that Yarats et al. (2022) has also included consistent evidence for this observation, where TD3 + BC (the base algorithm for ReBRAC) performed poorly on the jaco domain (Fig. 9 of Yarats et al. (2022)).

### F.5 Importance of the behavioral cloning regularization

Figure: Figure 12: The behavioral cloning regularization in the policy loss is a key component of InFOM.
Refer to caption: x13.png

To study the effect of the BC regularizer (Eq. [8](https://arxiv.org/html/2506.08902v2#S4.E8)), we conduct experiments comparing a variant of InFOM without the behavioral cloning regularization coefficient ($\alpha=0$) to our full algorithm with domain-dependent $\alpha$ values (Table [2](https://arxiv.org/html/2506.08902v2#A4.T2)). We select the same ExORL and OGBench tasks as in Fig. [6](https://arxiv.org/html/2506.08902v2#S5.F6) (walker flip, quadruped jump, cube double task 1, and scene task 1) and report the means and standard deviations of performance over 8 random seeds after fine-tuning. Results in Fig. [12](https://arxiv.org/html/2506.08902v2#A6.F12) suggest that behavioral cloning regularization ($\alpha>0$) in the policy loss is a key component of our algorithm.

### F.6 The effect of dataset sizes

Figure: Figure 13: The effect of pre-training dataset size on InFOM. Increasing pre-training dataset sizes boosts the final performances of InFOM.
Refer to caption: x14.png

#### Pre-training dataset size.

Since we aim to predict temporally distant future states from heterogeneous data (Sec. [4.1](https://arxiv.org/html/2506.08902v2#S4.SS1)), InFOM implicitly requires a sufficiently diverse dataset for effective pre-training. To study the relationship between the size of pre-training datasets and the performance of our algorithm, we conduct ablation experiments varying the pre-training dataset size in $\{1\mathrm{K},10\mathrm{K},100\mathrm{K},500\mathrm{K},1\mathrm{M}\}$. We compare the performances of InFOM on two ExORL tasks (cheetah run and quadruped jump) after fine-tuning on the same reward-labeled dataset. We report results across 8 random seeds, following the same evaluation protocol in Appendix [D.3](https://arxiv.org/html/2506.08902v2#A4.SS3).

Results in Fig. [14](https://arxiv.org/html/2506.08902v2#A6.F14) indicate that larger pre-training datasets yield higher returns on these tasks. We conjecture that pre-training InFOM on a diverse, reward-free dataset reduces sampling of out-of-distribution future states when estimating $Q_{z}$ (Eq. [6](https://arxiv.org/html/2506.08902v2#S4.E6)) during fine-tuning, resulting in a higher final performance.

#### Fine-tuning dataset size.

We also conduct ablation experiments studying the effect of fine‑tuning dataset sizes. Specifically, we select two OGBench tasks (cube single task 2 and puzzle 4x4 task 1) and vary the size of the fine‑tuning datasets in $\{5\mathrm{K},10\mathrm{K},25\mathrm{K},50\mathrm{K},100\mathrm{K},250\mathrm{K},500\mathrm{K}\}$. Again, we aggregate the performance of InFOM over 8 random seeds, following the same evaluation protocol in Appendix [D.3](https://arxiv.org/html/2506.08902v2#A4.SS3).

Results Fig. [14](https://arxiv.org/html/2506.08902v2#A6.F14) show that increasing the fine-tuning dataset size (within the chosen range) yields consistent improvements in success rates on the OGBench tasks. Our explanation for these observations is that the size of the fine‑tuning dataset affects the accuracy of the reward prediction.

### F.7 The sufficient number of future states in the Q estimation

Figure: Figure 15: Using a sufficient number of flow future states is important. Increasing the number of flow future states ($N$) in the $Q_{z}$ estimate boosts the accuracy while reducing variance, resulting in higher final performances of InFOM. We choose $N=16$ as a balance between the accuracy, variance, and computational constraints in our experiments.
Refer to caption: x16.png

Since we use MC future states from the InFOM, intuitively, the number of future states $N$ affects the accuracy and variance of the Q value estimation (Eq. [6](https://arxiv.org/html/2506.08902v2#S4.E6)). To investigate the effect of $N$, we conduct ablation studies on a total of 8 tasks, with 4 tasks from the ExORL benchmarks (cheetah walk, walker walk, walker flip, and quadruped jump) and 4 tasks from the OGBench benchmarks (cube double task 3, puzzle 4x4 task 1, cube double task 1, and scene task 1). Below, we report returns and success rates after fine-tuning, aggregating the results over 8 random seeds.

To further investigate how $N$ affects performance on various tasks, we conduct ablation studies on a total of $\mathbf{8}$ tasks, with two more ExORL tasks (walker flip and quadruped jump) and two more OGBench tasks (cube double task 1 and scene task 1). Below, we report returns and success rates together with the results in the initial responses.

Fig [15](https://arxiv.org/html/2506.08902v2#A6.F15) suggests that, in cheetah walk and puzzle 4x4 task 1, increasing the number of flow future states yields better performance with consistent variance. In walker walk and cube double task 3, a larger $N$ does mitigate the high variance in $Q_{z}$, at the cost of increasing computation. Taken together, these results indicate that a sufficiently large number of flow future states used in $Q_{z}$ achieve more accurate estimates of Q values, while reducing the variance. In practice, our choice of $N=16$ is a balance between the accuracy, variance, and computational constraints of the estimator.

### F.8 Additional hyperparameter ablations

Figure: (a) Latent dimension $d$
Refer to caption: x17.png

We conduct additional ablation experiments on walker flip, quadruped jump, cube double task 1, and scene task 1 to study the effect of some key hyperparameters in InFOM (Table [2](https://arxiv.org/html/2506.08902v2#A4.T2)). Following the same evaluation protocols as in Appendix [D.3](https://arxiv.org/html/2506.08902v2#A4.SS3), we report means and standard deviations across eight random seeds after fine-tuning each variant.

As shown in Fig. [16(a)](https://arxiv.org/html/2506.08902v2#A6.F16.sf1), our algorithm is sensitive to the latent intention dimension $d$. Additionally, the effect of the number of steps for the Euler method $T$ (Fig. [16(b)](https://arxiv.org/html/2506.08902v2#A6.F16.sf2)) saturates after increasing it to a certain threshold ($T=10$), suggesting the usage of a common value for all tasks.

Results in Fig. [16(c)](https://arxiv.org/html/2506.08902v2#A6.F16.sf3), Fig. [16(d)](https://arxiv.org/html/2506.08902v2#A6.F16.sf4), and Fig. [16(e)](https://arxiv.org/html/2506.08902v2#A6.F16.sf5) suggest that the expectile $\mu$ can affect the performance on ExORL tasks, while having minor effects on OGBench tasks. Importantly, the KL divergence regularization coefficient $\lambda$ and the behavioral cloning regularization coefficient $\alpha$ are crucial hyperparameters for InFOM, where domain-specific hyperparameter tuning is required. As discussed in Appendix [D.4](https://arxiv.org/html/2506.08902v2#A4.SS4), we generally select one task from each domain to sweep hyperparameters and then use one set of hyperparameters for every task in that domain.

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 2
presentation: 1
questions: 1- How does the inference work ? No algorithm is provided. 

2- How do you intend to use this method for pre-trained policies ?
soundness: 3
strengths: - The idea of leveraging "intent" for multi-step generation especially in the multi-agent setting or iterative refininement with human input. 
- Extensive experiments show practicability of the method.
summary: The present work leverages flow-matching models to predict trajectories of future states for RL tasks to help the actor finding better policies.
Extensive experiments shows on gym and robotic manipulation benchmarks  shows promissing results.
weaknesses: Major

 After several reading, I fail to understand how the algorithm works. It should not be that hard to understand: 
- flow-matching to predict future states is straightforward given trajectories either conditionned or guided on observed states.
- the predicted latent variable has to be used somehow by the policy, either by boosting an pre-trained one or by training from scratch.  

 However, the main body of the paper discusses Q-functions which in the end are conditioned on the intention. This leads the reader to infer that the intention variable is infered and fed to the agent along actor trajectories. 

My wild guess: since the policy is mainly trained via a behaviorial cloning loss (see algorithm 2 in appendix) which is barely mentioned in the main body,  using the predicted state actions, the flow-matching is merely doing data-augmentation on the off-line dataset. 

I thus think that the abstract is misleading, sections 2 and 3 should be rewritten entirely to focus method rather than the intention (ironically). Algorithm 2 is rather involved and should be explained fully in the main body of the paper, all losses should be introduced clearly. 


Minor 

1- Diffentiating through an ODE solver is associated to a citation to Park et al. 2025b.  However, this is a problem already accounted for in the Normalizing flow litterature, see FFJORD https://arxiv.org/abs/1810.01367 or even Neural ODE https://arxiv.org/abs/1806.07366. Considering that Ricky Chen is a major contributor to both flow matching and Normalizing flow literature, it is odd to cite Park et al for this aspect.

2- "The deterministic nature of ODEs equips flow-matching methods with simpler learning objectives and faster inference speed than denoising diffusion models (Lipman et al., 2023; 2024; Park et al., 2025b)" is problematic for two reasons. I doubt that the training objective is simpler for flow-matching compared to Diffusion model if by that the authors mean the target function to learn or even the numerical stability of the loss. It is true however that the numerical stability of the *inference* is better for FM compared to DDPM. It is also true that the MSE-based loss of FM is more numerically stable and lighter than the KL-based loss of ODE-based Normalizing flows   such as FFJORD, see for instance https://arxiv.org/abs/2107.07232. The second reasin is that Park et al. 2025b has little to do with the statement.

---

confidence: 2
contribution: 4
initial_rating: 8
presentation: 4
questions: Is there prior work one using other generative modeling approaches, specifically diffusion models for learning occupancy models? If yes, could the authors provide some intuition on comparison between using diffusion versus flow matching for learning occupancy models?
soundness: 4
strengths: * The paper is well written and easy to follow
* The experimental results showing the proposed approach outperforms the baselines in most of the case with the gap widening on more difficult tasks
* A large number of baselines have been included and sufficient experimental detail is provided.
* I really the like the qualitative analysis of the learnt latent intention model.
summary: This work leverages recent advances in generative models in the direction of flow-matching algorithms in order to tackle the paradigm of pre-training and finetuning RL models. This is done by learning occupancy models over the state space using flow matching (as done in Farebrother et al. (2025)) while  taking into account the fact that large pretraining datasets usually contain a mixture of intentions since they are collected by different users. The proposed approach InFOM explicitly models the intention as a latent variable using a VAE which is used to condition the occupancy model. While fine-tuning, the method generates intention-conditioned Monte-Carlo estimates of  the crictic from sampled future states and then distills them into a single critic via an upper-expectile loss. The authors show that across 36 state-based and 4 image-based tasks ), InFOM matches or outperforms strong pre-train-and-fine-tune baselines, reporting a 1.8× median return gain and a 36% success-rate increase, with particularly large gains on harder manipulation domains.
weaknesses: The proposed approach builds on top of Farebrother et al. (2025) which uses flow matching to learn occupancy models, by explicitly modeling intention as a latent variable. Including an ablation / baseline comparing the downstream performance with intention conditioning of the occupancy model vs without conditioning it seems to be missing.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Q1: How well does the flow occupancy model generalize to out-of-distribution intentions? 

Q2: Is there a risk that the generative occupancy model produces unrealistic futures for novel $z$, harming policy learning?
soundness: 3
strengths: 1. Novel integration of flow matching with intention-aware occupancy modeling.

2. Strong empirical results: consistent gains across diverse domains, including challenging sparse-reward and vision-based tasks.

3. Well-motivated design choices, especially the use of SARSA-style bootstrapping and expectile-based implicit GPI to stabilize training.
summary: The paper proposes InFOM, a method for unsupervised pre-training in offline reinforcement learning that learns a latent-variable generative model of long-horizon state occupancy conditioned on inferred user intentions. By combining variational inference with flow matching, the authors model the discounted occupancy measure. During fine-tuning, they use implicit generalized policy improvement (GPI) via expectile distillation. The method shows strong empirical gains over prior pre-training approaches.
weaknesses: 1. Choice of Prior over Intentions Lacks Justification. The paper assumes a standard Gaussian prior $p(z)=N(0,I)$ for the latent intention $z$ . While common in VAEs, this choice may be suboptimal for modeling user intentions, which are often discrete or categorical (e.g., “pick”, “place”, “navigate to A”).

Suggestion: The authors should consider discrete latent variables (e.g., via Gumbel-Softmax) to improve interpretability and align better with the semantics of “intentions” in multi-task or goal-directed settings. Ablations comparing continuous vs. discrete priors would strengthen the modeling claim. The current Gaussian prior may encourage over-smoothed representations, potentially conflating distinct behavioral modes (as hinted in Fig. 4, where InFOM already shows better clustering—but could it be sharper with discrete codes?).

2.  Incorrect ELBO Formulation: The parameter $\lambda$ should be at least one, not arbitrary. In equation (3), the coefficient $\lambda$ must satisfy $\lambda \geq 1$ to guarantee that the derived expression is the lower bound.

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 3
questions: - In Appendix C.1., in the derivation of the ELBO loss, does an inequality (c) (line 1155) require $\lambda \geq 1$ for this derivation? 
- Did the newly trained reward, Q-value, and policy models use the features learnt by the occupancy model?
- The improved effect of using the behavior cloning suggests that the dataset consists of high-quality data that is worth being "cloned". What is the issue if the dataset consists only of highly suboptimal but diverse rewards? And how strongly does the diversity of the dataset influence the pre-training performance?
- What are the results on the occupancy measure pretraining performance separately, compared to a standard TD-flows without intention decoding?
soundness: 3
strengths: - The paper is very well-written and describes the whole training pipeline in great detail.
- The approach bridges the well-developed flow-matching literature with an RL pretraining and subsequent fine-tuning, which distinguishes this approach from prior work on TD-flows.
- The approach of modeling the policy intentions as latent variables seems to be fresh and interesting, since it doesn't require any additional supervision, and distinguishes this approach from Multi-Task RL.
- Strong performance on many benchmarks;
summary: The work provides a way to pre-train a general RL occupancy model using an unlabeled dataset, using a technique close to the newly introduced Temporal-Difference Flows combined with variational inference over the "intention" latent variable, which helps learn the final occupancy model. Given the pretrained flow occupancy model, the authors proposed a way to train a new policy and its Q-value using generalized policy improvement and a Q-value estimate from a reward model and an occupancy measure model. The authors provided extensive experimental validation of their pre- and post-training pipeline across various robotics benchmarks, as well as various ablation studies on the algorithmic choices.
weaknesses: - The final algorithm combines four (4) neural networks and, at first glance, looks extremely complicated, which can be prone to the accumulation of errors;


Decision: Accept (Poster)