Intention-Conditioned Flow Occupancy Models
Authors: Chongyi Zheng
Seohong Park
Sergey Levine
Benjamin Eysenbach
Venue: ICLR 2026 Poster (2026)
Keywords: reinforcement learning
flow matching
latent variable model
pre-training and fine-tuning
Primary area: reinforcement learning

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 intention. This intention increases the expressivity of our model, and enables adaptation with generalized policy improvement. We call our proposed method **intention-conditioned flow occupancy models (InFOM)**. Comparing with alternative methods for pre-training, our experiments on $36$ state-based and $4$ image-based benchmark tasks demonstrate that the proposed method achieves $1.8 \times$ median improvement in returns and increases success rates by $36\\%$.

Published as a conference paper at ICLR 2026
INTENTION-CONDTTIONED FLOW OCCUPANCY MODELS
Chongyi Zheng'
Seohong Park
Sergey Levine
Benjamin Eysenbach'
'rrinceton University
2 University of California, Berkeley
shongyizeprincetoonm
ABSTRACT
Large-scale pre-training has fundamentally changed how machine learning research is done today: large foundation models trained once, and then be used by anyone in the community (including those without data or compute resources train a model from scratch) adapt and fine-tune to specific tasks. Applying this same framework to reinforcement learning (RL) is appealing because it of- fers compelling avenues for addressing core challenges in RL, including sample efficiency and robustness. However, there remains fundamental challenge to pre-train large models the context of RL: actions have long-term dependencies, so training foundation model that reasons across fime is important. Recent ad- vances in generative AI have provided new tools for modeling highly complex distributions. In this paper, we build 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 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-conditionee flow occupancy models (InFOM). Comparing with alternative methods for pre-training, our experiments on 36 state-based and image- based benchmark tasks demonstrate that the proposed method achieves 1.8x median improvement in returns and increases success rates by 36%.
Website:
wwwponnggiiihhhnngge -zheng.ee..ee gi github..hh..c 1 fnom com/ chongyy-zhennge infom
Code:
INTRODUCTION
Many of the recent celebrated successes of ma- chine learning have been enabled by training large foundation models on vast datasets, and then adapting those models to downstream tasks. Examples include today's 's chatbots (c.g.. Gem- ini (Team et al., 2023) and ChatGPT (Achiam al., 2023)) and generalist robotic systems (e.g., To (Black et al., 2024) and Octo (Team et al., 2024)). This mme-traiingg-fin-tume has been wildly successful in fields ranging from computer vision natural language process- ing (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 open problem. What fundamentally makes the RL problem difficult 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 build foundation models for RI. often neglect these important bits, focusing
pidata)
alssteesica (ata)
Kitosue
-.I.I. Inteatiec
paradigm
Figure 1: InFOM is a latent variable model for pre-
training and fine-tuning in reinforcement learning. (Lefi) 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 for details.

---

Published as a conference paper at ICLR 2026
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 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., 2024b; Farebrother et al., 2025). World models can achieve great performance in sample efficiency (Janner et al., 2019) and generalize to diverse tasks (Hafner al., 2023; Mendonca et al., 2021), although their capacity 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- nigh-dimensional tasks can also achieve great performance in predicting future events (Sikchi et al., 2024; Barreto et al., 2018; Zheng et al., 2024b; 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 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) for pre-training in RL that simultaneously learns 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., 2024b; 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 al., 2025) to train occupancy models, enabling highly flexible modeling of occupancy measures. We call the resulting algorithm inlention-conitionne flow occupancy models (InFOM). Experiments on 36 state- bate- and 4 image- based benchmark tasks show that InFOM outperforms alternative methods for pre- vee-traiiing and fine-tuning by 1.8x 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) and enables efficient policy extraction (Sec. 5.3).
2 RELATED WORK
Offiine unsupervised RL. The goal of offline unsupervised RL to pre-train policies, value func- tions, or models from an unlabeled (reward-fre 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 al., 2024b; 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 al., 2023), which also learns a model that predicts intentions. 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 Oeilell 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 simultancously.
Unsupervised representation learning for RI. 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 learn- ing (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- -rediccedi representa- tions (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 is often (Laskin et al., 2020), though not always (Zhang et al., 2021), unclear whether the learned
We will use "successor representations," "occuppnn "occupancy measures," and "occupancy models" interchangeably.
2

---

Published as a conference paper at ICLR 2026
representations will facilitate policy adaptation. In our experiments, demonstrate that learning occupancy models enables faster policy learning.
RL with generative models. Modern generative models have been widely adopted solve RL problems. Prior work has employed autoregressive models (Vaswani et al., 2017), iterative generative models (e.g., denoising diffusion (Sohl-Dickstein al., 2015; Ho al., 2020) and flow matching (Liu et al., 2023; Lipman et al., 2023; 2024)), or autoencoders (Kingma & Welling, 2013) to model trajec- tories (Chen al., 2021; Janner 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), policies (Wang et al., 2023; Hansen-Estruch et al., 2023; Park et al., 2025b), and values (Dong et al., 2025; Agrawalla et al., 2025). We employ state-of-ther
flow-matching objective (Farebrother
et al., 2025) to model discounted state occupancy measures.
Successor representations and successor features. Prior work has used successor representa- tions (Dayan, 1993) and successor features (Barreto et al., 2017) for transfer learning (Barreto et al., 2017; 2018; Borsa 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 al., 2023; Zheng al., 2025; Jain et al., 2023; Zhu al., 2024), and goal-conditioned RL (Eysenbach et al., 2020; 2022; Zheng et al., 2024b;a). Our method is closely related to prior methods that learn successor representations with generative models (Janner 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 al. (2025). which also uses flow-man matchingg to model the occupancy measures and partly employs the generalized policy improvement (GPI) for policy extraction. Unlike Farebrother et al. (2025), which uses aeward-backward representations to capture behavioral intentions and perform GPI over a finite set of intentions, our method employs latent variable model to learn intentions (Sec. 4.2) and uses an expectile loss to perform implicit GPI (Sec. 4.4). We empirically show that these choices lead to higher returns and success rates (Sec. 5.1, Sec. 5.3).
PRELIMINARIES
We consider a Markov decision process (MDP) (Sutton et al., 1998) defined by a state space S, an action space A, initial state distribution E A(S), reward function r S - IR, a discount factor E [0. 1), and a transition distribution p S x A + (SS), where A(-) 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 Th - r(sx) without loss of generality (Tomar et al., 2024; Frans et al., 2024; Thakoor et al., 2022). In Appendix A.1, 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 al., 2020; Touati & Ollivier, 2021; Zheng et al., 2024b; Myers et al., 2024; Blier et al., 2021) that quantifies the discounted visitation frequency of different states under the policy æ. Prior work (Dayan, 1993; Janner et al., 2020; Touati & Ollivier, 2021; Zheng et al., 2024b) 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:
(s, s, = - (15))))(8)))) + ('po)aaa))) (p) ss',''), a'wa(a''l")
where (-()) denotes the Dirac delta measure centered at s.2 The discounted state occupancy measure allows us to rewrite the Q-function as a linear function of rewards (Barreto al., 2017; Touati & Ollivier, 2021; Zheng et al., 2024b; Sikchi et al., 2024):
(=(s) = I - - E. - [r(s))]
(2)
"The recursive relationship in Eq. 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., 2024b; Thakoor et al., 2022).
3

---

Published as a conference paper at ICLR 2026
The alternative (dual (Sikchi al., 2024)) definition of Q-function (Eq. 2) allows us cast the policy evaluation step as first learning a generative model ps(ss s, a) to simulate the discounted state occupancy measure of x and then regressing the estimator Q towards the average reward at states sampled from Pyy (Toussaint & Storkey, 2006; Tomar al., 2024; Thakoor et al., 2022; Zheng et al., 2024b). See Sec. 4.4 for detailed formulation.
Flow matching and TD flows. Flow matching (Lipman et al., 2023; 2024; Liu al., 2023; Albergo & Vanden- 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 How-matching methods with more stable learning objectives and faster inference speed than denoising diffusion models (Lipman et al., 2023; 2024; Verine et al., 2023). In Appendix A.2, 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) into the condi- tional flow matching loss (Eq. 10), resulting a temporal difference flow matching procedure (TD) flows) (Farebrother et al., 2025). In Appendix A.3, discuss the detailed formulations of the TD flow objective for a target policy æ. Choosing the target policy ã to be the same as the behavioral policy B, we obtain 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.
INTENTION-CONDITIONEDD FLOW OCCUPANCY MODELS
In this section, will introduce 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-condittoned 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, a') collected by the behavioral policy 3 is provided for pre-training and reward-labeled dataset Drewand {(s, r}r collected by some other policy / on a downstream task used for fine-tuning. Importantly, the behavioral policy 3 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 as in some latent space Z. In practice, these intentions can refer to desired goal images or language instructions that index the behavioral policy 3 = {3(- :, 2) 2 E 2}. Because these latent intentions are unobserved to the pre- pre-training algorithm, similar to prior work (Hausman et al., 2017; Li et al., 2017: Henderson et al., 2017). In Appendix B.2, include discussions distinguishing our problem setting from meta RL and multi-task RL. problems.
want to infer them as latent random variable Z from the offline dataset,
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 & 4.3). During fine-tuning, we first recover set of intention-cooditioned 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). Our method builds upon an assumption regarding the consistency of latent intentions.
Assumption (Consistency). The unlabeled dataset D for pre-training is obtained by executing behavioral policy following a mixture of unknown intentions 2 E Z. We assume that consecutive transitions (s, a) and (8', ,a)) share the same intention.
4

---

Published as a conference paper at ICLR 2026
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 for justifications of this assumption.
4.2 VARIATIONAL INTENTION INFERENCE
The goal of our pre-training framework is learn 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 based on consecutive transitions (s, a, s', a') using the encoder Pe (z s', aa) and predict the occupancy measures of a future state s using the occupancy models qa(sf 8, a, 2). We want to maximize the likelihood of observing future state $ starting from ataten pair (s,a) (amortized variational inference (Kingma & Welling, 2013; Margossian & Blei, 2024)), both sampled from the unlabeled dataset D following the joint behavioral distribution pd (s, (, 8f) y (8, a a))²(( (sf 8,a):
max E (s.a,,)))) ga(s a)]
2 max pp (s.a,,y, (pplt [log qa(ss | a, 2)] (DeL(pele 1s', [[)()) Pe.
(3)
where p(z) (N 1) denotes an uninformative standard Gaussian prior over intentions, À 1 denotes the coefficient that controls the strength of the KL divergence regularization. In practice, we can use any À 2 0 because rescaling the input (s, a, f), similar to normalizing the range of images from {0, ... 255} to [0, 1] in the original VAE (Kingma & Welling, 2013), preserves the ELBO. We defer the full derivation of the evidence lower bound (ELBO) and the explanation of À to Appendix C.I. Inferring the intention from the next transition (s', a²) follows from consistency assumption (Assump. 1), and is important for avoiding overfitting (Frans et al., 2024). Importantly, Pe and Ja are optimized jointly with this objective. One way of understanding this ELBO is maximizing an information bottleneck with the chain of random variables (S", A') + Z + (S, Sy). See Appendix C.1 for the connection.
We use flow matching to reconstruct the discounted state occupancy measure rather than maximizing the likelihood directly, resulting minimizing a surrogate objective:
(qa, Pesld
min Lrow (qa, Pe) + (E (Pe (s','') p()))
We use LFoO to denote placeholder for the flow matching loss and will instantiate this loss for the flow occupancy models Qa 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 unlabeled dataset D and an intention encoder (=(= 8', a,a', the goal is to learn a generative occupancy model ga(sy y 8, a, that approximates the discounted state occupancy measure of the behavioral policy conditioned on z) Rs p (xy 8, o, 2). We will use [0, 1] X x x x Z denote the time-dependent vector field that corresponds to qu. There are two desired properties of the learned occupancy models: (I) distributing the peak probability density to multiple sy. i.c., 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) 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
different intentions, i.e., (5 |
5

---

Published as a conference paper at ICLR 2026
(Eq. 11) to learn intention-conditionn flow occupancy models:
LSARSA (Ud:Pe) = () CSRRAA current flow (La, Pe) + CSARSA future flow (Uds Pe),
(5)
SARSA current How (Ud:Pe)= = (s ai al (ap a'),
[lle(t, s,, -
(2 s' ,²),
CSARSA future flw (Ud: (Ud:Pe) = B (s,a,' a )'mpp' (s..s' a'), [llva(t, 8, - Tult, 3 8 , 2)
(n..s'
2m o ,²),
Importantly, incorporating the information from latent intentions into the flow occupancy models allows us to (7) use the simpler and more stable SARSA bootstrap instead of the Q-learning style bootstrap (Eq. 11) on large datasets, (2) generalize over latent intentions, avoiding counterfactual errors. Sec. 5.2 visualizes the latent intentions, and Appendix F.2 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). We first estimate a ser of ineention-conitionns
Q functions using regression and then use those intention-conditionee Q functions to extract a policy, utilizing generalized policy improvement (GPI) (Barreto et al., 201 7)) The key idea of GPI is that, in addition to taking the maximum over the actions, 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). We address this issue by replacing the greedy "max" with an upper expectile loss, resulting in an implicit generalized policy improvement procedure.
also take the
Generative value estimation. Given a reward- labeled dataset Dreward and the pre-trained flow occu- pancy models Jd, we estimate for a fixed latent intention 2 E 2, we first sample a set N future states from the flow occupancy (1) models, $ ,.... 8 sf
Q values for downstream task. Specifically,
(N) (t) qa(s) 8, a, a),,,....) and then constructs Monte Carlo (MC) estimation of the Q function using those generative samples:
N Σ F=1
Q2(s,0)
(9))) o ~ ga(sy | s, 2,,,
(6)
where r(-) 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.11 indicate that N = 16 works effectively in our experiments. Note that we choose to sample z from the prior p(z) instead of from the posterior qa(2) s',, a'), resembling drawing random samples from a variational auto auto-encoder (Kingma & Welling, 2013). We include an ablation study in Appendix F.5, comparing the effect of fine-tuning with latents from the prior p(z) and the posterior (a(= 8', a'). In practice, we find sampling from the prior p(z) worked well in our experiments.
Implicit generalized policy improvement. We can then use those MC estimation of Q functions to learn policy by invoking the generalized policy improvement. The naive GPI requires sampling finite set of latent intentions from the prior distribution p(z), 2(1), ,.),.. (M) 201) ~ p(2) and
greedily choose one Q₂ to update the policy (Barreto et al., 2017):
max Q200 (
arg max E
smpp (s), amx((ass) ,,, (..,,(M) (M); (p()p(()
260
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 Vaqa(sy |s, a, 2). 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
We omit the dependency of Q2 on 8
(1) ,8 simplify notations.
(N)

---

Published as a conference paper at ICLR 2026
cheetah
walker
quadruped
jaco
cube single cube double
scene
puzzle 4x4
Domains for evaluation. (Left) ExORL domains (16 state-based tasks). (Right) OGBench domains Figure 2:
(20 state-based tasks and image-based based tasks).
MC estimation intention- ineentio-conditional Q functions. This approach is appealing because gradients of the 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 12 (Kostrikov et al., 2022), resulting in the following critic loss
(()) = E ((s.)) (s.a), x=p(2) [L2 (Q=(s,a) - Q(s.
where (x) theoretical soundness of this distillation step. After distilling the into a single function, we can extract the policy by selecting actions to maximize Q with behavioral cloning regularization (Fujimoto & Gu, 2021) using the actor loss
- 1p- 1(x < 0)(2² and μ E [0.5, In Appendix C.2, we discuss the intuition and
Q functions
[(T) -E, -B (s.a) (s.a), (s.a), [s) [Q(s, a²) + a log (a s)],
(8)
where 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= values. Ablation experiments in Appendix F.7 and Appendix F.12 show that this behavioral cloning regularization is important for improving the policy performance. Taken together, call the expectile Q distillation step (Eq. 7) and the policy optimization step (Eq. 8) implicit generalized policy improvement (implicit GPI).
Algorithm summary. We use neural networks to parameterize the intention encoder Pa, the vector field of the occupancy models the reward predictor Tyy the critic Q and the policy Tw. We consider two stages: pre-training and fine-tuning. of InFOM. InFOM pre-trains (I) the vector field de using the SARSA flow loss (Eq. 5) and (2) the intention encoder Pe using the ELBO (Eq. 3). Alg. 2 shows the pseudocode InFOM for fine-tuning. InFOM mainly learns (1) the reward predictor Tn via simple regression, (2) the critic Qs using expectile distillation (Eq. 7), and (3) the policy Tu by conservatively maximizing the Qy (Eq. 8). The open soen-source implementation is available online.*
Alg. 1, we summarize the pre-training process
EXPERIMENTS
Our experiments start with comparing InFOM to prior methods that first pre-train on reward-free datasets and then fine-tune 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 varia- tional intention encoder show alignment with the underlying ground-truth intentions. Our ablation experiments reveal the effect of the implicit GPT policy extraction strategy. We also include additional experiments showing InFOM enables faster policy learning during fine-tuning in Appendix F.3. Our algorithm is robust to various choices of hyperparameters (Appendix F.1 Following prior work (Park et al., 2025b), all experiments report means and standard deviations across 8 random seeds for state-based tasks and random seeds for image-based tasks.
5.1 COMPARING TO PRIOR PRE- TRA-IIIIN AND FINE-NUNING 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 based and 4 image-based tasks across diverse robotic navigation and manipulation
hhttps: ://i thub. com/ chongyi zheng/ / foom
7

---

Published as a conference paper at ICLR 2026
800
IIOO (ours
D0M0 ReSAAC
InfoNcE
1.0
InFOR [0un]
SMO e BAARAS
Moro PAb MAC
1OL
18-10 NeSMAC
XOL
600
SUCO0SS rote
ratur
0.6
400
0.4
200
0.2
0.0
choetah
walker
quaoruped
jace
tube single cube double
scene
puzzle 4x4
visual tasks
(a) 16 state-based ExORL tasks from Yarats et al. (2022). (b) 20 state-based and 4 image-based OGBench tasks from Park et al. (2025a). We average over 5 tasks for each state-based domain and average over 4 visual tasks.
We average over 4 tasks for each domain.
Figure 3: Evaluation on EXORL and OGBench tasks. We compare InFOM against prior methods that utilize various learning paradigms on task-agnostic pre-training and task-specific fine-tuning. InFOM performs similarly to, if not better than, prior methods out of the 9 domains, including the most challenging visual tasks. We report means and standard deviations over random seeds (4 random seeds for image-based tasks) with error bars indicating one standard deviation. See Table 4 for full results.
domains and compare against 8 baselines (Fig. 2). 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., 2024b), and latent skills (Touati & Ollivier, 2021; Park et al., 2024b). We defer the detailed discussions about benchmarks and datasets to Appendix D. and the rationale for choosing different baselines to Appendix D.2. Whenever possible, we use the same hyperparameters for all methods (Table 1). See Appendix D.3 for details of the evaluation protocol and Appendix D.4 for implementations and hyperparameters of each method.
We report results in Fig. 3, 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. 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 (cheet ah and quadruped), while InFOM can obtain 20x improvement on jaco, where baselines only make trivial progress (Table 4). We suspect the outsized improvement on the jaco task is because of the high-dimen state space (twice that of the other EXORL tasks (Yarats et al., 2022)) and because it has sparse rewards; Appendix Fig. 13 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 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 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), which might be mitigated by increasing the number of sampled future states (See Appendix F.1 1). In Appendix F.I, we compare InFOM against selective baselines on real robotics datasets, showing 34% improvement.
5.2 VISUALIZING LATENT INTENTIONS
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 alterna- tive 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) set of forward- -arrarddbbcwwrddd 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 e task 1 from the OGBench benchmarks.
8

---

Published as a conference paper at ICLR 2026
time
FB FOM
HILP FOM
InFOM
E-SNE
t-SNE
pick
place
Figure 4: Visualization of latent intentions. (Top) The optimal policy picks up the blue block from the left and places
right. (Bottom) Using t-SNE (Maaten & Hinton, 2008), we visualize the latent intentions inferred by the variational intention encoder InFOM, comparing against latent representations inferred by HIIPP and FB for learning FOMs. The predictions from InFOM align with the underlying intentions. See Sec. 5.2 for details and Appendix E for more visualizations.
Fig. 4 shows the optimal trajectory, where the manipulator picks the blue block from the left and then places it on the right, and 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 capable of inferring latent intentions that align with the underlying ground-truth intentions. See Appendix E for more visualizations. In Appendix F.2, we include additional experiments comparing the downstream performance between InFOM and HILP FOM and FB + FOM. Results in Appendix Fig. 9 suggest that InFOM can outperform those two baselines on of tasks.
5.3
IMPORTANCE OF THE IMPLICIT GENERALIZED POLICY IMPROVEMENT
Our final experiments study different approaches for policy optimization. We hypothesize that our proposed method is more efficient and ro- bust than other policy extraction strategies. conduct ablation exper-
800
InfO tOurs)
1.0
InfOH SH
FOM
400
0.5
test this hypothesis, iments on one task in the ExORL benchmarks (quadi crued jump) and another task taken from the OGBench benchmarks (scene taask 1), again following the evaluation protocols in Appendix D.3. 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 intentions (z(2),... ..... (M)) We choose
0.0
quadrupadl jump
Cube
Scatee
Figure 5: Comparison alternative policy extrac- tion strategies. compare InFOM alternative pol- icy extraction strategies hased on the standard general- ized policy improvement ment. Our method is 44% more performant with 8x smaller variance than the variant using the standard GPL. See Sec. 5.3 for details.
one-step policy improve-
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 88; 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.
As shown in Fig. 5, InFOM achieves significantly higher returns and success rates than its variant based one-step policy improvement, suggesting the importance of inferring user intentions.

---

Published as a conference paper at ICLR 2026
Compared with its GPI counterpart, method is 44% more performant with 8x smaller variance (the error bar indicates one standard deviation), demonstrating that the implicit GPI indeed performs relaxed maximization over intentions while maintaining robustness.
Additional experiments. In Appendix F.3, we include additional ablations showing that InFOM enables faster policy learning. Appendix F.4 ablates InFOM against variant of InFOM with set of discrete latents trained vector quantization loss, showing that the continous latent space generally performs better. In Appendix F.8, relate the diversity of the pre-training datasets their sizes. The dataset size ablations in Appendix F.9 show that using sufficient pre-training and fine-tuning data is important. Appendix F.10 studies the effects of fine-tuning on suboptimal datasets. Our hyperparameter ablations can be found in Appendix F.12.
Farebrother et al. (2025) has already discussed using alternative prior generative modeling approaches to learn the occupancy measure. Specifically, they compare flow-based occupancy models against representative generative methods, including denoising diffusion (Ho et al., 2020), VAE (Kingma & Welling, 2013; Higgins et al., 2017), and GAN (Goodfellow et al., 2014). Results in Fig. of Farebrother et al. (2025) show that flow-based occupancy models (TD2-CFM in the figure) outperform alternative generative methods in modeling the occupancy measures. For this reason, we do not include comparisons against alternative generative occupancy models distinguish contributions.
Alternative generative occupancy models.
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 consecutive state- action pairs might not always accurately capture the original intentions in the trajectories. While we empirically showed that this simple approach 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 's and FQL's implementations (Park et al., 2025a;b). We include the common hyperparameters for all the methods in Appendix Table 1, the hyperparameters for InFOM in Appendix Table 2 and Appendix Table 3, and the hyperparameters for baselines in Appendix Table 3. All the experiments were run on single NVIDIA A6000 GPU and be finished in 4 hours for state-based tasks and 12 hours for image-based tasks. We provide open-soure implementations of InFOM and all baselines https: //gggithuuu com/ chongyi manonngi-zhhngg/. zhengg innoo.
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 experi- ments. 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 Princeton University. This research was supported by ONR 400014-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 this material are those of the authors and do not necessarily reflect the views of the National Science Foundation.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint 3xXi::3333.08774, 2023.
Bhavya Agrawalla, Michal Nauman, Khush Agrawal, and Aviral Kumar. floq: Training critics via flow-matching for scaling compute in value-based arXiv preprint arXiv:2509.0680 068663 2025.
Anurag Ajay, Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofir Nachum. OPAL: Offline
primitive discovery for accelerating offline reinforcement learning. In International Confer- ence on Learning Representations, 2021. URL VE9LGwJOIIN.
net/ forum?
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 Learning Reprexentations, 2023. URL https: /tpp//ppeevvveeeee/eee net forum? id=sP1 Fo2K9DFG.
Michael Samuel Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2023. URL
https: / openreviewwweeeeeeeeeeeeee net. forum? wedddi11geBbBBB
Alexander A. Alemi, Ian Fischer, Joshua V. Dillon, and Kevin Murphy. Deep variational information bottleneck. In International Conference on Learning Representations, 2017. URL https:
///ppnrevieww...////////o net httpr//veee.eeet/oooorrddddd-dae http:/ooommm--ii/oeeeeeeoo-
Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, and
François Fleuret. Diffusion for world modeling: Visual details matter in atari. Advances in Neural Information Processing Systems, 37:58757-588791 2024.
David Barber and Felix Agakov. The im algorithm:
variational approach information maximiza- tion. Advances in neural information processing systems, 16(320):201, 2004.
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.
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
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arxiv:1308..444 2013.
Kevin Black, Noah Brown, Danny Driess, Adnan Esmail, Michael Equi, Chelsea Finn, Niccolo Fusai, Lachy Groom, Karol Hausman, Brian Ichter, et al. No: A sssionlanguaage-ctiin flow model for general robot control. arXiv preprini arXiv:2410.2414 2.161..246..... 2024.
Léonard Blier, Corentin Tallec, and Yann Ollivier. Learning successor states and goal-dependent values: A mathematical viewpoint. arXiv preprint arXiv:2101.07123, 2021.
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 1812::2//2.07666 2018.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
/ttp://gittie com/ maxxmmm// mi/jax.
David Brandfonbrener, Will Whitney, Rajesh Ranganath, and Joan Bruna. Offline rl without off pol-cy evaluation. Advances in neural information processing systems, 34:4933 0944933-4446, 2021.
11

---

Published as a conference paper at ICLR 2026
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.
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.
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/
Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows on discrete state-spaces: Enabling multimodal flows with applications protein co-design. In International Conference on Machine Learning, pp. 5453-5512. PMLR, 2024.
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.
Boyuan Chen, Chuning Zhu, Pulkit Agrawal, Kaiqing Zhang, and Abhishek Gupta. Self-supervised reinforcement learning that transfers using random features. Advances Neural Information Processing Systems, 36:56411-56433 166666, 2023.
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-15099 2021.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Quan Dao, Hao Phung. Binh Nguyen, and Anh Tran. Flow matching in latent space. arXiv preprint vrXiv:2307.08695 2023.
Peter Dayan. Improving generalization for temporal difference learning: The successor representation. Neural computation, 5(4):613-624 1993.
Jacob Devlin, Ming-Woi Weig 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 (long and short papers), pp. 4171-4186, 2019.
Zihan Ding, Amy Zhang, Yuandong Tian, and Qinqing Zheng. Diffusion world model. arXiv e-prints, pp. ar02 2024.
Perry Dong, Chongyi Zheng, Chelsea Finn, Dorsa Sadigh, and Benjamin Eysenbach. Value flows. arXiv preprint 2r5 100.0760 2025.
Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. RIa: Fast reinforcement learning via slow reinforcement learning. arXiv preprint drXiv:6611.02773 2016.
Rick Durrett. Probability: theory and examples, volume 49. Cambridge university press, 2019.
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.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity all you need: Learning skills without reward function. In International Conference Learning Representations, 2019. URL https:://// https::/openrevviiieeee/ opwnrevieweeeeeeeeeeeee httpr/iwww..eeeeeeeeeeeeeee formm//ormmme S3x6
Benjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine. C-learning: Learning to achieve goals via recursive classification. arXiv preprint arX1v:2011.08909, 2020.
12

---

Published as a conference paper at ICLR 2026
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- 2022.
Jesse Farebrother, Matteo Pirotta, Andrea Tirinzoni. Remi Munos, Alessandro Lazaric, and Ahmed Touati. Temporal difference flows. In ICLR 2025 Workshop on World Models: Understanding, Mod- elling and Scaling, 2025. URL https: wpwprevveeee...eeeeeee net. www.rmmiddinn..pp 233.
Kevin Frans, Seohong Park, Picter Abbeel, and Sergey Levine. Unsupervised zero-shot reinforcement learning via functional reward encodings. In International Conference on Machine Learning, pp. 13927-13942. PMLR, 2024.
Kevin Frans, Danijar Hafner, Sergey Levine, and Picter Abbeel. One step diffusion via shortcut models. In The Thirteenth International Conference on Learning Representations, 2025. URL
https: ftttp::/cpeereeee wpwpreveeeeeeeeeeeeeeeeeee .nett/eeweeeeeeeee.......e..eeee hwt.:/eet/oouteaeeeeeeee [orumtid-=11266nnx.s
Scott Fujimoto and Shixiang Shane Gu. A minimalist approach offline reinforcement learning. Advances in neural information processing systems, 11-20132-20143 2021.
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.
Dibya Ghosh, Chethan Anand Bhateja, and Sergey Levine. Reinforcement learning from passive data via latent intentions. In International Conference on Machine Learning, pp. 1 11321-11339. PMLR, 2023.
Ian J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances neural information processing systems, 27, 2014.
Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, and David Duvenaud. Scalable reversible generative models with free-form continuous dynamics. International Conference on Learning Representations, 2019, URL https: pettps://openrevieee .ne.eewwweeeeeeeeeeeeeeeeee.eee. net forum ? rid=r0xgknCcK7.
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-21221 2020.
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 machine learning, pp. 1861-1870. Pmlr, 2018.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.044 2023.
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: 05030, 2019.
Philippe Hassen-Estruch, Ilya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. Idql: Implicit leaannnigg actor-critic method with diffusion policies. arXiv preprint arXiv:2304. 10573, 2023.
Karol Hausman, Yevgen Chebotar, Stefan Schaal, Gaurav Sukhatme, and Joseph Lim. Multi-modal imitation learning from unstructured demonstrations using generative adversarial nets. Advances in neural information processing systems, 30, 2017.
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 patlern recognition, pp. 9729-9738, 2020.
13

---

Published as a conference paper at ICLR 2026
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/CVE conference on computer vision and pattern recognition, pp. 16000-16009, 2022.
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 709.06683, 2017. URL https: //atts///apti /pp.i hatpp.semamnne semanti cscholas ar. org/
CorpusID:7079522)
Dan Hendrycks and Kevin Gimpel. Gaussian error 1rXi::1606.08455.. 2016.
linear units (gelus).
arXiv preprint
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, 2017. URL https::/ http///pppeeeeeeeeeeeeeeeeeeeeee net wow.oooionuyylaa.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilisticc models. Advances in neural information processing systems, 33:6840-6851 2020.
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:5 1491-5154,, 2023.
Amav Kumar Jain, Lueas Lehnert, Irina Rish, and Glen Berseth. Maximum state entropy exploration using predecessor and successor representations. Advances in Neural Information Processing Systems, 36:49991-500199 2023.
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.
Michael Janner, Igor Mordatch, and Sergey Levine. Gamma-models: Generative temporal difference learning for 1735, 2020.
prediction. Advances in neural information processing systems, 33:
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances neural information processing systems, 34:1273-1286, 2021.
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.
Scott Jeen, Tom Bewley, and Jonathan Cullen. Zero-shot reinforcement learning from low quality data. Advances in Neural Information Processing Systems, 37:16894-1666 2024.
Jaekyeom Kim, Seohong Park, and Gunhee Kim. Constrained gpi for zero-shot transfer in reinforce- ment learning. Advances in Neural Information Processing Systems, 35:4585-4597, 2022.
Junsu Kim, Seohong Park, and Sergey Levine. Unsupervised-to-nnliee reinforrereene learning. arXiv preprint 4008 14785, 2024.
Diederik Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980. 2014.
Diederik Kingma and Max Welling. Auto-encoding variational bayes. 1312v 14, 2013.
arXiv preprint
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit g-learning. In International Conference on Learning Representations, 2022. URL https: //openeeview....c/ wtt.:eteoooommmma--22223390
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in neural information processing systems, 33:1 19911199111 2020.
14

---

Published as a conference paper at ICLR 2026
Nathan Lambert, Kristofer Pister, and Roberto Calandra. Investigating compounding prediction emrors in learned dynamics models. arXiv preprint vrXiv:2203009631 2022.
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.
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 urXiv:2405.05941 2024.
Yunzhu Li, Jiaming Song, and Stefano Ermon. Infogail: Interpretable imitation learning from visual demonstrations. Advances in neural information processing systems, 30, 2017.
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 Repre- sentations, 2023. URL https: httpttps///peeeeviieeeee /pppreveeeeeeeeeeeeeeeeeee .ww.iiwwweeeeeeee........e..ee Eorum? L==
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:24/2.06264, 2024.
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: http///ppppeeeeeeeeee/ee//e net http://wwwweeeettteeeeeee
Jiasen Lu, Dhruv Batra, Devi Parikh, and Stefan Lee. Vilbert: Pretraining task agnkstiooiic visiolinguistic representations for iision-and-languag tasks. Advances in neural information processing systems, 32, 2019.
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:
//oppnreviewwcc////// net
Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 95Nv):2579-2205, 2008
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.1108%, 2017.
Charles Margossian and David M Blci. Amortized variational inference: When and why? In Uncertainty in Artificial Intelligence, pp. 2434-2449. PMLR, 2024.
Bogdan Mazoure, Benjamin Eysenbach, Ofir Nachum, and Jonathan Tompson. Contrastive value learning: Implicit models for simple offline RI. In 7th Annual Conference on Robot Learning, 2023. URL https:
/ttttps://operrvvi..e http///peevvveeeeeeeeeee net forum? d=oqo 6bJy.
Pietro Mazzaglia, Tim Verbelen, Bart Dhoedt, Alexandre Lacoste, and Sai Rajeswar. Choreographer: Learning and adapting skills in imagination. arXiv preprint arXiv:22/1.133.. 2022.
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.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller. Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 5 518(7540):529-533 2015.
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.
15

---

Published as a conference paper at ICLR 2026
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.
Mark Nemecek and Ronald Parr. Policy caches with successor features. In International Conference Machine Learning. pp.
PMLR, 2021.
Tianwei Ni, Benjamin Eysenbach, Erfan Seyedsalehi, Michel Ma, Clement Gehring, Aditya Mahajan, and Pierre-Luc Bacon. Bridging state and history representations: Understanding self- self-predictive rl. arXiv preprint ar arXiv:240I. 008898, 2024.
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.
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 O. In 2024 IEEE International Conference on Robotics and Automation (ICRA), pp. 6892 09026993..15 IEEE, 2024.
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.
Seohong Park, Dibya Ghosh, Benjamin Eysenbach, and Sergey Levine. Hiql: Offline goal-conditioned with latent states as actions. Advances in Neural Information Processing Systems, 36:34866 34891, 2023a.
Seohong Park, Oleh Rybkin, and Sergey Levine. Metra: Scalable unsupervised with metric-aware abstraction. arXiv preprint arXiv:2310.08887, 2023b.
Seohong Park, Kevin Frans, Sergey Levine, and Aviral Kumar. value learning really the main bottleneck in offline rI? arXiv preprint arXiv:2406.09322 2024a.
Seohong Park, Tobias Kreiman, and Sergey Levine. Foundation policies with hilbert representations. In International Conference on Machine Learning, pp. 39737-39761. PMLR, 2024b.
Seohong Park, Kevin Frans, Benjamin Eysenbach, and Sergey Levine. OGBench: Benchmarking offline goal-conditioned RL. In The Thirleenth International Conference on Learning Representa- tions, 2025a. URL https://p// ht/p//ppeeviiieeeeeee/ee net/ fww.mmrrmiiooo 9m00-M99222jg11...
Seohong Park, Qiyang Li, and Sergey Levine. Flow q-learning. arXiv preprint arXiv:2502.02538. 2025b.
Karl Pertsch, Youngwoon Lee, and Joseph Lim. Accelerating reinforcement learning with learned skill priors. In Conference on robot learning, pp. 188-204, PMLR, 2021.
Jan Peters and Stefan Schaal. Reinforcement learning by reward-weighted regression for operational space control. Proceedings of the 24th international conference on Machine learning, pp. 745-750, 2007.
Jan Peters, Katharina Mulling, and Yasemin Altun. Relative entropy policy search. Proceedings of the AAAI Conference on Artificial Intelligence, volume 24, pp. 1607-1612, 2010.
Vitchyr Pong, Ashvin V Nair, Laura M Smith, Catherine Huang, and Sergey Levine. Offline meta-reiinorreement learning with online self-supervision. In International Conference on Machine Learning, pp. 17811-111 1-7829. Pmlr, 2022.
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.
16

---

Published as a conference paper at ICLR 2026
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, al. Improving language understanding by generative pre-training. 2018.
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. PmLR, 2021.
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy meta-reinforrcement learning via probabilistic context variables. In International conference on machine learning, pp. 5331-5340. PMLR, 2019.
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 htttps:///ppeeee /t/p//ppeeeeeeeeeeeeeeeeeeeeeeee wwt.viiwweee/eeeee..ee www.eeiwwwweeetrueeeee id md-mmpa
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 confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Gavin A Rummery and Mahesan Niranjan. On-line Q-learning using connectionist systems, vol- ume 37, University of Cambridge, Department of Engineering Cambridge, UK. 1994.
Andrew Michael Saxe, Yamini Bansal, Joel Dapello, Madhu Advani, Artemy Kolchinsky, Bren- dan Daniel Tracey, and David Daniel Cox. On the information bottleneck theory of deep learning. In International Conference on Learning Representations, 2018. URL https: // //openreview .nt. wwtt/eernnteeeeeee
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.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 170.007077. 2017.
Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-efficient reinforcement learning with self-predictive representations. arXiv preprint arXiv:2007.05999, 2020.
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.
Archit Sharma, Shixiang Gu, Sergey Levine, Vikash Kumar, and Karol Hausman. Dynamics-aware unsupervised discovery of skills. arXiv preprint arXiv: 1907.01677 2019.
Harshit Sikchi, Qinqing Zheng, Amy Zhang, and Scott Niekum. Dual RL: Unification and new meth- ods for reinforcement and imitation learning. In The Twelfth International Conference Learning Representations, 2024. URL https: ht/p//ppeevviieeeee net forum iidxtt 9Bu6 6rqv.
Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with context- based representations. In International conference machine learning, pp. 9767-9777. PMLR. 2021.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference machine learning, pp. 2252222666 pmlr, 2015.
Yang Song, Jascha Sohl-Dickstein Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. ccore-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021. URL htttttss://ppeeeeeeeeeeeee //tp///eeeeeeee.//////e/e net forum? Id=PxTIG12RRHS
17

---

Published as a conference paper at ICLR 2026
Richard S Sutton, Andrew G Barto, et al. Reinforcement learning: An introduction, volume 1. MIT press Cambridge, 1998.
Erik Talvitie. Model regularization for stable sample rollouts. In UAI, pp. 780-789, 2014.
Denis Tarasov, Vladislav Kurenkov, Alexander Nikulin, and Sergey Kolesnikov. Revisiting the minimalist approach offline reinforcement learning. Advances in Neural Information Processing Systems, 36: 1592-11620, 2023a.
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.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmalcki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv:1801.0069) 2018.
Gemini Team, Rohan Anil, Sebastian Borgeaud, Jean- Sann-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. 1805, 2023.
Octo Model Team, Dibya Ghosh, Homer Walke, Karl Pertsch, Kevin Black, Oier Mees, Sudeep Dasari, Joey Hejna, Tobias Kreiman, Charles Xu, et Octo: An open-source generalist robot policy, arXiv preprint
2024.
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-21301 PMLR, 2022.
Naftali Tishby, Fernando C Pereira, and William Bialek. The information bottleneck method. arXiv preprint physics/0004057, 2000.
Manan Tomar, Philippe Hansen-Estruchh Philip Bachman, Alex Lamb, John Langford, Matthew E Taylor, and Sergey Levine. Video occupancy models. arXiv preprint arXiv:2407.09533. 2024.
Ahmed Touati and Yann Ollivier. Learning one representation optimize all rewards. Advances in Neural Information Processing Systems, 34:13-23, 2021.
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: // //oopprreveww.. net// forum?id-WWEappDDQ
Marc Toussaint and Amos Storkey, Probabilistic inference for solving discrete and continuous state markov decision processes. Proceedings of the 23rd international conference on Machine learning, pp. 945-955999 945-952, 2006.
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 wrXiv:2307.09288 2023.
Khadichabonu Valieva and Bikramjit Banerjee. Quasimetric value functions with dense rewards. arXiv preprint urXiv:2409.087244 2024.
Aaron Van Den Oord, Oriol Vinyals, al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforeement learning with double q- learning. Proceedings of the AAAI conference on artificial intelligence, volume 30, 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones. Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017
18

---

Published as a conference paper at ICLR 2026
Alexandre Verine, Benjamin Negrevergne. Yann Chevaleyre, and Fabrice Rossi. On the expressivity of bi-lipschitz normalizing flows. In Asian Conference on Machine Learning, pp. 1054-1069. PMLR. 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. Conference on Robot Learning. pp. 1723-1736. PMLR, 2023.
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.
Zhendong Wang, Jonathan Hunt, and Mingyuan Zhou. Diffusion policies as expressive policy class for offline reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023. URL https: wpwpp/veeeeeeeeeeeeeeeeee net/ /aee/foormmte
Christopher JCH Watkins and Peter Dayan. Q-learning. Machine learning, 8:279-292, 1992.
Tete Xiao, Ilija Radosavovic, Trevor Darrell, and Jitendra Malik, Masked visual pre-training for motor control. arXiv preprint arXiv:2203. 06173, 2022.
Denis Yarats, David Brandffonbrener, 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: www.oopeevvieeeeee net/ forum? ? d=Su- 4 125.
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. Advances neural information processing systems, 33: 5824-5836 2020.
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. 2561 1-25.35. PMLR, 2022.
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.
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: /ttp///ppp/r/e...c/////// net/ forum? id= 2FCwDKRREU.
Chongyi Zheng, Benjamin Eysenbach, Homer Rich Walke, Patrick Yin, Kuan Fang. Ruslan Salakhut- dinov, and Sergey Levine. Stabilizing contrastive RL: Techniques for robotic goal reaching from offline data. In The Twelfth International Conference on Learning Representations, 2024a. URL
https:/ /wwp//eeeewwweeeeeee/e net forum? X
Chongyi Zheng, Ruslan Salakhutdinov, and Benjamin Eysenbach. Contrastive difference predictive coding. The Twelfth International Conference on Learning Representations, 2024b. URL https:
//openreview. net worum wetee//torddddaaaaeeeeeeeeaeeeee -at.diddaaaaa-------
Chongyi Zheng, Jens Tuyls, Joanne Peng, and Benjamin Eysenbach. Can MISL fly? analysis and ingredients for mutual information skill learning. The Thirteenth International Confer- Learning Representations, 2025. URL https:/ /wwpnevvewwew.o///////p/. net forum?id= XoIeVdFO7U.
Zhiyuan Zhou, Andy Peng, Qiyang Li, Sergey Levine, and Aviral Kumar. Efficient online re- inforcement learning fine-tuning need not retain offline data. In The Thirteenth International Conference on Learning Representations, 2025. URL https: /ttp:////pee/eeeeeee net forum? 1d=HNOCYZbAPW.
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.
19

---

Published as a conference paper at ICLR 2026
Algorithm 1 Intention Intention-Condie Conditioned Flow Occupancy Model (pre-training).
1: Input The intention encoder Po, the vector field eyg the target vector field eg. the policy Two and the reward-free dataset D.
2: for each iteration do
3: Sample a batch of as,,a, a') D}-
4: Sample a batch of (c~ (N~N((, π)} and a batch of {t UNIF((O, 1))).
5: Encode intentions {a ~ Po(= ss, a')} for each (5', V SARSA flow occupancy model loss.
6:
s + (1 -1-28 t)e + ts
7:
8 EulerMethod oz. e, s' "e',d'22) 2),
(1 (1 (1- (1- t.g.
8:
CSARSA current flow (0, Φ) + ((s,a,,,),...)))
| ,,, s, a, - (s - =
9:
CAARSA
(0, + E E,,,,,,,1...557)
3, - op(1, 57, 8', a', 2)13
10:
CSARSA flow (0, 3) + (1 - Ceurrem (0, 6) Ltunur (0, 6).
b Eq.
V Intention encoder loss.
11:
CELBO 6) <- LSARSA now (0, Φ) + (s DEL (2 s s',a') NV r))-
b Eq. 4
V (Optional) Behavioral cloning loss.
12:
(Bc(w) &- (Eas.a)) [log Tu (a })]]-
13:
Update the vector field and the intention encoder & by minimizing (e, Φ).
14:
Update the policy by minimizing (wc(
15:
Update the target vector field using the Polyak average of e.
16:
Return Uo Pv. and Tw..
A PRELIMINARIES
A.1 VALUE FUNCTIONS AND THE ACTOR-CRITIC FRAMEWORK
The goal of RL is to learn a policy æ (AA) that maximizes the expected discounted return (T) = Er~w(rr) [h r'r.]. where a trajectory sampled by the policy. We will use 3: S + ((AA) to denote the behavioral pol- icy. Given a policy и, we measure the expected discounted return starting from state-action
Algorithm 3 Euler method for solving the flow
ODE (Eq. 9).
1: Input The vector field and the noise e. (Op- tional) The number of steps T with default T 10.
2: Initialize t = 0 and x² =
3: for each step f = 0, 1,.... 1
pair (s, a) and a state as the (unnormalized)
4: x1+1 +- x²' + ott/
Q-function and the value function, respectively:
5: Return f= xT
(", BYwa(T)
rh Su = s,
("(s) = E Bans(als) 1Q" a))
A=0
Prior actor-critic methods (Schulman et al., 2015; 2017: Haarnoja al., 2018: Fujimoto et al., 2018; Kumar et al., 2020; Fujimoto & Gu, 2021) typically maximize the RL objective J(x) ) by (1) learning an estimate Q of QF via the temporal difference (TD) loss (policy evaluation) and then (2) improving the policy æ by selecting actions that maximizes Q (policy improvement):
gl++1 f- arg max E (s.a. (o)p (sann))) [(((s,a) -
Q
k+1 arg max (s) (()t)(s,a)))
where ls indicates the number of updates and 3 is the behavioral policy representing either a replay buffer (online RL) fixed dataset (offline RL).
20

---

Published as a conference paper at ICLR 2026
Algorithm 2 Iivention-Conditiin Flow Occupancy Model (fine-tuning).
1: Input The intention encoder Ps. the vector field Ug. the target vector field uy the reward predictor Typ the critic u, the policy Tw (random initialization or initialized using aw from Alg. 1), and the reward-labeled dataset Drewand-
2:
for each iteration do
3: Sample a batch of ((s, a, r,,'', ~ D Dreward )-
4: Sample a batch {E (0, 1)} and a batch {t UNIF((O, 1))}.
5: Sample prior intentions {= ~ p(2)].
6: Sample batch of {(e("). (())) (N, ())).
V SARSA flow occupancy model loss and intention encoder loss.
7: CELBO (0, 1 as in Alg. 1.
v Reward predictor loss. -
8: Rewardd (1) < (s.r)
V Critic loss.
9:
f Euler (va, (+)), 88 a, 2) (Alg. 3) for each (s, a, 2, (.).
(0)
Q2 (8, - (I-7 Σ 85 (3)
10:
D Eq.
11:
CCmue(W) t E (s.2,2.86) ,r. ,7) 1700 - Qu(s,
>> Eq. 7
V Actor loss.
12:
Lxstur (w) + -B(,,a), waw (amls) (s, a²) + a log (a [)].
b Eq.
13:
Update the vector field and the intention encoder & by minimizing LFLBO Φ).
14:
Update the reward predictor n, the critic vi, and the policy by minimizing LReward (1).
Cerue (V), and Cacuo( (w) respectively.
15:
Update the target vector field 0 using the Polyak average of O.
16:
Return Pv. Tn. Qs, and Tw.
A.2
FL.OW MATCHING
The goal of flow matching methods is to transform simple noise distribution (e.g., d-dimensional standard Gaussian) into a target distribution Px over some space X' C Rd that we want to approximate. Specifically, flow matching uses time-dependent vector field : [0, 1] Rd Bd to construct dime-dependent diffeomorphic flow Φ [o, 1] x Rd + R (Lipman et al., 2023; 2024) that realizes the transformation from single noise generative sample é. following the ODE
d dt o(t,c) v(t e(t, o(0,e) =, e)1,c) =.
We will use to denote time step for flow matching and sample the noise from a standard Gaussian distribution N (0, /) throughout our discussions.? Prior work has proposed various formulations for learning the vector field (Lipman 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).
(V) EE
UNIF ar-px (nn
[lle(t, (x (x-<))13 c)
(10)
where UNIF([O, 1]) the uniform distribution over the unit interval and x² tx + (1 f)e is linear interpolation between the ground-trui truth sample and the Gaussian noise Importantly, we can generate a sample from the vector field v by numerically solving the ODE (Eq. 9). We will use the Euler method (Alg. 3) as our ODE solver following prior practice (Grathwohl et al., 2019; Chen al., 2018; Lipman al., 2023; Liu et al., 2023; Park al., 2025b; Frans et al., 2025).
Inn theory, the noise can be drawn from any distribution, not necessarily limited to Gaussian (Liu et al., 2023).
21

---

Published as a conference paper at ICLR 2026
A.3 TEMPORAL DIFFERENCE FLOWS
Given a policy T, prior work (Farebrother et al.. 2025) models the occupancy measure pa by optimizing the vector field v: 1] x S x S x A + S using the following loss:
LTD Row (v) = (1 - 7) current flow (v) + 7LTD future flow (lo((o) (v)
(11)
LTD current flow (v) = (0.
[lle(t, s", s. s",,t, - (
(s.a)rpp (s,a)
LTD future flow
(v) E (=UN))((.1))
(10((,87,88 o(l, 8 8 a')1
(5,p,)")""" (s,a,")) aam's((a'ss")
where (s, a) and p a, s') denote the joint distribution of transitions, s'= is + (1 t)e is a linear interpolation between the current state and the noise c, and ō denotes the Polyak 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 8, by applying the Euler method (Alg. 3) to ū at the next itate-action pair (s', a'), where a' is sampled from the target policy of interest, and the noisy future state s = f + (1 t)e a linear interpolation between this future state 8y and the noise e. Intuitively, minimizing LTD current flow reconstructs the distribution of current state while minimizing LTD future flow bootstraps the vector field at a noisy target future state s similar Q-learning (Watkins & Dayan, 1992). Choosing the target policy π be the same as the behavioral policy 3, we obtain SARSA (Rummery & Niranjan, 1994) variant of the loss optimizing the SARSA flows. We call the loss in Eq. 11 the TD flow loss and use the SARSA variant of it to learn generative occupancy models.
B FURTHER DISCUSSIONS ON THE PROBLEM SETTING
B. THE CONSISTENCY ASSUMPTION ON INTENTIONS
We now discuss the reason for making the consistency assumption (Assumption 1) on latent intentions. Since we use heterogeneous behavioral policy to collect the unlabeled dataset, each unknown user intention indexed their own behavioral policy 3 Z + ((A). The key observation is that the occupancy measure of each intention-cone intention-conditioned conditiooed behavioral policy follows its own Bellman equations (Similar to Eq. 1):
(s 2) = (1-7)0.(33))
[p) [s),d'=))] a' d'~3(a'13333
suggesting that the same latent 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). We note that prior work (Touati & Ollivier, 2021) also adapts the same formulation 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 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 aoaard-lablld data during the meta-training phase, where task-specific rewards provide supervision for policy learning. In contrast, during pre- trainingg 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 al., 2020). During pre-training, (I) InFOM does not have access to task descriptions 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.
'The TD flow loss is called the CDM-CFM loss in Farebrother al. (2025), and rename it for simplicity.
22

---

Published as a conference paper at ICLR 2026
C THEORETICAL ANALYSES
C.I 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', a') and future state sf sampled from the same trajectory following the behavioral joint distribution p² s,, s', = pd (s) s)p a)p(s' (. s'). We use (s', a') to encode the intention by the encoder Pe(z | a) and (s, a, Sz. 2) to learn the occupancy models qa(s) 8, a,a, a,,, employing an ELBO of the likelihood of the prior data:
(x,(,))))))) lloga((s s,
= Ex ( sa))
[log qa(s) s, a)]
= =sp
log Ex(z) [aa(s s, a, =]]
@ Eps
log log Bps)) qa(55 Is a, Pa(= 8', ,')' Pe(2 s')
( E Fe a() (og 18 =)]
(RL(p | &') ) Il [()))]
(s, an [ (sler llog 8,,, ADRI (po 1s', [[)]
-ELBO qa),
where in (a) we introduce the amortized variational encoder (= s', .a). in (b) we apply the Jensen's ' inequality (Durrett, 2019), and in (c) introduce coefficient À 2 control the strength of the KI. divergence regularization. In practice, we can use any À 2 0 because rescaling the input (s, a, sy), similar to normalizing the range of images from {0...- 255} [0, 1] in the original VAE (Kingma & Welling, 2013), preserves the ELBO. Formally, following prior work (Higgins et al., 2017), maximizing this ELBO can also be interpreted as an optimization problem that simultaneously predicts future states while penalizing the intention encoder:
max (e,,,,, ,') (xp((s))t [log ya(sy 88, a, a)=
s.t. (pI( &,,a)) II p(=)) S const..
Pe.fd
Rewriting this constrained optimization problem as the Lagrangian produces
Ep a²) (x ,a) llog ga(s) s, 2)] ADeL(pe s', s',a') p(2)),
where we introduce coefficient À 2 0 to control the strength of the KI. divergence regularization.
Alternatively, the constrained optimization problem can also be cast as variational lower bound on an information bottleneck with the chain of random variables (S", A') + Z + A, (Tishby et al., 2000; Alemi al., 2017; Saxe al., 2018):
(S, XI AS',
([S,A,S5:: (s ( [DKL (po(= &') || po()))
(S, - AI (s) DKL(P s a') p()))]
E (saasawa
(s, =)] AEp (D (81, [[)( +
p=(=''',a')
E E (naasyy a') [logqa(s} s, s.a,2)) ABpr(sa (Pa s''a') p()))) + const.
(a(=)('a')
where in (a) we use the definition of (S' A': Z) and (2) is the marginal distribution of latent intentions defined as Pe(2) = f (s a (z s', d'd''ds' da', in (b) apply the non-negative property of the KL divergence (2) || p(2)), in (c) we apply the standard variation lower bound of the mutual information (Barber & Agakov, 2004; Poole al., 2019) to incorporate the
23

---

Published as a conference paper at ICLR 2026
decoder (occupancy models) qa(s, a, 8 2), and in (d) we choose the variational decoder to satisfy log (s, a, s 2) log (s. a) + log / constant. Therefore, the lower bound in Eq. can also be interpreted as maximizing the information bottleneck (5,, A, Sy: Z) (5" A'; Z) with À 2 0.
2) and consider the entropy (S, A,,,) Sy)
C.2 INTUTIONS AND DISCUSSIONS ABOUT THE IMPLICIT GENERALIZED POLICY IMPROVEMENT
The intuition for the expectile distillation loss (Eq. 7) is that the scalar function Q(:,)) 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(=) when μ 1 (See Sec. 4.4 in Kostrikov et al. (2022) for details). Therefore, given an infinite amount of samples (N + 0o) and an expectile μ - 1, the Q converges to the greedy value functions:
(s, = max
zmp((>) (1 (1 Budte/ha.)lrsss))
If further assume that the flow occupancy models optimal, i.e., (s s, a, (s 8, a, 2), then the optimal Q corresponds to greedy value function under the behavioral policy B:
Q" (s, a)
max (s, x>p(2)
Unlike Q-learning, which converges to the optimal Q-function sequentially (Watkins & Dayan, 1992; Sutton et al., 1998), the implicit GPI proposes new policy that is strictly no worse than the set of policies that correspond to each Q₂ in parallel (See Sec. 4.1 in Barreto et al. (2017) for further discussions). Unlike one-step policy improvement (Wang et al., 2018; Brandfonbrener al., 2021; Peters & Schaal, 2007; Peters et al., 2010), implicit GPI converge to the optimal policy for 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 step policy improvement using a flow occupancy model 4a($; a) that is not conditioned on latent intentions. This flow occupancy model captures discounted state occupancy measure of the (average) behavioral policy. After training the flow occupancy model, FOM + one-step PI fits Q function and extracts behavioral-regullrizel policy:
Q f arg min
1-7 B
((Q(s r(s)))
Q
<- arg max E (s.a)) [Q(s, a²) + a lo s)]
Intuitively, the first objective fits the behavioral Q function based on the dual definition (Eq. 2). and the second objective trains policy maximize this behavioral Q function, invoking one- step policy improvement. While this simple objective sometimes achieves strong performance on some benchmark tasks (Brandfonbrener al., 2021; Eysenbach al., 2022), it does not guarantee convergence to the optimal policy due to the use of a behavioral value function.
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. 2).
We use 16 state-based tasks from the EXORL (Yarats et al., 2022) benchmarks based on the DeepMind Control Suite (Tassa al., 2018). These tasks involve controlling four robots (cheetah, walker, quadruped, and jacc) to achieve different locomotion behaviors. For each domain, the specific tasks cheetah (run, run backward, wall k, wa 1 backward), wa ker
EXORL.
24

---

Published as a conference paper at ICLR 2026
fwalk, run, stand, flip), quadruped (run, jump, stand, walk), jaco (reach top left, reach top right, reach bottom left, reach bottom x right]. J. For all tasks in cheetah, walker, and quadruped, both the episode length and the maximum return 1000. For all tasks in jaco, both the episode length and the maximum return are 250. Following prior work (Park 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 al., 2024), we will use 5M unlabeled transitions collected by some exploration methods (e.g., RND (Burda et al., 2019)) for pre- 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 ja 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 s 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 (pick and place cube diagonally), task (pick and place cube off-diagonally)), cube double (task 1 (pick and place one cube), task (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 eube 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 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, vi sual cube double task 1, visual scene task 1, and visual puzzle 4x4 task from each domain, respectively. The observations are 64 x 64 x 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 statebbase and image-basee based tasks from OGBench, we will use IM unlabeled transitions collected by Mar-Maakov 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 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 different strategies for pre-training and fine-tuning in RL. First, implicit Q-Learning (IQL) (Kostrikoy et al., 2022) and revisited behavior-regularized actor-critic (ReBRAC) (Tarasov al., 2023a) are state-of-the-art offline RL algorithms based on the standard actor-critic framework (Appendix A) 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 celf-distillation with no labels (DINO) (Caron et al., 2021) as our representation learning loss and name the resulting bascline DINO + ReBRAC. Third, our next baseline, model- -ased 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., 201 8). 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., 2024b), 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
25

---

Published as a conference paper at ICLR 2026
Table I: Common hyperparameters for our method and baselines.
Hyperparameter
Value
learning rate
x 10
optimizer
Adam (Kingma, 2014)
pre-training gradient steps
x 10 for state-based tasks, x 10 for image-based tasks
fine-tuning gradient steps batch size
5 x 10 for state-based tasks, 1 x 10" for maage-baseed tasks
256
MLP hidden layer sizes
(512, 512, 512, 512)
MLP activation function
GELU (Hendrycks & Gimpel, 2016)
discount factor
0.99
target network update coefficient
x 10-3
double Q aggregation
min
policy update frequency in fine-tuning image encoder
1/4
small IMPALA encoder (Espeholt et al., 2018; Park et al., 2025b)
image augmentation method
random cropping
image augmentation probability
1.0 for DINO RcBRAC, 0.5 for all other methods
image frame stack
3
as the discriminative counterparts. After pre-training the ratio predictors, importance sampling is required to recover the Q function (CRI. + IS & TD InfoNCE IS) (Mazoure et al., 2023; Zheng al., 2024b), 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 zero-shot manner and then tune the policy using offline RL (Kim et al., 2024), differing from the implicit GPI as in our method. We will use IQL the fine-tuning algorithm and call the resulting methods FB IQL and HILP + IQL.. For image-basee tasks, we selectively compare to four baselines: ReBRAC, CRL + IS, DINO ReBRAC, and FB IQI.
D.3 EVALUATION PROTOCOLS
We compare the performance of InFOM against the eight baselines (Sec. 5.1) 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 tme-ttnnng and report the evaluation results averaged over 400K, 450K, and 500K gradient steps instead. For image-basee 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 9).
D.4 IMPLEMENTATIONS AND HYYPERPARAMETERS
In this section, we discuss the implementation details and hyperparameters for InFOM and the eight baselines. Whenever possible, we use the same set hyperparameters for all methods (Table 1) 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 times less frequently than other models during fine-tuning For specific hyperparameters of cach 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-basee based tasks, tune hyperparameters for each task individually.
cheetah: cheetah run
walker: walker walk
quadruped: quadruped jump
26

---

Published as a conference paper at ICLR 2026
Table 2: Hyperparameters for InFOM. See Appendix D.4 descriptions of each hyperparameter.
Hyperparameter
Value
latent intention dimension d' number of steps for the Euler method T
See Table 3
10
number of future states N
16
normalize the Q loss term in ((T) (Eq. 8) expectile
No
See Table 3
KL divergence regularization coefficient X
See Table 3
behavioral cloning regurlaization coefficient o
See Table 3
Table 3: Domain-specitte specific hy perparameters for our method and the baselines. We individually tune these hyperparameters for each domain and the set of hyperparameters for tasks in the same domain. See Appendix D tasks used to tune these hyperparameters and descriptions each hyperparameter. "." indicates that the hyperparameter does not exist.
TxPOM (Oun))
I0F:
RODRAC
TINNO - R/BRAC
MIIPO RADRAC
CR). 1
TTI EVE-NCI- IS
TTE TQT
FHI IQI
Gase
Meat
cheetab
125
0.9
0.06
0.5
[1]
0.1
0.1
124
0.05
0.008
walkes
312
6.1
0.5
10
0.1
0.1
124
0.06
10
quadroped
0.5
08
6.2
0.1
0.003
cube single
6,95
0.06
50
0.64
550
20
30
10
6.95
0.2
300
possle CK4
128
0.87
6.1
0.1
0.1
6.1
200
30
10
viouel cUN Mangiu
6.35
0.01
0.1
siouel scune Laak
125
6.95
0.1
990
0.
0.0
0.1
30
10
vioual pucciu 4ad Laak
138
6.95
0.1
900
0.
001
30
10
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 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 multilayer perceptron (MLP) that takes the next state-action pair (s', a') 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 optimiza- tion. We use the reparameterization trick (Kingma & Welling, 2013) to backpropagate the gradients from the flow-matching loss and the KI. divergence regularization (Eq. 4) into the intention encoder. Our initial experiments suggest that the dimension 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 A, we sweep over {2.0, 1.0, 0.2, 0.1, 0.05, 0.025, 0.01, 0.005) to find the best À for each domain. Second, we use flow- mot-hing vector fields to model the flow occupancy models. The vector field is an MLP that takes in a noisy future state sf. a state-action pair (s, a). and latent intention and outputs the vector field with the same dimension as the state. We apply layer normalization the vector field to stabilize optimization. As mentioned in Sec. 3, we use flow-matching objectives based optimal transport (linear path) and sample the time step from the uniform distribution over the unit interval. Following prior work (Park et al., 2025b), we use fixed T = 10 steps (step size 0.1) for the Euler method and do not apply sinusoidal embedding for the time. To make a fair comparison with other baselines, we also pre-train a behavioral cloning policy that serves initialization for fine-tuning
For fine-tuning 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
27

---

Published as a conference paper at ICLR 2026
MLP that predicts the scalar reward of 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) and sweep the expectile over {0.9, 0.95, 0.99} to find the best 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 or in Eq. 8 is important, and we sweep over (300, 30, 3, 0.3) to find the best for each domain. Following prior practice (Park et al., 2025b), we do not normalize the Q loss term in the actor loss [(x) (Eq. 8) as in Fujimoto & Gu (2021). Other choices of the policy network include the diffusion model (Ren et al., 2025; Wang ct 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 and Table 3 summarize the hyperparameters for InFOM
We reuse the IQL (Kostrikov et al., 2022) implementation and the Re- BRAC (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 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.013 to find the best AWR inverse temperature a 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.13 to search for the best actor BC coefficient aactoor and use the range {100, 10, 1, 0.1, 0} to search for the best critic BC coefficient Ceritic- We use the default values for other hyperparameterss following the implementation from Park et al. (2025b). See Table 3 for domain-specifie hyperparameters.
IQL and ReBRAC.
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, 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 Kisudent and use a fixed temperature 0.04 for teacher representations on all domains. We use representation space with 512 dimensions. We update the target representation centroid with a fixed ratio of 0.1. During pre-training, we learns the DINO representations along with 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 Qactor and Qeritic as ReBRAC. For image- based tasks, we apply random cropping to the same image twice with probability of 1.0 and use those two augmentations to compute the teacher and the student representations. See Table for domain-specific temperatures for student representations.
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 MIP to predict the residual between the next state s' and the current state s conditioned on the current state-action pair (s. We pre-train the one-step model with a behavioral cloning policy. During fine-tuning, we use the model with 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 Nimaginary and sweep over {1, 2, 4] for the number of steps in each rollout Himaginary We use the same BC coefficient as in ReBRAC. See Table for the domain-specific number of imaginary rollouts and number of steps in each rollout.
MBPO + ReBRAC.
28

---

Published as a conference paper at ICLR 2026
CRL + IS and TD InfoNCE + IS. We mostly reuse the CRL (Eysenbach et al., 2022) implemen- tation 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 implemen- tations (Zheng et al., 2024b). For both methods, we pre-train the classifiers that predict the ratio between the occupancy measures and the marginal densities over future states with behavioral cloning policy. We use the SARSA variant of TD InfoNCE during pre-training. After pre-training the classifiers, we learn reward predictor and apply importance sampling weights predicted by the classifiers to set of future states sampled from the fine-tuning datasets to estimate Q. This Q estimation then drives policy optimization. We use 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, 0.3, 0.03) for the BC coefficient a (Table 3).
FB IQL and HILP IQL. We implement FB (Touati & Ollivier, 2021) and HILPP (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 prwaardbbckwwrr repre- sentations and the intention-conditionel policies in an actor-critic manner. We use a coefficient 1 for the orthonormality regularization of the backward representations. We use 512-dimensional representations. We sample the latent intentions for pre-training from either 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 /512. We sweep over {100, 10, 1, 0.1} for the BC coefficient Chepe. For HILP,, we pre-train the Hilbert representations e 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, 1,0.1) } to find the best AWR inverse temperature a. We also use 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 v512, and then use the representation of the next state d5(") and the representation of the current state (s)) to compute the intrinsic reward (Ø(s') d(s))
During fine-tuning, we first infer a task-specific backward representation or a Hilhert representation using 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.13 for the AWR inverse temperature OAwR for IQL. For HILP, we reuse the same AWR inverse temperature in representation learning for IQL. See Table 3 for domain-specific BC coefficients and AWR inverse temperatures.
E ADDITIONAL VISUALIZATIONS OF LATENT INTENTIONS
We include additional visualization of latent intentions on quadruped jump in Fig. 6.
F ADDITIONAL EXPERIMENTS
F.1 EVALUATION ON ROBOTICS BENCHMARKS
To further study the pre-training and fine-tuning effects of our method on realistic datasets. Specifi- cally, we choose the RT-1 dataset (Brohan al., 2022), which contains 73499 episodes of transitions. This dataset was collected by commanding Google robot to pick, place, and move 17 objects in the Google micro-kitchens, covering 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 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 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
29

---

Published as a conference paper at ICLR 2026
Table Evaluation on ExORL and OGBench benchmarks. Following OGBench (Park al., 2025a), we bold values at and above 95% of the best performance for each task.
Tank
laFOM (Ouss)
IQL
ReBRAC
DINO ReBRAC
MIPPO REBRAC
C 15
TD INRANCE
FB EQL.
HILP
IQL
chectah zun
07.0 7.5
5010
07.0
12.9
106.7:2.4
73.3266.7
06.2 2-83
81.32 16.9
90.3 1.9
chastah zun backward
1
11.0
01:32
cheetah
526.3
walker xalk
467.5 62
908.6--33
908-1
208.0---37
377.6 4.5
313.3 7-8
211-3 1 13.3
015.3
925.4
walker run
116.3 15.3
82..4.06
07.5 1.3
96.54-1.2 1.0
107.6
91.5++3..
91.0 1.5
17.1+-12
97-4--24
nalkez gtand
081.2- 0 72.1
0.3.1:2..3
360.5 1.1
4.10
455.1 222
400.0:2:
8 6.0
$1682277
443.4:84
waiker sllp
341. 3-7
quadzoped zun
STL.1
quadzuped pump
573.6
718.3
688
183-2 8.
407.1
728
715.6
quadzuped salk
10 7.9
376.1 1:377
7:144:772
380.0 3.7
3015 9:6
241. 2: 29.3
388
7.0
491.4
jace reach top left
742.5 48.7
74.7 19.6
60.0 2.4.
12.62235
60.1 6.
23.1
31.5 2:30
25.0
40.4
Jaco reach top righe
6RT.5
12.0-1331
10.8
6.3
10.3
16.0 100---11
13.
zeach botten
740.7
Jaco
righe
single
02.5 4.45
$3.02:83 8.7
07.8
18
2.
13.5 185-888
12.3 8.86
62.9
cube single tack
76:44 13-1
51.7 15.1
0:
942 7.0
ST4
6..5..6
16.7 8.86
76.5 15-4
cuss singls task
$6.4 90.5
41.6 6.8
94.8 14.82:00 0.8
152 1.
98.1 2.
125
11.7 2.4
16.0 12.2
35.6 16.7
cuze single tack
1.0
16.7
36
double
10.5
17.8
22=
41415
6.7
0.2
double
12.5 10.7
1.3.112
0.0:
00:100
6.0 000
000000
0.000.0.
0.0000
00.10. 0.0
cube double Look
11.6 2-8
0.3-2.0.
0.00000
00200
6.0 006
011--00
000000
00000
00:000 0.0
cuse double task
0.:00. 0.4
0.0
0.0:0.0
00:00
0.0 2000
0012000
0.00000
0.0 000000
00.10. 0.0
cube double tack
1:5-111 1.0
00:00
00
00--00
02 09-0-
0.000 110
00
234
282.0
1.515
4-1
0.7 0.5
1.7
0.3 03
9.4 0.0.
0854003
05--1.0
0.-00
00
scens task
1.207
0.2
28:08
02 0.
2.0 1.1
122..4
0.
0.2
0.0
scene tack
03 0-1
0.0
0J0 0.0
00
0.0 110
03 01
0.000000 0.0
0.0 110
00
pozzie axt
24:2 14.4
2.5
13.81:3.11
033
00 0.0
pussle
022-04
2:2.
0.4
pussie 3xi
123:2:2
0.7
0.0:
00:0:00
6.000.0.0
05 0.6
0.0
6.0.2:0..
0.1
visual cube singie Lask
02.1 2008
106=772
15.4 14.6
12.02.00
31.0 15,0
vinusl cube co02
50.1
22.2
time
FB FOM
HILP FOM
InFOM
t-SNE dim
t-SNE dim 0
t-NEE dim
balance
jump
Figure 6: Visualization latent intentions on quadrupod-jung
et al., 2022; Dao et al., 2023), we therefore pre-train 3-VAE (Higgins et al., 2017) to encode images into latent embedding space and then learn algorithms 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 4 random seeds.
30

---

Published as a conference paper at ICLR 2026
20
InFOM (Ours)
ReBRAC
pairwise mean squared error
HILP + FOM
BC
DINO + ReBRAC
FB FOM
0.6
InFOM (Ours)
retur
0.4
10
0.2
0.0
google robot pick coke can
quadruped jump
scene task 1
Figure 7: Evaluation robotics datasets. InFOM outperforms the best baseline by 34% when trained top embeddings from a fixed image encoder. See Appendix F.1 details.
Figure 8: Comparison to prior intention encod- ing mechanisms after pre-training. We com- pare InFOM to prior intention encoding mech- anisms based on unsupervised skill discovery (HILP (Park et al., 2024b)) or successor feature learning (FB (Touati & Ollivier, 2021)) after pre- training. FB FOM is equivalent to TD flows with GPI in Farebrother al. (2025). InFOM achieves lower prediction errors both tasks.
Results in Fig. suggest that InFOM outperforms the best baseline by 34% when trained on top of embeddings from 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
We now conduct experiments ablating a key component in our method: the variational intention encoder. To investigate whether this framework induces simple and performant way infer diverse user intentions from unlabeled dataset, we compare it to various intention encoding mechanisms proposed by prior methods. Specifically, we consider replacing the variational intention encoder with either (/) set of Hilbert representations and Hilbert foundation policies (Park et al., 2024b) representations and representation-condiiinn 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).
(HILP FOM) or (2) a set of forward-backward
We first compare the future state predictions from InFOM against HILP + FOM and FB FOM on two ExORL tasks (quad drueed jump and scene task 1) after pre-training. Specifically, compute the pairwise mean squared error (MSE) between predicted future states and ground-truth future states along a trajectory. We first sample 100 trajectories from the pre-training datasets, and then, for each trajectory, we sample 400 future states from InFOM and the two baselines starting from the same initial (s, a) pair. We compute the pairwise MSE between each sampled future state and the corresponding sequence of ground-trun mound future states within the same trajectory. The prediction error reported as the pairwise MSE averaged over all transitions in the 100 trajectories and the 400 sampled future states. Results in Fig. 8 show that InFOM achieves lower prediction errors than two FOM baselines.
We then compare the performance of InFOM against HILP FOM and FB + FOM after fine-tuning. 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. Results in Fig. 9 indicate that InFOM can outperform prior intention encoding methods of tasks, while being simpler. Both HILP and FB capture intentions with full unsupervised RL objectives based on an actor- critic backbone. In contrast, capture intentions simply training intention encoder based latent
31

---

Published as a conference paper at ICLR 2026
cheetah run
quadruped jump
120
800
INFON LOuno 66 FOM
aUcOss rote
return
0.
0-
0.0
tine turning steps
1ed
tine fin-tuning stop
IFFM (Ours)
CRL +
MBPO ReBRAC
Figure 9: Comparison to prior intention encoding mechanisms after fine-tuning, We compare InFOM prior intention encoding mechanisms based after fine-tuning, We observe that InFOM outperforms prior methods on 3 out of the 4 tasks.
ReBRAC
DINO + ReBRAC
IQL
Figure 10: Convergence speed during fine- tuning. On tasks where InFOM and baselines perform similarly, our flow occupancy models en- able faster policy learning.
variable model over adjacent transitions, without relying on a potentially complicated offline RL. procedure (Tarasov et al., 2023b; Park al., 2024a).
F.3 FLOW OCCUPANCY MODELS ENABLE FASTER POLICY LEARNING
We then investigate whether the proposed method leads to faster policy learning 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 for details of these baselines). We choose these baselines because they perform similarly our method, helping 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, InFOM converges faster than prior methods that only pre-train behavioral cloning policies (ReBRAC) or self-supervised state represen- tations (DINO ReBRAC), demonstrating the effectiveness of extracting temporal information. The observation that methods utilizing a one-step transition model (MBPO ReBRAC) or future state classifier (CRL+ IS) learn more slowly than our method highlights the importance of predicting long-horts 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 DISCRETE INTENTIONS
The choice of the prior over latent variables p(z) is still an open question in the literature. Prior work has used standard Gaussian distribution (Frans al., 2024), uniform Mises-Fishe distribution Park et al. (2024b); Touati & Ollivier (2021); ); Zheng et al. (2025), a continuous uniform distribution (Sharma et al., 2019), and discrete uniform distribution (Eysenbach et al., 2019).
To further investigate the effect of using a discrete set of latent intentions for InFOM, we run additional ablation experiments. We selected a set of discrete latent embeddings 2 {-,,-.. .2x) (a lookup table with K 256), and used a vector quantization (VQ) loss to learn those embeddings together with InFOM in VQ-VAE (Van Den Oord et al., 2017). Specifically, given a consecutive transition (s, s', .a), the flow-based intention decoder qa(2 a) remains the same. while the intention encoder Pe (z s', a') can now be decomposed into two components: (I) the deterministic encoder Penc S x A - Rd and the quantizer Pquam Rd - Z. The role of the quantizer is to query the closest discrete latent intentions from the encoder outputs using the nearest neighbor,
Pquami (paum (Sonces's,0 a.)) =k. where k = argmin K Pe (s', aa)) 2:|12-
32

---

Published as a conference paper at ICLR 2026
1,0
1.0
InfdW (Ound)
InTOM (Cund) nTON the funed
InTOW dliceee eeerroon..
Succemn rate
success rota
6.0
Wakertfo
2uadruse3 Jume
siggl double
quadruped jump
Soub 1ask
scone
Figure 11: Using discrete intentions slightly de- creases InFOM's performance on ExORL tasks (-11%), while drastically decreasing the mean suc- cess rate of InFOM on OGBench tasks (-78%),
Using the posterior ga(= s', a') sam- ple the latents does not significantly change perfor- mance InFOM (+7%), suggesting that our method is robust against unseen latents. We choose to use the prior p(z) for sampling latents to estimate Q= through- out experiments.
Figure 12:
Using this quantizer, we replace the surrogate objective in Eq. with the following SARSA flow loss with vector quantization loss:
CSARSA flow (Pency Pquam; 9d) + Lvo Pquam Qa);
Lvo(Pene Pquant qa) Epp sa Pene (s', lse 2k (s', al)
+ DE (a',a)) (1) (s',a') (zx(s', a')ss'""')Jsll)
where [-Isg denotes the stop gradient operator, and use straight-through gradients (Bengio et al., 2013) to optimize the SARSA flow loss. During fine-tuning, we use all the discrete latents {21,-.. _
construct intentioncoodiitinn Q= estimations (Eq. 6) and distill them into the critic Q as in Eq. 7.
We conducted ablation experiments on two ExORL tasks (walker flip and quadruped jump) and two OGBench tasks (cube single task 2 and cube double task 1) and report performances aggregated over 8 random seeds. Results in Fig. suggest that using discrete intentions slightly decreases InFOM's performance on ExORL tasks (-11%), while drastically decreasing the mean success rate of InFOM on OGBench tasks (-78%). These results indicate that using a continous latent space generally leads to better performance in our experiments.
F.5 FINE-TUNING WITH POSTERIOR INTENTIONS
In Sec. 4.4, when estimating the cecation-conditionel Q= for a specific task, we have already sampled the latent from the prior p(z) instead of the posterior qa(z a'). Sampling from the prior, in general, increases the possibility of drawing out-of-distribution latents. We hypothesize that InFOM can generalize over unseen latents different (s, a) pairs. To quantitatively test this hypothesis, we conduct additional ablation experiments to study the effect of estimating intention-condittinn Qa using in-distribution latents the final performance of InFOM. Specifically, we replace the distillation loss in Eq. 7 with a variant that samples from the posterior qa(= 8', a')::
E(Q)= = E (s,,,, (o'' (s.a,s' ,a),, (=]s"_a)" L (QI a) Q(s,a)).-
We choose to conduct ablation experiments on two EXORL tasks (wa leer flip and quadruped jump) and two OGBench tasks (cube double task 1 and scene task 1), aggregating the return and the success rate over 8 random seeds. Results in Fig. 12 indicate that using the posterior to sample the latents for each Q, does not significantly change the performance InFOM (+7%). Conversely, these results suggest that InFOM is robust against unseen latents for different (s, a) pairs and using the prior p(z) )) to sample latents provides sufficient learning signals drive fine-tuning. We choose to use the prior p(z) for sampling latents to estimate Q2 throughout our experiments.
F.6 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 jaco reach top left and jaco reach bottom right, studying whether using dense rewards will miti- gate the performance gap. Specifically, the dense reward function is defined as r(s.g) g) -||ss gll2
33

---

Published as a conference paper at ICLR 2026
800
1.0
quad ruped jump
InToWin SC InFOM (Our)
Feasson cofiecnt -
Faorsu come uidn corficart
SUcoss rete
205
14tm
400
State
186
590.34
500.1
walkor Alp
quadruppdd
0.0
"obba double tase
soone tack
pre-traning Calaset sioss
pre-training dataset sizes
Figure 14: The behavioral cloning regularization in the policy loss is key component of InFOM.
Figure 15; The diversity of the pre-training datasets has a positive correlation with their sizes.
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 random seeds.
Results in Fig. 13 highlight that using a dense reward function results 3.6x smaller perfor- mance 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 ob- servation, where TD3 BC (the base algorithm for ReBRAC) performed poorly on the jaco domain (Fig. 9 of Yarats et (2022)).
ReBRAC
ReBRAC + dense reward InFOM (Ours)
800
return
400
F.7 IMPORTANCE OF THE BEHAVIORAL. CLONING REGULARIZATION
0
jaco reach top left
jaco reach bottom right
To study the effect of the BC regularizer (Eq. 8), we conduct experiments comparing a variant of InFOM without the behavioral cloning regular- ization coefficient (a 0) to our full algorithm with domain domain-dependene dependent values (Table 2). We select the same EXORL and OGßench tasks as in Fig. 5 (wa 1 ker fl i,, 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. 14 suggest that behavioral cloning regularization (a >> in the policy loss key component of our algorithm.
Figure 13: Reward function structure can impose challenges. The baseline ReBRAC achieves 3.6x higher performance on variants of jaco tasks with dense reward function.
F.8 DIVERSITY OF THE PRE-TRAINING DATASETS
To quantify the diversity of the pre-training dataset, conduct a statistical analysis the datasets for two ExORL tasks (cheetah run and quadruped jump), analyzing the relationship between the size of the dataset and the diversity of the dataset. Following prior work (Park et al., 2023b), discretize the continuous state space as nigh-dimensiooal grid (up to 10-2) and use the number of unique grid points covered by the dataset to measure the diversity. Results in Fig. 15 show that increasing the dataset size induces higher diversity in the pre-training datasets, with an average correlation coefficient of 0.76 those tasks. Thus, we can study the effect of diverse pre- training datasets on InFOM's 's performance by varying the pre- pre-training dataset size.
F.9 THE EFFECT OF DATASET SIZES
Pre-training dataset size. Since we aim to predict temporally distant future states from het- erogeneous data (Sec. 4.1), InFOM implicitly requires sufficiently diverse dataset for effective pre-training. To study the relationship between the size of pre- mre-training datasets and the perfor- mance our algorithm. we conduct ablation experiments varying the pre- trainininng dataset size in (1K, 10K, 100K, 500K, 1M). We compare the performances of InFOM on two ExORI tasks
34

---

Published as a conference paper at ICLR 2026
cheetah
quadraped jump
single
puzzzic
RUCCMMS rate
1000
10
366
236
506
U0KK
1996
235
1900
2566
pre-trating dal csst sines
prestraining g datissert sizes
tie tuniiuing dataset sizts
tin tun datasct stzes
Figure 16: The effect of pre-training dataset size InFOM. Increasing pre- pre-training dataset sizes boosts the final performances of InFOM.
Figure 17: The effect of fine-tuning dataset size on InFOM. Increasing the fine- line-tuning dataset size yields consistent improvements in success rates.
(cheetah 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.
Results in Fig. 16 indicate that larger pre-training datasets yield higher returns these tasks. We conjecture that pre-training InFOM on a diverse, reward-free dataset reduces the possibility of sampling out-of-distribution (unseen) intentions, resulting in 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 s single task 2 and puzzle 4x4 task 1) and vary the size of the fine-tuning datasets in {5K, 10K, 25K, 50K, 100K, 250K, 500K). Again, we aggregate the performance of InFOM over 8 random seeds, following the same evaluation protocol in Appendix D.3.
Results Fig. 17 show that increasing the fine-tuning dataset size (within the chosen range) yields con- sistent improvements 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.10 FINE-TUNING ON SUBOPTIMAL DATASETS
We hypothesize that using highly suboptimal fine-tuning datasets will decrease the downstream performance of InFOM. To study the effect of fine-tuning on suboptimal datasets, we conduct ablation experiments two ExORL tasks (cheetah have dense reward functions and can still produce diverse rewards. To construct suboptimal datasets, we use the reward quantile filter each transition in the 10 EXORL dataset collected by RND (see Appendix D.1 for details) and then sample 5 X 10t reward-labeled transitions from the remaining transitions. After constructing these suboptimal datasets, we use them to fine-tune InFOM. Results in Fig. 18 indicate that fine-tuning InFOM on highly suboptimal datasets (0.2 reward quantile) achieved only 9% performance of the original InFOM, while using datasets with 0.8 reward quantile can already achieve 85% performance of the original InFOM. These results suggest that using sufficiently optimal dataset is important for improving the fine-tuning performance.
and quadruped jump) because they
F.11 THE SUFFICIENT NUMBER OF FUTURE STATES IN THE Q ESTIMATION
Since we use MC future states from the InFOM to estimate the intention-conditione Q2 (Eq. 6). the model may produce unrealistic future states. Thus, the number of future states N affects the accuracy and variance of the Q value estimation (Eq. 6). To investigate the effect of N, we con- duct ablation studies on total
tasks, with tasks from the EXORL benchmarks (cheetah wall wall walke ker wa 1 k, wa 1 ker flip, and quadruped jump) and 4 tasks from the OG- Bench 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 random seeds.
Fig. 19 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, larger N does mitigate the high variance in Qu, at the cost of increasing computation. Taken together, these results indicate that a sufficiently large number of flow future states used in Q- achieves accurate estimation of Q values, while reducing the variance. In
35

---

Published as a conference paper at ICLR 2026
EXORL
cheetsh
cusdruped
SUCCESS nac
number of flow future
number flow future states
war maller nafc
pauk dullock
Ineduring cataset reward cuanties
The-tuning datases rewore quantiles
cestrane Jano
tenc soone Hal0
Fine-tuning on suboptimal datasets. Fine-tuning on highly suboptimal datasets (0.2 re- ward quantile) decreased the performance of InFOM, while using sufficiently optimal (0.8 reward quan- tile) dataset can already retain the performance.
Figure 18:
Figure 19: Using sufficient number of flow future states is important. Increasing the number of flow fu- states (NV in the Q, estimate boosts the accuracy while reducing variance, resulting in higher final perfor- mances of InFOM. We choose between the accuracy. variance, and computational con- straints in experiments.
16 as a balance
contrast, a smaller number of N potentially yields errors in Q= from unrealistic future states, resulting in high variance. In practice, our choice of N 16 is balance between the accuracy, variance, and computational constraints of the estimator.
F.12 ADDITIONAL HYPERPARAMETER ABLATIONS
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 In- FOM (Table 2). Following the same evaluation protocols as Appendix D.3, we report means and standard deviations across eight random seeds after fine-tuning each variant.
As shown in Fig. 20a, our algorithm is sensitive to the latent intention dimension d. Additionally, the effect of the number of steps for the Euler method T (Fig. 20b) saturates after increasing it to certain threshold (T = 10), suggesting the usage of a common value for all tasks.
Results in Fig. 20c, Fig. 20d, and Fig. 20e suggest that the expectile on EXORL tasks, while having minor effects on OGBench tasks. Importantly, the KL divergence regularization coefficient À and the behavioral cloning regularization coefficient parameters for InFOM, where domain-specifi hyperparameter tuning is required. As discussed in Appendix D.4, generally select task from each domain sweep hyperparameters and then use one set of hyperparameters for every task in that domain.
affect the performance
crucial hyper-
36

---

Published as a conference paper at ICLR 2026
800
10 IDan
400
400
0.5
0.0
0.0
duuble
ousdruped
(a) Latent dimension d
(b) Number of steps for the Euler method T
Jl D.9 (Ours)
#= 0.99
0.99 (Ours)
0.95
0.95
success rate
1.0
return
400
0.5
0.0
walker flip
quadruped jump
cube double task scene task 1
(c) Expectile μ
3-000
A-Al
800
6:21
A-31
2-8-35
(Pura
A.0.OL
(Cun)
0.6
EaCOCos rale
walker flp
quadruped jump
Scane
(d) KL divergence regularization coefficient À
600
800
350
393
SUCCESS rate
300
8:00
0.5
0.0
0.0
wolker fpp
quadruped jump
oube doudle tock
500% tack
(e) Behavioral cloning regularization coefficient
Figure 20: Hyperparameter ablations. We conduct ablations to study the effect of key hyperpa- rameters of InFOM as listed in Table 2 on walker flip, quadruped jump, cube double task 1, and scene task 1.
37


Reviews:
[{"review_id": "VHNOk3hSQL", "reviewer": "Reviewer_DcLe", "summary": "The present work leverages flow-matching models to predict trajectories of future states for RL tasks to help the actor finding better policies.\nExtensive experiments shows on gym and robotic manipulation benchmarks  shows promissing results.", "questions": "1- How does the inference work ? No algorithm is provided. \n\n2- How do you intend to use this method for pre-trained policies ?", "limitations": "", "rating": 2, "confidence": 2, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "- The idea of leveraging \"intent\" for multi-step generation especially in the multi-agent setting or iterative refininement with human input. \n- Extensive experiments show practicability of the method.", "weaknesses": "Major\n\n After several reading, I fail to understand how the algorithm works. It should not be that hard to understand: \n- flow-matching to predict future states is straightforward given trajectories either conditionned or guided on observed states.\n- the predicted latent variable has to be used somehow by the policy, either by boosting an pre-trained one or by training from scratch.  \n\n 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. \n\nMy 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. \n\nI 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. \n\n\nMinor \n\n1- 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.\n\n2- \"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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "L4tQCYDgHZ", "reviewer": "Reviewer_Eojh", "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.", "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?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "* The paper is well written and easy to follow\n* The experimental results showing the proposed approach outperforms the baselines in most of the case with the gap widening on more difficult tasks\n* A large number of baselines have been included and sufficient experimental detail is provided.\n* I really the like the qualitative analysis of the learnt latent intention model.", "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gv4AZVywaW", "reviewer": "Reviewer_Skrx", "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.", "questions": "Q1: How well does the flow occupancy model generalize to out-of-distribution intentions? \n\nQ2: Is there a risk that the generative occupancy model produces unrealistic futures for novel $z$, harming policy learning?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Novel integration of flow matching with intention-aware occupancy modeling.\n\n2. Strong empirical results: consistent gains across diverse domains, including challenging sparse-reward and vision-based tasks.\n\n3. Well-motivated design choices, especially the use of SARSA-style bootstrapping and expectile-based implicit GPI to stabilize training.", "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”).\n\nSuggestion: 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?).\n\n2.  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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yv7hCY69Py", "reviewer": "Reviewer_q22z", "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.", "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? \n- Did the newly trained reward, Q-value, and policy models use the features learnt by the occupancy model?\n- 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?\n- What are the results on the occupancy measure pretraining performance separately, compared to a standard TD-flows without intention decoding?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is very well-written and describes the whole training pipeline in great detail.\n- 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.\n- 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.\n- Strong performance on many benchmarks;", "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;", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)