Imagining Goals with Geometric Horizon Models for Policy Generalization

# Imagining Goals with Geometric Horizon Models for Policy Generalization
Source: [https://arxiv.org/html/2610.19129](https://arxiv.org/html/2610.19129)
Dhruv SreenivasAffiliation:MilaAdriana HugessenAffiliation:MilaKhimya KhetarpalAffiliation:MilaAffiliation:Université de Montréal,Google DeepMind,†Equal AdvisingGlen BersethAffiliation:Mila

###### Abstract

Training from offline data has allowed for substantial progress in domains such as robotics, leading to general\-purpose policies that can be easily applied zero\-shot or efficiently finetuned for downstream tasks\. However, these policies can still have poor generalization, both due to the choice of modeling objective and from learning from a static dataset\. In this work, we focus on the challenging task of zero\-shot goal generalization, where a policy is evaluated on unseen tasks that require combining its existing knowledge \(compositional generalization\)\. An avenue for improving a policy’s generalization is by generating new experience with world models; however, such generation has proven difficult for longer horizons\. Thus, to alleviate this issue, we propose TD\-Aug, which samples from a geometric horizon model and allows for directly imagining novel outcomes that can be achieved by composing existing knowledge\. We demonstrate that training on these future outcomes as goals for goal\-conditioned BC policies significantly improves generalization in stitching\-based OGBench tasks\. Project page:[https://td\-aug\.github\.io](https://td-aug.github.io/)

††footnotetext:Correspondence:daniel\.lawson@mila\.quebec## 1Introduction

Training policies with goal\-conditioned behavior cloning \(BC\) provides a simple strategy to learn from prior experience\. However, generalization remains a challenge, limiting performance outside of the training data\. Generally, world models\([Ha and Schmidhuber, 2018](https://arxiv.org/html/2610.19129#bib.bib28)\)present an interesting strategy for improving policies from offline data\([Hafner et al\., 2025](https://arxiv.org/html/2610.19129#bib.bib15)\)\. These models can be used to generate synthetic experience\([Sutton, 1991](https://arxiv.org/html/2610.19129#bib.bib20);[Bruce et al\., 2024](https://arxiv.org/html/2610.19129#bib.bib14)\), allowing policies to get more out of limited real data, which may be important in domains such as robotics\. In this work, we consider generating experience using a type of world model to imagine counterfactual future outcomes to enable improved goal generalization\.

We tackle a gap identified for BC policies showing poor performance on tasks involving stitching\([Ghugare et al\., 2024](https://arxiv.org/html/2610.19129#bib.bib37)\)\. In this setting, after policies are trained on a subset of tasks, they are then evaluated zero\-shot on harder tasks involving composition\. We illustrate an example in Figure[2](https://arxiv.org/html/2610.19129#S2.F2), where a policy is trained on how to solve short horizon tasks \(left\), and then tested on a compositional task, such as going across the maze \(middle\)\. To solve this problem, we propose to learn a generative model that predicts what states will be encountered after mixing behaviors of the independent data collecting policies \(right\)\. Specifically, the generative model we consider is a geometric horizon model \(GHM\)\([Janner et al\., 2020](https://arxiv.org/html/2610.19129#bib.bib45);[Thakoor et al\., 2022](https://arxiv.org/html/2610.19129#bib.bib46)\)trained using Temporal Difference Flows \(TD\-Flows\)\([Farebrother et al\., 2025](https://arxiv.org/html/2610.19129#bib.bib43)\)\. Then, using samples generated from our GHM as a data augmentation for goals, we teach our policy to perform compositional tasks\. We provide an overview in Figure[1](https://arxiv.org/html/2610.19129#S1.F1), and refer to our method asTD\-Aug\. We further justify this procedure by showing that augmenting BC with goals from the successor measure relates to value weighted regression\.

To implement TD\-Aug, we study several GHM objectives based on both temporal difference \(TD\)\([Farebrother et al\., 2025](https://arxiv.org/html/2610.19129#bib.bib43);[Farebrother et al\., 2026](https://arxiv.org/html/2610.19129#bib.bib21)\)and geometric Monte Carlo \(MC\)\([Eysenbach et al\., 2022](https://arxiv.org/html/2610.19129#bib.bib35)\)sampling, demonstrating that TD methods are crucial for generalization\. Using horizon consistency\([Farebrother et al\., 2026](https://arxiv.org/html/2610.19129#bib.bib21)\), we propose a new TD\-Flow objective which combines both MC and TD sampling during training, leading to improved long\-horizon generation\.

Empirically, in Section[6\.1](https://arxiv.org/html/2610.19129#S6.SS1), we first validateTD\-Augon low\-dimensional navigation tasks in OGBench\([Park et al\., 2025a](https://arxiv.org/html/2610.19129#bib.bib36)\)that require stitching for success, substantially improving behavior cloning generalization, from31%31\\%to79%79\\%average success\. Then, we demonstrate that our approach can also improve generalization on high\-dimensional, image\-based tasks \(29%29\\%to56%56\\%\) in Section[6\.2](https://arxiv.org/html/2610.19129#S6.SS2)and complicated manipulation tasks \(13%13\\%to52%52\\%\) in Section[6\.3](https://arxiv.org/html/2610.19129#S6.SS3)\.

Figure 1:Overview of TD\-Aug\.Wea\)train a GHM on an offline dataset to approximate samples from the successor measure of the behavior mixture policy\. Then,b\)we train BC policies on a mix of ground\-truth and generated goals,improving zero\-shot generalization on compositional tasks\.Algorithm 1TD\-AugRepresentation Learning \(Optional\)

Use either states

ss, or learned representations

z=ϕ⁡\(s\)z=\\phi\(s\), as the space for GHM and policy \(Section[6\.2](https://arxiv.org/html/2610.19129#S6.SS2)\)\.

Geometric Horizon Model Training

Train

MϕβM^\{\\beta\}\_\{\\phi\}to approximate the successor measure of the behavior mixture policy, using

ℒMCHC​\-​β\\mathcal\{L\}\_\{\\mathrm\{MCHC\}\\text\{\-\}\\beta\}\(Equation[9](https://arxiv.org/html/2610.19129#S4.E9)\)\.

Policy Training

Train

πθ\\pi\_\{\\theta\}by BC on a mixture of generated goals from

MϕβM^\{\\beta\}\_\{\\phi\}with proportion

α\\alpha, as well as real samples \(Equation[11](https://arxiv.org/html/2610.19129#S5.E11)\)\.

## 2Related Work

Data augmentationwith temporal structure has been an effective way to re\-use experience in reinforcement learning\. Models can generate data, for example by sampling transitions for model\-free approaches\([Sutton, 1991](https://arxiv.org/html/2610.19129#bib.bib20);[Ha and Schmidhuber, 2018](https://arxiv.org/html/2610.19129#bib.bib28);[Feinberg et al\., 2018](https://arxiv.org/html/2610.19129#bib.bib17);[Hafner et al\., 2020](https://arxiv.org/html/2610.19129#bib.bib19);[Lu et al\., 2023](https://arxiv.org/html/2610.19129#bib.bib34);[Zhou et al\., 2024](https://arxiv.org/html/2610.19129#bib.bib41)\)\. In the goal\-conditioned setting, Hindsight Experience Replay \(HER\)\([Andrychowicz et al\., 2017](https://arxiv.org/html/2610.19129#bib.bib18)\)allows re\-using transitions with different achieved outcomes as goals\([Kaelbling, 1993](https://arxiv.org/html/2610.19129#bib.bib3)\)\. When training policies with BC, sampling hindsight goals is beneficial, but does not covercompositional generalization\([Ghugare et al\., 2024](https://arxiv.org/html/2610.19129#bib.bib37)\), where a policy is evaluated on state\-goal pairs not seen together in the same training trajectories\. To mitigate this gap, states are stitched across trajectories using clustering, but this only covers 1\-step stitches and relies on heuristics, such as a good distance metric between states\. In contrast, we use generative models to capture a future state distribution that directly stitches across trajectories\.

Long\-Horizon Prediction:In the context of predicting future states, instead of generating states sequentially, prediction can be made more accurate by sampling states in parallel, i\.e\. jointly sampling sub\-trajectories\([Janner et al\., 2022](https://arxiv.org/html/2610.19129#bib.bib26)\), which could be chained together\([Luo et al\., 2025](https://arxiv.org/html/2610.19129#bib.bib33);[Lee and Choi, 2025](https://arxiv.org/html/2610.19129#bib.bib32)\), for longer\-horizon or compositional predictions\. In comparison, we avoid having to denoise intermediate states or utilize heuristic strategies to chain sub\-trajectories together\. This is achieved while still generating long\-horizon samples inO⁡\(1\)O\(1\)steps through amortization during GHM training\.

Representation learningcan be employed to improvecompositional generalizationof BC by encoding temporal dependencies in the representations of a policy, such as with objectives that relate to approximating successor measures or temporal distance\([Myers et al\., 2025](https://arxiv.org/html/2610.19129#bib.bib38);[Lawson et al\., 2025](https://arxiv.org/html/2610.19129#bib.bib29);[Park et al\., 2026b](https://arxiv.org/html/2610.19129#bib.bib25)\)\. This has the advantage of avoiding state prediction, but still has challenges, such as with stability and long\-horizon generalization\. In practice, generative modeling and representation learning can be used together, which is useful for predicting high\-dimensional states\.

Geometric horizon models \(GHMs\), also known asγ​\-models\\gamma\\textrm\{\-models\}\([Janner et al\., 2020](https://arxiv.org/html/2610.19129#bib.bib45)\), propose a generative variant of temporal difference learning, allowing for approximate sampling from thesuccessor measure\([Thakoor et al\., 2022](https://arxiv.org/html/2610.19129#bib.bib46);[Tomar et al\., 2024](https://arxiv.org/html/2610.19129#bib.bib22);[Wiltzer et al\., 2024](https://arxiv.org/html/2610.19129#bib.bib23);[Farebrother et al\., 2025](https://arxiv.org/html/2610.19129#bib.bib43)\), a continuous generalization ofsuccessor representations\([Dayan, 1993](https://arxiv.org/html/2610.19129#bib.bib24)\)of a policy\. Specifically, we consider Temporal Difference Flows\([Farebrother et al\., 2025](https://arxiv.org/html/2610.19129#bib.bib43)\)\(TD\-Flows\), which use flow\-matching objectives for learning such models\. Most recent work with TD\-Flows\([Farebrother et al\., 2025](https://arxiv.org/html/2610.19129#bib.bib43);[Farebrother et al\., 2026](https://arxiv.org/html/2610.19129#bib.bib21)\)has considered using GHMs for decision\-time planning\. GHMs can also be used for policy improvement\([Zheng et al\., 2026](https://arxiv.org/html/2610.19129#bib.bib27)\), or model\-based value expansion\([Feinberg et al\., 2018](https://arxiv.org/html/2610.19129#bib.bib17)\), where future synthetic states can be used to estimate bootstrap targets for value estimation\([Janner et al\., 2020](https://arxiv.org/html/2610.19129#bib.bib45)\)\. Recent work has used MVE for offline RL using GHMs and Universal Horizon Models \(UHMs\)\([Chung et al\., 2026](https://arxiv.org/html/2610.19129#bib.bib16)\), but does not consider the goal\-conditioned setting\. In contrast, we particularly study GHMs for goal augmentation, and how this alone can induce better zero\-shot goal generalization after offline training\.

![Refer to caption](https://arxiv.org/html/2610.19129v1/pointmaze_giant_train_eval_mchc_panels.png)Figure 2:Ground\-truth and generated samples\.Left, onpointmaze\-giant\-stitch, we display ground truth trajectories from the dataset\.Middle, we display long\-horizon evaluation tasks\.Right, we generate samples from an MCHC TD\-Flow, as in Section[4](https://arxiv.org/html/2610.19129#S4), conditioned on the initial \(outlined\) state and action\. These samples correspond to reachable future states, including those that can only be reached through stitching several trajectories with longer total horizon than short training trajectories\. We provide visualizations across more environments in Figure[8](https://arxiv.org/html/2610.19129#A4.F8)\.
## 3Background

### 3\.1Goal\-Conditioned Reinforcement Learning

We study offline learning in goal\-conditioned Markov Decision Processes \(MDPs\), defined as the tuple\(𝒮,𝒜,𝒢,P,r,γ,μ\)\(\\mathcal\{S\},\\mathcal\{A\},\\mathcal\{G\},P,r,\\gamma,\\mu\), where𝒮\\mathcal\{S\}is the state space,𝒜\\mathcal\{A\}is the action space,𝒢\\mathcal\{G\}is the goal space,P:𝒮×𝒜→Δ⁡\(𝒮\)P:\\mathcal\{S\}\\times\\mathcal\{A\}\\rightarrow\\Delta\(\\mathcal\{S\}\)is the transition kernel,rris the reward function,γ∈\[0,1\)\\gamma\\in\[0,1\)is the discount, andμ∈Δ⁡\(𝒮\)\\mu\\in\\Delta\(\\mathcal\{S\}\)is an initial state distribution\. We consider that goals coincide with states \(𝒢=𝒮\\mathcal\{G\}=\\mathcal\{S\}\), and the reward is the Dirac delta at the goal:rg​\(st,at,st\+1\)=δg​\(st\+1\)r\_\{g\}\(s\_\{t\},a\_\{t\},s\_\{t\+1\}\)=\\delta\_\{g\}\(s\_\{t\+1\}\)\. We aim to learn a policyπ:𝒮×𝒢→Δ⁡\(𝒜\)\\pi:\\mathcal\{S\}\\times\\mathcal\{G\}\\rightarrow\\Delta\(\\mathcal\{A\}\)using a dataset𝒟\\mathcal\{D\},*without online interaction*, to maximize expected return:

J\(π\)=𝔼g∼p\(g\),s0∼μ,at∼π\(⋅∣st,g\),st\+1∼P\(⋅∣st,at\)\[∑t=0∞γtrg\(st,at,st\+1\)\]\\displaystyle J\(\\pi\)=\\mathbb\{E\}\_\{g\\sim p\(g\),s\_\{0\}\\sim\\mu,a\_\{t\}\\sim\\pi\(\\cdot\\mid s\_\{t\},g\),s\_\{t\+1\}\\sim P\(\\cdot\\mid s\_\{t\},a\_\{t\}\)\}\\left\[\\sum\_\{t=0\}^\{\\infty\}\\gamma^\{t\}r\_\{g\}\(s\_\{t\},a\_\{t\},s\_\{t\+1\}\)\\right\]
For the subsetX⊆𝒮X\\subseteq\\mathcal\{S\}, we define the*normalized successor measure*of a policyπ\\pias:

Mγπ​\(X∣s,a\)≜\(1−γ\)​∑k≥0γk​Pr​\(sk\+1∈X∣s0=s,a0=a,π\)\\displaystyle M^\{\\pi\}\_\{\\gamma\}\(X\\mid s,a\)\\triangleq\(1\-\\gamma\)\\sum\_\{k\\geq 0\}\\gamma^\{k\}\\mathrm\{Pr\}\(s\_\{k\+1\}\\in X\\mid s\_\{0\}=s,a\_\{0\}=a,\\pi\)\(1\)wherePr\\mathrm\{Pr\}denotes the time\-specific probability that we are inXXstarting from\(s,a\)\(s,a\)when followingπ\\pi\. We write the density of the successor measuremγπm^\{\\pi\}\_\{\\gamma\}viaMγπ​\(d​s\+∣s,a\)=mγπ​\(s\+∣s,a\)​d​s\+\\smash\{M^\{\\pi\}\_\{\\gamma\}\(\\mathrm\{d\}s\_\{\+\}\\mid s,a\)=m^\{\\pi\}\_\{\\gamma\}\(s\_\{\+\}\\mid s,a\)\\,\\mathrm\{d\}s\_\{\+\}\}, and the action\-marginalized successor measure is defined asMγπ\(⋅∣s\)≜𝔼a∼π\(⋅∣s\)\[Mγπ\(⋅∣s,a\)\]\\smash\{M^\{\\pi\}\_\{\\gamma\}\(\\cdot\\mid s\)\\triangleq\{\\mathbb\{E\}\}\_\{a\\sim\\pi\(\\cdot\\mid s\)\}\[M^\{\\pi\}\_\{\\gamma\}\(\\cdot\\mid s,a\)\]\}\. We also define the discounted state occupancy density as the marginal ofmγπ\\smash\{m^\{\\pi\}\_\{\\gamma\}\}over the initial state distribution:dπ\(s\)≜𝔼s0∼μ,a0∼π\(⋅∣s0\)\[mγπ\(s∣s0,a0\)\]\\smash\{d^\{\\pi\}\(s\)\\triangleq\{\\mathbb\{E\}\}\_\{s\_\{0\}\\sim\\mu,a\_\{0\}\\sim\\pi\(\\cdot\\mid s\_\{0\}\)\}\[m^\{\\pi\}\_\{\\gamma\}\(s\\mid s\_\{0\},a\_\{0\}\)\]\}\. With the goal reaching reward, the successor measure’s density relates to thegoal\-conditioned value functionviaQγπ​\(s,a,g\)=\(1−γ\)−1​mγπ​\(g∣s,a\)Q^\{\\pi\}\_\{\\gamma\}\(s,a,g\)=\(1\-\\gamma\)^\{\-1\}m^\{\\pi\}\_\{\\gamma\}\(g\\mid s,a\)\.

### 3\.2Flow Matching

Following success in modeling policies and geometric horizon models\([Farebrother et al\., 2025](https://arxiv.org/html/2610.19129#bib.bib43)\), we utilizeflow matching\([Lipman et al\., 2023](https://arxiv.org/html/2610.19129#bib.bib44)\)\. With a*source*distributionp0∈Δ⁡\(ℝd\)p\_\{0\}\\in\\Delta\(\\mathbb\{R\}^\{d\}\)and a*target*distributionp1∈Δ⁡\(ℝd\)p\_\{1\}\\in\\Delta\(\\mathbb\{R\}^\{d\}\), flow matching learns a velocity fieldvτ:ℝd→ℝdv^\{\\tau\}:\\mathbb\{R\}^\{d\}\\rightarrow\\mathbb\{R\}^\{d\}such that forx0∼p0x\_\{0\}\\sim p\_\{0\}, the samplex1=x0\+∫01vτ​\(xτ\)​𝑑τx\_\{1\}=x\_\{0\}\+\\smash\{\\int\_\{0\}^\{1\}\}v^\{\\tau\}\(x\_\{\\tau\}\)\\mathrm\{d\}\\taumatches the target distribution \(x1∼p1x\_\{1\}\\sim p\_\{1\}\)\.

In practice, we use*conditional*flow matching\([Lipman et al\., 2023](https://arxiv.org/html/2610.19129#bib.bib44)\):

ℒCFM\(θ\)=𝔼x1∼𝒟,τ∼U\(0,1\),xτ∼pτ\|1\(⋅∣x1\)\[∥vθτ\(xτ\)−uτ\(xτ∣x1\)∥22\]\.\\displaystyle\\mathcal\{L\}\_\{\\textrm\{CFM\}\}\(\\theta\)=\\mathbb\{E\}\_\{x\_\{1\}\\sim\\mathcal\{D\},\\tau\\sim U\(0,1\),x\_\{\\tau\}\\sim p\_\{\\tau\\mid 1\}\(\\cdot\\mid x\_\{1\}\)\}\\left\[\\lVert v^\{\\tau\}\_\{\\theta\}\(x\_\{\\tau\}\)\-u\_\{\\tau\}\(x\_\{\\tau\}\\mid x\_\{1\}\)\\rVert\_\{2\}^\{2\}\\right\]\.
We use the*conditional optimal transport*\(CondOT\) coupling, drawing a straight line betweenx0x\_\{0\}andx1x\_\{1\},xτ=\(1−τ\)​x0\+τ​x1x\_\{\\tau\}=\(1\-\\tau\)x\_\{0\}\+\\tau x\_\{1\}\. This gives the true conditional velocity field target asuτ​\(xτ∣x1\)=x1−x0u\_\{\\tau\}\(x\_\{\\tau\}\\mid x\_\{1\}\)=x\_\{1\}\-x\_\{0\}\. Our source distributionp0p\_\{0\}is the standard Gaussian,𝒩⁡\(0,I\)\\mathcal\{N\}\(0,I\)\.

### 3\.3Training Policies

We train policies on datasets𝒟=\{\(s0i,a0i,⋯,sTi,aTi\)\}i=1N\\mathcal\{D\}=\\\{\(s\_\{0\}^\{i\},a\_\{0\}^\{i\},\\cdots,s\_\{T\}^\{i\},a\_\{T\}^\{i\}\)\\\}\_\{i=1\}^\{N\}comprised of trajectories collected by a set of unknown policies\{βj​\(a\|s\)\}\\\{\\beta\_\{j\}\(a\|s\)\\\}\. These policies can be marginalized by a mixture policyβ⁡\(a\|s\)=∑jβj​\(a\|s\)​p​\(βj\|s\)\\beta\(a~\|~s\)=\\sum\_\{j\}\\beta\_\{j\}\(a~\|~s\)p\(\\beta\_\{j\}~\|~s\), which is a Markovian policy that has occupancy equal to the expectation under the individual policies:dβ​\(s\)=𝔼p⁡\(βj\)​\[dβj​\(s\)\]\\smash\{d^\{\\beta\}\(s\)=\\mathbb\{E\}\_\{p\(\\beta\_\{j\}\)\}\\left\[d^\{\\beta\_\{j\}\}\(s\)\\right\]\}\([Ghugare et al\., 2024](https://arxiv.org/html/2610.19129#bib.bib37)\)\. In general, there is not equality in the conditional occupancy \(the SM\),Mβ\(⋅\|s,a\)≠𝔼βj∼p⁡\(βj\|s,a\)\[Mβj\(⋅\|s,a\)\]M^\{\\beta\}\(\\cdot\|s,a\)\\neq\\mathbb\{E\}\_\{\\beta\_\{j\}\\sim p\(\\beta\_\{j\}\|s,a\)\}\\left\[M^\{\\beta\_\{j\}\}\(\\cdot\|s,a\)\\right\]\.

We consider the standardgoal\-conditioned BCobjective, which samples a current state from the dataset \(OPENdβ\)d^\{\\beta\}\), and trains the policy to reach goals sampled as future states in the same trajectory, corresponding to samples from the successor measure of an individual policyMβjM^\{\\beta\_\{j\}\}:

ℒBC​\(θ\)=−𝔼s∼dβ,a∼β\(⋅\|s\)βj∼p\(βj\|s,a\),s\+∼Mβj\(⋅\|s,a\)​\[log⁡πθ​\(a\|s,g=s\+\)\]\\displaystyle\\mathcal\{L\}\_\{\\text\{BC\}\}\(\{\\theta\}\)=\-\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}s\\sim d^\{\\beta\},a\\sim\\beta\(\\cdot~\|~s\)\\\\ \\beta\_\{j\}\\sim p\(\\beta\_\{j\}\|s,a\),~s\_\{\+\}\\sim M^\{\\beta\_\{j\}\}\(\\cdot\|s,a\)\\end\{subarray\}\}\\left\[\\log\{\\pi\_\{\\theta\}\}\(a\|s,g=s\_\{\+\}\)\\right\]\(2\)We perform goal\-conditioned behavior cloning using flow matching, representing the policy by a conditional velocity fieldvθτ​\(aτ∣s,g\)v^\{\\tau\}\_\{\\theta\}\(a\_\{\\tau\}\\mid s,g\):

ℒBC​\(θ\)=𝔼\(s,a,s\+\),τ∼U⁡\(0,1\),a0∼p0,aτ=\(1−τ\)​a0\+τ​a​\[‖vθτ​\(aτ∣s,g=s\+\)−\(a−a0\)‖22\]\\displaystyle\\mathcal\{L\}\_\{\\text\{BC\}\}\(\\theta\)=\\mathbb\{E\}\_\{\(s,a,s\_\{\+\}\),\\tau\\sim U\(0,1\),a\_\{0\}\\sim p\_\{0\},a\_\{\\tau\}=\(1\-\\tau\)a\_\{0\}\+\\tau a\}\\left\[\\\|v^\{\\tau\}\_\{\\theta\}\(a\_\{\\tau\}\\mid s,g=s\_\{\+\}\)\-\(a\-a\_\{0\}\)\\\|\_\{2\}^\{2\}\\right\]\(3\)To sample fromπθ\(⋅\|s,g\)\\pi\_\{\\theta\}\(\\cdot\|s,g\), usinga0∼p0a\_\{0\}\\sim p\_\{0\}, we integrate withvθτv\_\{\\theta\}^\{\\tau\}, viaa1=a0\+∫01vθτ​\(aτ\|s,g\)​𝑑τa\_\{1\}=a\_\{0\}\+\\int\_\{0\}^\{1\}v^\{\\tau\}\_\{\\theta\}\(a\_\{\\tau\}\|s,g\)\\mathrm\{d\}\\tau\.

### 3\.4Temporal Difference Flows

Recently, there has been significant interest in developing generative models of the successor measure, called geometric horizon models \(GHMs\), that can sample states from the*geometric horizon*of any given policy\([Janner et al\., 2020](https://arxiv.org/html/2610.19129#bib.bib45);[Thakoor et al\., 2022](https://arxiv.org/html/2610.19129#bib.bib46)\)\. These models can be trained via TD learning, as for all\(s,a\)\(s,a\), the successor measureMγπM^\{\\pi\}\_\{\\gamma\}satisfies:

Mγπ\(⋅\|s,a\)=\(1−γ\)P\(⋅\|s,a\)\+γ𝔼s′∼P\(⋅\|s,a\),a′∼π\(⋅\|s′\)\[Mγπ\(⋅\|s′,a′\)\]\\displaystyle M^\{\\pi\}\_\{\\gamma\}\(\\cdot\|s,a\)=\(1\-\\gamma\)P\(\\cdot\|s,a\)\+\\gamma\\mathbb\{E\}\_\{s^\{\\prime\}\\sim P\(\\cdot\|s,a\),a^\{\\prime\}\\sim\\pi\(\\cdot\|s^\{\\prime\}\)\}\[M^\{\\pi\}\_\{\\gamma\}\(\\cdot\|s^\{\\prime\},a^\{\\prime\}\)\]\(4\)Optimizing this objective is challenging with prior generative modeling techniques such as GANs\([Goodfellow et al\., 2014](https://arxiv.org/html/2610.19129#bib.bib47)\)or normalizing flows\([Rezende and Mohamed, 2015](https://arxiv.org/html/2610.19129#bib.bib48)\), especially at large discountsγ\\gamma\. This led to the proposal to use flow matching to learn GHMs\([Farebrother et al\., 2025](https://arxiv.org/html/2610.19129#bib.bib43)\), which showed strong success in simulated control domains\. TD\-Flows parameterize the GHM as a velocity fieldvϕτ:ℝd×𝒮×𝒜→ℝd\\smash\{v^\{\\tau\}\_\{\\phi\}\}:\\mathbb\{R\}^\{d\}\\times\\mathcal\{S\}\\times\\mathcal\{A\}\\rightarrow\\mathbb\{R\}^\{d\}with parametersϕ\\phi, such that with sourcex0∼p0∈Δ⁡\(ℝd\)x\_\{0\}\\sim p\_\{0\}\\in\\Delta\(\\mathbb\{R\}^\{d\}\), the velocity field induces the flowψτ:ℝd×𝒮×𝒜→ℝd\\psi^\{\\tau\}:\\mathbb\{R\}^\{d\}\\times\\mathcal\{S\}\\times\\mathcal\{A\}\\rightarrow\\mathbb\{R\}^\{d\}:

ψϕτ​\(x0\|s,a\)=x0\+∫0τvϕi​\(ψϕi​\(x0\|s,a\)\|s,a\)​𝑑i\.\\displaystyle\\psi^\{\\tau\}\_\{\\phi\}\(x\_\{0\}\|s,a\)=x\_\{0\}\+\\int\_\{0\}^\{\\tau\}v^\{i\}\_\{\\phi\}\(\\psi^\{i\}\_\{\\phi\}\(x\_\{0\}\|s,a\)\|s,a\)\\mathrm\{d\}i\.\(5\)We aim to learnvϕτv^\{\\tau\}\_\{\\phi\}, such that forx0∼p0x\_\{0\}\\sim p\_\{0\}, the samplesψϕ1​\(x0\|s,a\)\\psi^\{1\}\_\{\\phi\}\(x\_\{0\}\|s,a\)are distributed according toMγπ\(⋅\|s,a\)\\smash\{M^\{\\pi\}\_\{\\gamma\}\(\\cdot\|s,a\)\}\. Thus, we will use the notationMϕπ\(⋅\|s,a\)\\smash\{M^\{\\pi\}\_\{\\phi\}\(\\cdot\|s,a\)\}for our approximate successor measure parameterized byvϕτ\\smash\{v^\{\\tau\}\_\{\\phi\}\}\. Exploiting the Bellman structure of the successor measure,[Farebrother et al\. \(2025\)](https://arxiv.org/html/2610.19129#bib.bib43)show that a flow that models said measure can be learned by minimizing theTD2​\-CFM\\text\{TD\}^\{2\}\\text\{\-CFM\}loss:

ℒTD2\(ϕ;ϕ¯\)=\(1−γ\)ℒtd\-ss\(ϕ\)\+γℒtd\-bs\(ϕ;ϕ¯\),where:\\displaystyle\\mathcal\{L\}\_\{\\text\{TD\}^\{2\}\}\(\\phi;\\bar\{\\phi\}\)=\(1\-\\gamma\)\\mathcal\{L\}\_\{\\textrm\{td\-ss\}\}\(\\phi\)\+\\gamma\\mathcal\{L\}\_\{\\textrm\{td\-bs\}\}\(\\phi;\\bar\{\\phi\}\),\\text\{ where:\}\(6\)ℒtd\-ss​\(ϕ\)=𝔼τ∼U⁡\(0,1\),\(s,a,s′\)∼𝒟β,x0∼p0,xτ=\(1−τ\)​x0\+τ​s′​\[‖vϕτ​\(xτ\|s,a\)−\(s′−x0\)‖2\]\\displaystyle\\mathcal\{L\}\_\{\\textrm\{td\-ss\}\}\(\\phi\)=\\mathbb\{E\}\_\{\\tau\\sim U\(0,1\),\(s,a,s^\{\\prime\}\)\\sim\\mathcal\{D\}^\{\\beta\},x\_\{0\}\\sim p\_\{0\},x\_\{\\tau\}=\(1\-\\tau\)x\_\{0\}\+\\tau s^\{\\prime\}\}\\left\[\\\|v^\{\\tau\}\_\{\\phi\}\(x\_\{\\tau\}\|s,a\)\-\(s^\{\\prime\}\-x\_\{0\}\)\\\|^\{2\}\\right\]ℒtd\-bs\(ϕ\)=𝔼τ∼U\(0,1\),\(s,a,s′\)∼𝒟β,a′∼π\(⋅\|s′\),x0∼p0,xτ=ψτϕ¯\(x0\|s′,a′\)\[∥vϕτ\(xτ\|s,a\)−vϕ¯τ\(xτ\|s′,a′\)∥2\]\\displaystyle\\mathcal\{L\}\_\{\\textrm\{td\-bs\}\}\(\\phi\)=\\mathbb\{E\}\_\{\\tau\\sim U\(0,1\),\(s,a,s^\{\\prime\}\)\\sim\\mathcal\{D\}^\{\\beta\},a^\{\\prime\}\\sim\\pi\(\\cdot\|s^\{\\prime\}\),x\_\{0\}\\sim p\_\{0\},x\_\{\\tau\}=\\psi^\{\\tau\}\_\{\\bar\{\\phi\}\}\(x\_\{0\}\|s^\{\\prime\},a^\{\\prime\}\)\}\\left\[\\\|v^\{\\tau\}\_\{\\phi\}\(x\_\{\\tau\}\|s,a\)\-v^\{\\tau\}\_\{\\bar\{\\phi\}\}\(x\_\{\\tau\}\|s^\{\\prime\},a^\{\\prime\}\)\\\|^\{2\}\\right\]whereϕ¯\\bar\{\\phi\}are target flow parameters, as commonly used in value\-based RL\([Mnih et al\., 2015](https://arxiv.org/html/2610.19129#bib.bib49)\)\. Recalling equation[4](https://arxiv.org/html/2610.19129#S3.E4),ℒtd\-ss\\mathcal\{L\}\_\{\\textrm\{td\-ss\}\}targets the single\-step transition kernel, whileℒtd\-bs\\mathcal\{L\}\_\{\\textrm\{td\-bs\}\}bootstraps using the successor measure at the next state\. Together, this leads to approximatingMγπ\\smash\{M^\{\\pi\}\_\{\\gamma\}\}\.

We can also approximate the successor measureMβ\(⋅\|s,a\)\\smash\{M^\{\\beta\}\(\\cdot\|s,a\)\}of the mixture policyβ\\beta, rather than an arbitrary policyπ\\pi\. While this requires samplinga′∼β\(⋅\|s′\)a^\{\\prime\}\\sim\\beta\(\\cdot\|s^\{\\prime\}\)in the Bellman update, in practice, we will generally bootstrap with the next action in the dataset, yielding the term:

ℒtd\-bs​\(ϕ\)=𝔼τ∼U⁡\(0,1\),\(s,a,s′,a′\)∼𝒟β,x0∼p0,xτ=ψϕ¯τ​\(x0\|s′,a′\)​\[‖vϕτ​\(xτ\|s,a\)−vϕ¯τ​\(xτ\|s′,a′\)‖2\]\\mathcal\{L\}\_\{\\textrm\{td\-bs\}\}\(\\phi\)=\\mathbb\{E\}\_\{\\tau\\sim U\(0,1\),\(s,a,s^\{\\prime\},a^\{\\prime\}\)\\sim\\mathcal\{D\}^\{\\beta\},x\_\{0\}\\sim p\_\{0\},x\_\{\\tau\}=\\psi^\{\\tau\}\_\{\\bar\{\\phi\}\}\(x\_\{0\}\|s^\{\\prime\},a^\{\\prime\}\)\}\\left\[\\\|v^\{\\tau\}\_\{\\phi\}\(x\_\{\\tau\}\|s,a\)\-v^\{\\tau\}\_\{\\bar\{\\phi\}\}\(x\_\{\\tau\}\|s^\{\\prime\},a^\{\\prime\}\)\\\|^\{2\}\\right\]We refer to this objective asℒTD2​\-​β\\mathcal\{L\}\_\{\\text\{TD\}^\{2\}\\text\{\-\}\\beta\}\. This is a simpler objective than the fully off\-policy objective in equation[6](https://arxiv.org/html/2610.19129#S3.E6), as we bootstrap over fully in\-support actions\. If one does not need to stitch across trajectories, we can utilize a simpler approach to modelMγβM^\{\\beta\}\_\{\\gamma\}using empirical Monte Carlo samples of future states\. This allows us to learn a generative model directly with supervised learning:

ℒMC​\(ϕ\)=𝔼\(s,a\)∼𝒟β,k∼geom⁡\(1−γ\),s\+=s→kτ∼U⁡\(0,1\),x0∼p0,xτ=\(1−τ\)​x0\+τ​s\+​\[∥vϕτ​\(xτ∣s,a\)−\(s\+−x0\)∥22\]\.\\displaystyle\\mathcal\{L\}\_\{\\textrm\{MC\}\}\(\\phi\)=\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}\(s,a\)\\sim\\mathcal\{D\}^\{\\beta\},k\\sim\\mathrm\{geom\}\(1\-\\gamma\),s\_\{\+\}=s\_\{\\rightarrow k\}\\\\ \\tau\\sim U\(0,1\),x\_\{0\}\\sim p\_\{0\},x\_\{\\tau\}=\(1\-\\tau\)x\_\{0\}\+\\tau s\_\{\+\}\\end\{subarray\}\}\\left\[\\lVert v^\{\\tau\}\_\{\\phi\}\(x\_\{\\tau\}\\mid s,a\)\-\(s\_\{\+\}\-x\_\{0\}\)\\rVert\_\{2\}^\{2\}\\right\]\.\(7\)Here,s→ks\_\{\\rightarrow k\}denotes the state that iskksteps further fromssalong the sampled trajectory in the dataset\. If we had infinitely long trajectories, collected directly by the mixture policy, this objective would approximateMβ\(⋅\|s,a\)M^\{\\beta\}\(\\cdot\|s,a\)\. However, in practice, we have datasets with truncated horizon, and collected from a set of data\-collecting policies\. This leadsℒMC​\(ϕ\)\\mathcal\{L\}\_\{\\textrm\{MC\}\}\(\\phi\)to instead approximate a mixture of successor measures𝔼βj∼p⁡\(βj\|s,a\)Mγβj\(⋅\|s,a\)\\smash\{\\mathbb\{E\}\_\{\\beta\_\{j\}\\sim p\(\\beta\_\{j\}\|s,a\)\}\{M^\{\\beta\_\{j\}\}\_\{\\gamma\}\}\(\\cdot\|s,a\)\}\([Lawson et al\., 2025](https://arxiv.org/html/2610.19129#bib.bib29)\), preferring the use ofℒTD2​\-​β\\mathcal\{L\}\_\{\\text\{TD\}^\{2\}\\text\{\-\}\\beta\}\.

## 4Improving Long\-Horizon, On\-Policy Generation by Combining MC and TD Sampling

In this section, we propose a variant of the TD\-Flow training procedure \(ℒTD2​\-​β\\mathcal\{L\}\_\{\\text\{TD\}^\{2\}\\text\{\-\}\\beta\}\) to address scaling to long horizon prediction, which is affected by the discountγ\\gammaused during bootstrapping\. Intuitively, fitting the GHM is easier at lower discounts, as the loss is more heavily weighted to the ground truth next\-state prediction, rather than the bootstrap\. Recent work proposes the Temporal Difference Horizon Consistency \(TD\-HC\) objective\([Farebrother et al\., 2026](https://arxiv.org/html/2610.19129#bib.bib21)\), which proposes how to use a "local horizon" GHM \(e\.g\. a GHM conditioned on a discountω<γ\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}<\\gamma\) to obtain an estimate of the successor measure to approximate a higher discount GHM \(γ\{\\color\[rgb\]\{0,0,1\}\\gamma\}\)\([Janner et al\., 2020](https://arxiv.org/html/2610.19129#bib.bib45);[Farebrother et al\., 2026](https://arxiv.org/html/2610.19129#bib.bib21)\), i\.e\.ω<γ\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}<\{\\color\[rgb\]\{0,0,1\}\\gamma\}through a Bellman\-like relation:

Mγπ\(⋅\|s,a\)≜\\displaystyle M^\{\\pi\}\_\{\{\\color\[rgb\]\{0,0,1\}\\gamma\}\}\(\\cdot\|s,a\)\\triangleq\(1−γ\)P\(⋅\|s,a\)\+γ1−γ1−ω𝔼s′∼P\(⋅\|s,a\),a′∼π\(⋅\|s′\)\[Mωπ\(⋅\|s′,a′\)\]\\displaystyle\(1\-\{\\color\[rgb\]\{0,0,1\}\\gamma\}\)P\(\\cdot\|s,a\)\+\{\\color\[rgb\]\{0,0,1\}\\gamma\}\\frac\{1\-\{\\color\[rgb\]\{0,0,1\}\\gamma\}\}\{1\-\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\}\\mathbb\{E\}\_\{s^\{\\prime\}\\sim P\(\\cdot\|s,a\),a^\{\\prime\}\\sim\\pi\(\\cdot\|s^\{\\prime\}\)\}\[M^\{\\pi\}\_\{\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\}\(\\cdot\|s^\{\\prime\},a^\{\\prime\}\)\]\+\\displaystyle\+γγ−ω1−ω𝔼s′∼P\(⋅\|s,a\),a′∼π\(⋅\|s′\),s\+∼Mπω\(⋅\|s′,a′\),a\+∼π\(⋅\|s\+\)\[Mγπ\(⋅\|s\+,a\+\)\]\\displaystyle\{\\color\[rgb\]\{0,0,1\}\\gamma\}\\frac\{\{\\color\[rgb\]\{0,0,1\}\\gamma\}\-\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\}\{1\-\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\}\\mathbb\{E\}\_\{s^\{\\prime\}\\sim P\(\\cdot\|s,a\),a^\{\\prime\}\\sim\\pi\(\\cdot\|s^\{\\prime\}\),s\_\{\+\}\\sim M^\{\\pi\}\_\{\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\}\(\\cdot\|s^\{\\prime\},a^\{\\prime\}\),a\_\{\+\}\\sim\\pi\(\\cdot\|s\_\{\+\}\)\}\[M^\{\\pi\}\_\{\{\\color\[rgb\]\{0,0,1\}\\gamma\}\}\(\\cdot\|s\_\{\+\},a\_\{\+\}\)\]\(8\)This can be converted to a TD\-Flow loss,ℒTDHC\\mathcal\{L\}\_\{\\mathrm\{TDHC\}\}, by modifying equation[6](https://arxiv.org/html/2610.19129#S3.E6), shown in Appendix[B\.1](https://arxiv.org/html/2610.19129#A2.SS1)\.

We propose a variant ofℒTDHC\\mathcal\{L\}\_\{\\mathrm\{TDHC\}\}adapted for the on\-policy setting to approximateMγβM^\{\\beta\}\_\{\{\\color\[rgb\]\{0,0,1\}\\gamma\}\}\. While equation[7](https://arxiv.org/html/2610.19129#S3.E7)supervises the flow with ground\-truth samples, which we will show leads to high\-quality generations, it does not generate samples that can be reached through stitching trajectories\. In contrast, equation[6](https://arxiv.org/html/2610.19129#S3.E6)stitches, but uses sparse ground\-truth signal, which decreases as the discount is increased\. This may make it harder to have higher\-quality samples and to scale to longer horizons\. To obtain the best of both objectives, we consider ahorizon\-consistent objective utilizing MC samplesfrom the dataset\. While we will learnMγβM^\{\\beta\}\_\{\{\\color\[rgb\]\{0,0,1\}\\gamma\}\}, we will not learnMωβM^\{\\beta\}\_\{\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\}, and instead utilize MC samples from dataset trajectories\. This leads to the objective:

ℒMCHC\-​β\(ϕ;ϕ¯\)=\(1−γ\)ℒtd\-ss\(ϕ\)\+γ1−γ1−ωℒmc\(ϕ\)\+γγ−ω1−ωℒtd\-mc\-bs\(ϕ;ϕ¯\),where:\\displaystyle\\mathcal\{L\}\_\{\\text\{MCHC\-\}\\beta\}\(\\phi;\\bar\{\\phi\}\)=\(1\-\{\\color\[rgb\]\{0,0,1\}\\gamma\}\)\\mathcal\{L\}\_\{\\textrm\{td\-ss\}\}\(\\phi\)\+\{\\color\[rgb\]\{0,0,1\}\\gamma\}\\frac\{1\-\{\\color\[rgb\]\{0,0,1\}\\gamma\}\}\{1\-\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\}\\mathcal\{L\}\_\{\\textrm\{mc\}\}\(\\phi\)\+\{\\color\[rgb\]\{0,0,1\}\\gamma\}\\frac\{\{\\color\[rgb\]\{0,0,1\}\\gamma\}\-\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\}\{1\-\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\}\\mathcal\{L\}\_\{\\textrm\{td\-mc\-bs\}\}\(\\phi;\\bar\{\\phi\}\),\\text\{ where:\}\(9\)ℒtd\-ss​\(ϕ\)=𝔼τ∼U⁡\(0,1\),\(s,a,s′\)∼𝒟β,x0∼p0,xτ=\(1−τ\)​x0\+τ​s′​\[‖vϕτ​\(xτ\|s,a\)−\(s′−x0\)‖2\]\\displaystyle\\mathcal\{L\}\_\{\\textrm\{td\-ss\}\}\(\\phi\)=\\mathbb\{E\}\_\{\\tau\\sim U\(0,1\),\(s,a,s^\{\\prime\}\)\\sim\\mathcal\{D\}^\{\\beta\},x\_\{0\}\\sim p\_\{0\},x\_\{\\tau\}=\(1\-\\tau\)x\_\{0\}\+\\tau s^\{\\prime\}\}\\left\[\\\|v^\{\\tau\}\_\{\\phi\}\(x\_\{\\tau\}\|s,a\)\-\(s^\{\\prime\}\-x\_\{0\}\)\\\|^\{2\}\\right\]ℒmc​\(ϕ\)=𝔼τ∼U⁡\(0,1\),\(s,a\)∼𝒟β,k∼geom⁡\(1−ω\)s\+=s→k\+1,x0∼p0,xτ=\(1−τ\)​x0\+τ​s\+​\[‖vϕτ​\(xτ\|s,a\)−\(s\+−x0\)‖2\]\\displaystyle\\mathcal\{L\}\_\{\\textrm\{mc\}\}\(\\phi\)=\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}\\tau\\sim U\(0,1\),\(s,a\)\\sim\\mathcal\{D\}^\{\\beta\},k\\sim\\mathrm\{geom\}\(1\-\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\)\\\\ s\_\{\+\}=s\_\{\\rightarrow k\+1\},x\_\{0\}\\sim p\_\{0\},x\_\{\\tau\}=\(1\-\\tau\)x\_\{0\}\+\\tau s\_\{\+\}\\end\{subarray\}\}\\left\[\\\|v^\{\\tau\}\_\{\\phi\}\(x\_\{\\tau\}\|s,a\)\-\(s\_\{\+\}\-x\_\{0\}\)\\\|^\{2\}\\right\]ℒtd\-mc\-bs​\(ϕ\)=𝔼τ∼U⁡\(0,1\),\(s,a\)∼𝒟β,k∼geom⁡\(1−ω\)s\+,a\+=s→k\+1,a→k\+1,x0∼p0,xτ=ψϕ¯τ​\(x0\|s\+,a\+\)​\[‖vϕτ​\(xτ\|s,a\)−vϕ¯τ​\(xτ\|s\+,a\+\)‖2\]\\displaystyle\\mathcal\{L\}\_\{\\textrm\{td\-mc\-bs\}\}\(\\phi\)=\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}\\tau\\sim U\(0,1\),\(s,a\)\\sim\\mathcal\{D\}^\{\\beta\},k\\sim\\mathrm\{geom\}\(1\-\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\)\\\\ s\_\{\+\},a\_\{\+\}=s\_\{\\rightarrow k\+1\},a\_\{\\rightarrow k\+1\},x\_\{0\}\\sim p\_\{0\},x\_\{\\tau\}=\\psi^\{\\tau\}\_\{\\bar\{\\phi\}\}\(x\_\{0\}\|s\_\{\+\},a\_\{\+\}\)\\end\{subarray\}\}\\left\[\\\|v^\{\\tau\}\_\{\\phi\}\(x\_\{\\tau\}\|s,a\)\-v^\{\\tau\}\_\{\\bar\{\\phi\}\}\(x\_\{\\tau\}\|s\_\{\+\},a\_\{\+\}\)\\\|^\{2\}\\right\]\\
This loss is not just supervised by the next ground truth states, but also future MC samples\. Additionally, instead of bootstrapping from the next statest\+1s\_\{t\+1\}, we will bootstrap from this MC sample, which is obtained as a future state afterst\+1s\_\{t\+1\}, i\.e\.s\+=st\+k\+1s\_\{\+\}=s\_\{t\+k\+1\}, wherek∼geom​\(1−ω\)k\\sim\\text\{geom\}\(1\-\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\)\. As we utilize MC samples, which are both truncated, and are collected by a single policy, they do not correspond to the trueMωβM^\{\\beta\}\_\{\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\}\. Thus, this objective is a biased approximation ofMγβM^\{\\beta\}\_\{\{\\color\[rgb\]\{0,0,1\}\\gamma\}\}asω\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}increases, but allows increasingγ\{\\color\[rgb\]\{0,0,1\}\\gamma\}while still training on higher proportions of ground\-truth data\.

## 5Imagining Goals for Policy Generalization

### 5\.1Generalization Problem

We consider a generalization problem for BC policies\. Following the setting from Section[3\.3](https://arxiv.org/html/2610.19129#S3.SS3), we train on trajectories collected by unknown policies\{βj​\(a\|s\)\}\\\{\\beta\_\{j\}\(a\|s\)\\\}, with mixture policyβ⁡\(a\|s\)\\beta\(a\|s\)\. We will now define training and test distributions forπθ\\pi\_\{\\theta\}equation[2](https://arxiv.org/html/2610.19129#S3.E2), following[Ghugare et al\. \(2024\)](https://arxiv.org/html/2610.19129#bib.bib37):

Ridπ=𝔼s∼dβ,βj∼p\(βj\|s\),sg∼Mβj\(⋅\|s\)​\[Rπθ​\(s,sg\)\]⏟in\-distribution training tasks,Rgenπ=𝔼s∼dβ,sg∼Mβ\(⋅\|s\)​\[Rπθ​\(s,sg\)\]⏟tasks requiring compositional generalization\\displaystyle R^\{\\pi\}\_\{\\text\{id\}\}=\\underbrace\{\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}s\\sim d^\{\\beta\},\\beta\_\{j\}\\sim p\(\\beta\_\{j\}\|s\),s\_\{g\}\\sim M^\{\\beta\_\{j\}\}\(\\cdot\|s\)\\end\{subarray\}\}\\bigl\[R^\{\{\\pi\_\{\\theta\}\}\}\(s,s\_\{g\}\)\\bigr\]\}\_\{\\text\{ in\-distribution training tasks \}\},R^\{\\pi\}\_\{\\text\{gen\}\}=\\underbrace\{\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}s\\sim d^\{\\beta\},s\_\{g\}\\sim M^\{\\beta\}\(\\cdot\|s\)\\end\{subarray\}\}\\bigl\[R^\{\{\\pi\_\{\\theta\}\}\}\(s,s\_\{g\}\)\\bigr\]\}\_\{\\text\{tasks requiring compositional generalization\}\}\(10\)WithRθπ​\(s,sg\)R^\{\\pi\}\_\{\\theta\}\(s,s\_\{g\}\)defining the success rate of the policy to reachsgs\_\{g\}fromss,RidπR^\{\\pi\}\_\{\\text\{id\}\}averages performance over in\-distribution tasks andRgenπR^\{\\pi\}\_\{\\text\{gen\}\}describes average success on a class of out\-of\-distribution tasks requiring compositional generalization\. These include tasks which may not be completed by a single data collecting policyβj\\beta\_\{j\}, but still have sufficient data coverage to complete through composing or stitching trajectories collected by different policies, i\.e\. reachable withβ\\beta\.

While equation[2](https://arxiv.org/html/2610.19129#S3.E2)corresponds to improving onRidπR^\{\\pi\}\_\{\\text\{id\}\}, the BC objective does not coverRgenπR^\{\\pi\}\_\{\\text\{gen\}\}, leading to poor empirical performance\. We can see that if we had access to goals coming from the successor measure of the mixture policyβ\\betaduring training, rather than that of individual policies, our policy could learn to reach goals required for betterRgenπR^\{\\pi\}\_\{\\text\{gen\}\}\.

### 5\.2Method

To address this type of policy generalization, we can directly approximate and sample from the target generalization distribution\. We note that because we do not have full rollouts from the mixture policy, we must utilize TD learning to stitch, relying onℒTD2​\-​β\\mathcal\{L\}\_\{\\text\{TD\}^\{2\}\\text\{\-\}\\beta\}orℒMCHC\-​β\\mathcal\{L\}\_\{\\text\{MCHC\}\\text\{\-\}\\beta\}, rather thanℒMC​\(ϕ\)\\mathcal\{L\}\_\{\\textrm\{MC\}\}\(\\phi\)\. Thus, we trainon\-policy TD\-Flows, which will yield an approximationMϕβ\(⋅\|s,a\)\\smash\{M^\{\\beta\}\_\{\\phi\}\(\\cdot\|s,a\)\}for the successor measure of the mixture policy\. After obtainingMϕβ\\smash\{M^\{\\beta\}\_\{\\phi\}\}, we performgoal augmentationwhere we train on a mix of original goals plus samples from the generative model, covering the compositional distribution:

maxπθ⁡𝔼s∼dβ,a∼β\(a\|s\),s\+∼Paug\(⋅\|s,a\)​\[log⁡πθ​\(a\|s,g=s\+\)\]\\displaystyle\\max\_\{\\pi\_\{\\theta\}\}\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}s\\sim d^\{\\beta\},a\\sim\\beta\(a\|s\),~s\_\{\+\}\\sim P\_\{\\text\{aug\}\}\(\\cdot\|s,a\)\\end\{subarray\}\}\\left\[\\log\{\\pi\_\{\\theta\}\}\(a\|s,g=s\_\{\+\}\)\\right\]\(11\)wherePaug\(⋅\|s,a\)≜αMβϕ\(⋅\|s,a\)⏟imagined compositional samples\+\(1−α\)𝔼p⁡\(βj\|s,a\)Mβj\(⋅\|s,a\)⏟ground truth future states\\displaystyle\\mathrm\{where~\}P\_\{\\text\{aug\}\}\(\\cdot\|s,a\)\\triangleq\\underbrace\{\\alpha\{\{M\}^\{\\beta\}\_\{\\phi\}\(\\cdot\|s,a\)\}\}\_\{\\text\{imagined compositional samples\}\}\+\\underbrace\{\(1\-\\alpha\)\\mathbb\{E\}\_\{p\(\\beta\_\{j\}\|s,a\)\}\{M^\{\\beta\_\{j\}\}\(\\cdot\|s,a\)\}\}\_\{\\text\{ground truth future states\}\}\(12\)
We generally refer our goal augmentation procedure asTD\-Aug, and by default useℒMCHC\-​β\\mathcal\{L\}\_\{\\text\{MCHC\-\}\\beta\}\. We can also instantiate our method withℒTD2​\-​β\\mathcal\{L\}\_\{\\text\{TD\}^\{2\}\\text\{\-\}\\beta\}, which we denote later withTD\-Aug2\\textbf\{TD\-Aug\}\_\{2\}\. We provide an overview of our method in Algorithm[1](https://arxiv.org/html/2610.19129#alg1)and Figure[1](https://arxiv.org/html/2610.19129#S1.F1)\. An advantage of this procedure is that we are not limited to a fixed number of stitches, or heuristic sampling strategies\([Ghugare et al\., 2024](https://arxiv.org/html/2610.19129#bib.bib37)\)\. We also performO⁡\(1\)O\(1\)sampling steps per sample to cover an expected horizonH=\(1−γ\)−1H=\(1\-\\gamma\)^\{\-1\}, while one\-step models needO⁡\(H\)O\(H\)steps, potentially resulting in compounding error and longer time for rollouts\.

We relate this procedure to offline goal\-conditioned RL, specifically 1\-step policy improvement over the mixture policy with value\-weighted regression\([Peters and Schaal, 2007](https://arxiv.org/html/2610.19129#bib.bib30);[Wang et al\., 2018](https://arxiv.org/html/2610.19129#bib.bib5);[Peng et al\., 2019](https://arxiv.org/html/2610.19129#bib.bib4);[Wang et al\., 2020](https://arxiv.org/html/2610.19129#bib.bib6)\), for a brief proof, see Appendix[C](https://arxiv.org/html/2610.19129#A3):

###### Theorem 1\(Equivalence to 1\-step improvement\)\.

Withα=1\\alpha=1, and exact samples fromMβ\(⋅\|s,a\)M^\{\\beta\}\(\\cdot\|s,a\), our augmentation strategy in equation[11](https://arxiv.org/html/2610.19129#S5.E11)is equivalent to value\-weighted regression used for one\-step improvement over behavior policyβ\\beta:

θ∗∈arg⁡maxθ​𝔼s∼dβ,s\+∼dβa∼β\(⋅∣s\)​\[Qβ​\(s,a,g=s\+\)dβ​\(s\+\)​log⁡πθ​\(a∣s,g=s\+\)\]\\displaystyle\\theta^\{\*\}\\in\\arg\\max\_\{\\theta\}\\;\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}s\\sim d^\{\\beta\},s\_\{\+\}\\sim d^\{\\beta\}\\\\ a\\sim\\beta\(\\cdot\\mid s\)\\end\{subarray\}\}\\left\[\\frac\{Q^\{\\beta\}\(s,a,g=s\_\{\+\}\)\}\{d^\{\\beta\}\(s\_\{\+\}\)\}\\,\\log\\pi\_\{\\theta\}\(a\\mid s,g=s\_\{\+\}\)\\right\]\(13\)

## 6Experiments

Figure 3:Evaluating sampling with increasing horizon\. We compute earth mover’s distance \(EMD\) between empirical samples and generated samples on held\-out evaluation data forpointmaze\-large\-stitch\.In this section, we show thatTD\-Augcan improve combinatorial generalization across a variety ofstitchtasks in the OGBench\([Park et al\., 2025a](https://arxiv.org/html/2610.19129#bib.bib36)\)benchmark suite\. In Section[6\.1](https://arxiv.org/html/2610.19129#S6.SS1), we \(a\) evaluate the quality of generated samples as horizon scales for different GHM objectives on low\-dimensional navigation tasks, and \(b\) show thatTD\-Augis useful for improving BC generalization, even on long\-horizon tasks requiring many stitches\. Then, \(c\) we show that our method can scale to visual environments by learning GHMs in latent space in Section[6\.2](https://arxiv.org/html/2610.19129#S6.SS2), and Section[6\.3](https://arxiv.org/html/2610.19129#S6.SS3), we demonstrate that we can also improve stitching on manipulation tasks\.

### 6\.1State\-based Environments

We first evaluate our method onpointmaze\-stitchnavigation tasks from OGBench, which allow us to isolate the challenge of long\-horizon stitching from high\-dimensional dynamics\. In these tasks, agents are trained on short trajectories, which travel up to 4 maze cells\. Then, agents are evaluated on longer compositional tasks such as those that span the whole width of the maze\. We find that the original stitching datasets, especially for point, are not terminated upon goal\-reaching, resulting in many transitions that idle at the goal\. Thus, we re\-collect the stitch datasets that terminate upon reaching a sub\-goal\. On new datasets, baselines perform similarly, or slightly better than before\.

To begin our study, we firstcompare the quality of generative models\.Qualitatively, we provide visualizations onpointmaze\-\{medium,large,teleport,giant\}\-stitchin Figure[8](https://arxiv.org/html/2610.19129#A4.F8)\. We can see that MC Flows can accurately capture training trajectories and multi\-modality in the dataset\. We can see this captures both multi\-modality in the policy, and environment stochasticity, as shown inteleport, where upon entering a black hole, the agent is randomly sent to a random white hole\. However, the MC Flow does not implicitly stitch across trajectories\. On the other hand, the TD\-based objectives \(TD2​\-​β,MCHC\-​β\{\\text\{TD\}^\{2\}\\text\{\-\}\\beta\},\{\\text\{MCHC\}\\text\{\-\}\\beta\}\), can generate samples that require stitching, capturing states reachable beyond than the training trajectory horizon\. However, theTD2​\-​β\\text\{TD\}^\{2\}\\text\{\-\}\\betasamples are more diffuse, or may go backwards from the direction implied by the ground truth samples\. In contrast, theMCHC\-​β\\text\{MCHC\}\\text\{\-\}\\betasamples balance implicit stitching with the quality of MC samples\.

Next, we perform aquantitativecomparison between model samples on a held\-out evaluation dataset onpointmaze\-large\-stitch\. From an initial state, we can see how models \(trained on short trajectories\) generate samples that are closer in distribution to longer, held\-out trajectories, measuring stitching\. This is performed by sampling initial \(state,action\) tuples from a held\-out dataset\. Then, from these tuples we sample future states geometrically, obtaining empirical samples\. From these same initial states, we sample from our model, as shown in Figure[7](https://arxiv.org/html/2610.19129#A4.F7)\. Then, we compute the earth mover’s distance \(EMD\) between these two sets of points, and then average across different initial positions\. We repeat this for models trained at various effective horizonsH=\(1−γ\)−1H=\(1\-\\gamma\)^\{\-1\}, measuring against ground truth samples at the corresponding horizon, as shown in Figure[3](https://arxiv.org/html/2610.19129#S6.F3)\. In addition to the GHM based models, we also generate samples from autoregressive one\-step prediction, which we further discuss in Appendix[B\.3](https://arxiv.org/html/2610.19129#A2.SS3)\. While some models have similar EMD at the initial horizon \(100\), we can see that the errors ofTD2​\-​β\\text\{TD\}^\{2\}\\text\{\-\}\\beta, and 1\-step baselines, quickly increases with horizon\. Additionally, the EMD of MC increases because it is trained on short trajectories, and thus has no mechanism to stitch\. In contrast,MCHC\-​β\\text\{MCHC\}\\text\{\-\}\\betahas the lowest EMD as horizon increases, showing the benefit of mixing both MC targets and TD when training on\-policy GHMs\.

Table 1:Generalization through compositional data augmentation\.With goal conditioned flow behavior cloning \(GCFBC\), we compare training on the original dataset to various data augmentation strategies\. We calculate success rate averaged over the last 3 checkpoints, and then report the mean, std over 3 seeds\. Weboldthe top performing methods\.![Refer to caption](https://arxiv.org/html/2610.19129v1/main_tau_latent.png)Figure 4:Extension ofTD\-Augto visual environments\.Next, we study how geometric horizon models canempiricallyimprove a policy’scombinatorial generalizationcapabilities in Table[1](https://arxiv.org/html/2610.19129#S6.T1)\. Without any augmentation, a goal\-conditioned flow\-based behavior cloning \(GCFBC\) baseline\([Park et al\., 2025b](https://arxiv.org/html/2610.19129#bib.bib42)\)obtains moderate zero\-shot generalization in smaller mazes, but struggles in larger mazes as the horizon is increased\. To improve generalization, we apply our generative goal augmentation procedure,TD\-Aug\. In addition toTD\-AugusingℒMCHC\-​β\\mathcal\{L\}\_\{\\text\{MCHC\}\\text\{\-\}\\beta\}, we compare toTD\-Aug2\\text\{TD\-Aug\}\_\{2\}usingℒTD2​\-​β\\mathcal\{L\}\_\{\\text\{TD\}^\{2\}\\text\{\-\}\\beta\}, as well asMC\-AugusingℒMC\\mathcal\{L\}\_\{\\text\{MC\}\}\. We compare to clustering\-based goal augmentation \(Cluster\-Aug\)\([Ghugare et al\., 2024](https://arxiv.org/html/2610.19129#bib.bib37)\), which can create stitches between two trajectories, and against autoregressive 1\-step prediction with history to sample from the SM \(Dyn\-Aug\), as utilized in Figure[3](https://arxiv.org/html/2610.19129#S6.F3)\. For these methods, we list additional training details in Appendix[A](https://arxiv.org/html/2610.19129#A1), including hyperparameters for BC in Table[5](https://arxiv.org/html/2610.19129#A1.T5), and for TD\-Flows in Table[4](https://arxiv.org/html/2610.19129#A1.T4)\. Although the task horizon varies, we use the same hyperparameters for each environment inpointmazeforTD\-Aug\. This includes the effective horizon, which we set to500500\(γ=0\.998\\gamma=0\.998\), while forMC\-AugandTD\-Aug2\\text\{TD\-Aug\}\_\{2\}we use100100\(γ=0\.99\\gamma=0\.99\), as they perform worse at higher horizons\. After training GHMs, all policies are trained for10610^\{6\}steps, and each method is trained with batches using goals that are of a proportion1−α1\-\\alphaof future states in the same trajectory \(i\.e\. hindsight goal sampling\), and with proportionα\\alphawith generated goals\. The BC baseline usesα=0\\alpha=0, and augmentation usesα=0\.5\\alpha=0\.5\. We can see thatTD\-Augsignificantly improves generalization, including on long\-horizon tasks, such aspointmaze\-giant\-stitch\. In contrast,MC\-Augmaintains close to original performance, demonstrating that generating additional synthetic short\-horizon samples does not improve generalization\. We also see that withTD\-Aug, GCFBC outperforms prior offline GCRL methods inpointmaze, including GCIQL\([Kostrikov et al\., 2022](https://arxiv.org/html/2610.19129#bib.bib39)\), QRL\([Wang et al\., 2023](https://arxiv.org/html/2610.19129#bib.bib40)\), Contrastive RL \(CRL\)\([Eysenbach et al\., 2022](https://arxiv.org/html/2610.19129#bib.bib35)\), and HIQL\([Park et al\., 2023](https://arxiv.org/html/2610.19129#bib.bib31)\)\. Observing the weaker performance of CRL onstitchtasks, which also uses MC samples like GCBC, we show contrastive learning on generated samples leads to improved performance in Appendix[E](https://arxiv.org/html/2610.19129#A5)\.

Table 2:Visual generalization through compositional data augmentation\.With autoencoder \(AE\) GCFBC, we compare training on the original visual datasets to various data augmentation strategies\. We calculate success rate averaged over the last 3 checkpoints, and then report the mean, std over 3 seeds\. Weboldthe highest performing method overall, and highlight inbrownthe top BC method\.![Refer to caption](https://arxiv.org/html/2610.19129v1/visual_antmaze_medium_ae_mchc_samples_seed4.png)Figure 5:Samples forvisual\-antmaze\-medium\.We visualize samples generated withMCHC\-​β\\text\{MCHC\-\}\\betaatγ=0\.995\\gamma=0\.995\. We display both real, geometrically sampled, states, as well as reconstructions from the GHM’s predicted latents using the AE decoder\.
### 6\.2Composition in Visual Navigation

In this section, we extendTD\-Augtovisual environmentswithvisual\-antmaze\-stitch\. Instead of learning the GHM in pixel space, we first train an autoencoder \(AE\), obtaining a low\-dimensional latent space𝒵\\mathcal\{Z\}\. We provide details for training the AE in Table[6](https://arxiv.org/html/2610.19129#A1.T6)\. Then, we utilize the same architectures for the GHM, and policy, as performed in the non\-visual setting \(Section[6\.1](https://arxiv.org/html/2610.19129#S6.SS1)\), but now operate over latent space, as shown in Figure[4](https://arxiv.org/html/2610.19129#S6.F4)\. As a baseline, we report both standard BC in OGBench \(GCBC\), using early state\-goal fusion, and an autoencoder BC baseline, labeledAE\-GCFBC\. The standard baseline concatenates state and goal images, and passes them through a small IMPALA\-based\([Espeholt et al\., 2018](https://arxiv.org/html/2610.19129#bib.bib9)\)CNN\. In contrast, the autoencoder BC baseline separately encodes goals and states into𝒵\\mathcal\{Z\}, and then concatenates them to predict actions\. This architecture is preferable with goal augmentation, as we generate goals in latent space\([Dao et al\., 2023](https://arxiv.org/html/2610.19129#bib.bib51)\)and feed them to the policy without decoding\. With this setup, in Table[2](https://arxiv.org/html/2610.19129#S6.T2), we report adding augmentation to AE\-GCFBC\. We can see that without augmentation, this architecture under\-performs GCBC\. While we improve withMC\-Aug, we can see thatℒTD2​\-​β\\mathcal\{L\}\_\{\\text\{TD\}^\{2\}\\text\{\-\}\\beta\}, used inTD\-Aug2\\text\{TD\-Aug\}\_\{2\}leads to less performance, which is in contrast to the non\-visual setting\. We attribute this to the difficulty of bootstrapping in higher\-dimensional approximate \(𝒵\\mathcal\{Z\}\) space\. Thus, we find the combination of MC and TD to be the most effective, whereTD\-Augimproves AE\-GCFBC and outperforms the GCBC baseline\.

Table 3:Manipulation with compositional data augmentation\.We evaluate GCFBC with action chunking on OGBench manipulation tasks, using theoraclegoal representation\. We use the originalplaydataset, and then introduce a newstitchdataset, consisting of only short trajectories\. Then, we addTD\-Augto close the generalization gap fromstitch\. We calculate average success on the last 3 checkpoints, and report the mean, std over 3 seeds\. Wehighlightthe top method onstitch\.
### 6\.3Composition in Manipulation

In this section, we evaluate TD\-Aug on OGBench manipulation tasks requiring stitching, which is a challenging setting for modeling the successor measure\. We first considercube\-\{double,triple\}\-play,puzzle\-4x4\-playtasks withoraclegoal representation\([Park et al\., 2025a](https://arxiv.org/html/2610.19129#bib.bib36);[Park et al\., 2026a](https://arxiv.org/html/2610.19129#bib.bib1)\)\. In addition to the originalplaydatasets, we construct a new challengingstitchdataset, where each trajectory contains one atomic task \(i\.e\. pick\-and\-place, or button press\), but the policy is evaluated on the original, long\-horizon, multi\-step tasks\. While GCFBC using action chunking,π\(at:t\+h\|s,g\)\\pi\(a\_\{t:t\+h\}\|s,g\), obtains good average success\(44%\)\(44\\%\)on theplaytasks, it struggles when training on thestitch\(13%\)\(13\\%\)\. Thus, we apply TD\-Aug to close the gap\. We also train GHMs, with action chunking,Mϕβ\(⋅\|s,at:t\+h\)\\smash\{M^\{\\beta\}\_\{\\phi\}\(\\cdot\|s,a\_\{t:t\+h\}\)\}, and predicth​\-sizedh\\text\{\-sized\}jumps in the full state space, with additional details in Appendix[B\.2](https://arxiv.org/html/2610.19129#A2.SS2)\. Onstitch, with TD\-Aug, GCFBC’s performance substantially increases from13→5213\\rightarrow 52\. We provide visualizations of the GHM’s samples in Figure[10](https://arxiv.org/html/2610.19129#A4.F10)\.

## 7Conclusion

In this work, we have shown that TD\-based geometric horizon models allow for enhanced generalization for behavior cloning through goal augmentation\. There are a variety of directions that could build on this work\. For example, behavior cloning is a strong pre\-training algorithm that has been shown to scale well in real\-world robotics settings\([O’Neill et al\., 2024](https://arxiv.org/html/2610.19129#bib.bib11);[Black et al\., 2024](https://arxiv.org/html/2610.19129#bib.bib50);[Intelligence et al\., 2025](https://arxiv.org/html/2610.19129#bib.bib10)\)\. Therefore, a key open question is how to efficiently and stably scale TD\-based GHM training to high\-dimensional, partially\-observable robotics tasks\.

### Acknowledgments

We want to acknowledge funding support from Natural Sciences and Engineering Research Council \(NSERC\) of Canada, Google Research, Fonds de recherche du Québec \(FRQNT\), and the Canadian Institute for Advanced Research \(CIFAR\) and compute support from Digital Research Alliance of Canada, Mila IDT, and NVidia\. We also thank Jesse Farebrother and Pablo Samuel Castro for advice throughout the project\.

## References

- Andrychowiczet al\.\(2017\)M\. Andrychowicz, F\. Wolski, A\. Ray, J\. Schneider, R\. Fong, P\. Welinder, B\. McGrew, J\. Tobin, P\. Abbeel, and W\. ZarembaHindsight experience replay\.Advances in neural information processing systems30\.Cited by:[§2](https://arxiv.org/html/2610.19129#S2.p1.1)\.
- Blacket al\.\(2024\)K\. Black, N\. Brown, D\. Driess, A\. Esmail, M\. Equi, C\. Finn, N\. Fusai, L\. Groom, K\. Hausman, B\. Ichter, S\. Jakubczak, T\. Jones, L\. Ke, S\. Levine, A\. Li\-Bell, M\. Mothukuri, S\. Nair, K\. Pertsch, L\. X\. Shi, J\. Tanner, Q\. Vuong, A\. Walling, H\. Wang, and U\. Zhilinskyπ0\\pi\_\{0\}: A vision\-language\-action flow model for general robot control\.External Links:2410\.24164,[Link](https://arxiv.org/abs/2410.24164)Cited by:[§7](https://arxiv.org/html/2610.19129#S7.p1.1)\.
- Bruceet al\.\(2024\)J\. Bruce, M\. D\. Dennis, A\. Edwards, J\. Parker\-Holder, Y\. Shi, E\. Hughes, M\. Lai, A\. Mavalankar, R\. Steigerwald, C\. Apps,et al\.Genie: generative interactive environments\.InForty\-first International Conference on Machine Learning,Cited by:[§1](https://arxiv.org/html/2610.19129#S1.p1.1)\.
- Chunget al\.\(2026\)H\. Chung, J\. Lee, and S\. OhOffline reinforcement learning with universal horizon models\.External Links:2605\.15603,[Link](https://arxiv.org/abs/2605.15603)Cited by:[§2](https://arxiv.org/html/2610.19129#S2.p4.1)\.
- Daoet al\.\(2023\)Q\. Dao, H\. Phung, B\. Nguyen, and A\. TranFlow matching in latent space\.arXiv preprint arXiv:2307\.08698\.Cited by:[§6\.2](https://arxiv.org/html/2610.19129#S6.SS2.p1.1)\.
- Dayan \(1993\)P\. DayanImproving generalization for temporal difference learning: the successor representation\.Neural computation5\(4\),pp\. 613–624\.Cited by:[§2](https://arxiv.org/html/2610.19129#S2.p4.1)\.
- Espeholtet al\.\(2018\)L\. Espeholt, H\. Soyer, R\. Munos, K\. Simonyan, V\. Mnih, T\. Ward, Y\. Doron, V\. Firoiu, T\. Harley, I\. Dunning,et al\.Impala: scalable distributed deep\-rl with importance weighted actor\-learner architectures\.InInternational conference on machine learning,pp\. 1407–1416\.Cited by:[§6\.2](https://arxiv.org/html/2610.19129#S6.SS2.p1.1)\.
- Eysenbachet al\.\(2022\)B\. Eysenbach, T\. Zhang, S\. Levine, and R\. R\. SalakhutdinovContrastive learning as goal\-conditioned reinforcement learning\.Advances in Neural Information Processing Systems35,pp\. 35603–35620\.Cited by:[Appendix E](https://arxiv.org/html/2610.19129#A5.p1.1),[§1](https://arxiv.org/html/2610.19129#S1.p3.1),[§6\.1](https://arxiv.org/html/2610.19129#S6.SS1.p4.1)\.
- Farebrotheret al\.\(2026\)J\. Farebrother, M\. Pirotta, A\. Tirinzoni, M\. G\. Bellemare, A\. Lazaric, and A\. TouatiCompositional planning with jumpy world models\.External Links:2602\.19634,[Link](https://arxiv.org/abs/2602.19634)Cited by:[§B\.1](https://arxiv.org/html/2610.19129#A2.SS1.p1.1),[§1](https://arxiv.org/html/2610.19129#S1.p3.1),[§2](https://arxiv.org/html/2610.19129#S2.p4.1),[§4](https://arxiv.org/html/2610.19129#S4.p1.1)\.
- Farebrotheret al\.\(2025\)J\. Farebrother, M\. Pirotta, A\. Tirinzoni, R\. Munos, A\. Lazaric, and A\. TouatiTemporal difference flows\.InForty\-second International Conference on Machine Learning,External Links:[Link](https://openreview.net/forum?id=j6H7c3aQyb)Cited by:[Table 4](https://arxiv.org/html/2610.19129#A1.T4),[§1](https://arxiv.org/html/2610.19129#S1.p2.1),[§1](https://arxiv.org/html/2610.19129#S1.p3.1),[§2](https://arxiv.org/html/2610.19129#S2.p4.1),[§3\.2](https://arxiv.org/html/2610.19129#S3.SS2.p1.1),[§3\.4](https://arxiv.org/html/2610.19129#S3.SS4.p1.2),[§3\.4](https://arxiv.org/html/2610.19129#S3.SS4.p1.3)\.
- Feinberget al\.\(2018\)V\. Feinberg, A\. Wan, I\. Stoica, M\. I\. Jordan, J\. E\. Gonzalez, and S\. LevineModel\-based value estimation for efficient model\-free reinforcement learning\.arXiv preprint arXiv:1803\.00101\.Cited by:[§2](https://arxiv.org/html/2610.19129#S2.p1.1),[§2](https://arxiv.org/html/2610.19129#S2.p4.1)\.
- Fujimoto and Gu \(2021\)S\. Fujimoto and S\. S\. GuA minimalist approach to offline reinforcement learning\.Advances in neural information processing systems34,pp\. 20132–20145\.Cited by:[Appendix E](https://arxiv.org/html/2610.19129#A5.p2.1)\.
- Ghugareet al\.\(2024\)R\. Ghugare, M\. Geist, G\. Berseth, and B\. EysenbachClosing the gap between TD learning and supervised learning \- a generalisation point of view\.\.InThe Twelfth International Conference on Learning Representations,External Links:[Link](https://arxiv.org/abs/2401.11237)Cited by:[Table 7](https://arxiv.org/html/2610.19129#A1.T7),[§1](https://arxiv.org/html/2610.19129#S1.p2.1),[§2](https://arxiv.org/html/2610.19129#S2.p1.1),[§3\.3](https://arxiv.org/html/2610.19129#S3.SS3.p1.1),[§5\.1](https://arxiv.org/html/2610.19129#S5.SS1.p1.1),[§5\.2](https://arxiv.org/html/2610.19129#S5.SS2.p2.1),[§6\.1](https://arxiv.org/html/2610.19129#S6.SS1.p4.1)\.
- Goodfellowet al\.\(2014\)I\. J\. Goodfellow, J\. Pouget\-Abadie, M\. Mirza, B\. Xu, D\. Warde\-Farley, S\. Ozair, A\. Courville, and Y\. BengioGenerative adversarial nets\.InAdvances in Neural Information Processing Systems,Z\. Ghahramani, M\. Welling, C\. Cortes, N\. Lawrence, and K\.Q\. Weinberger \(Eds\.\),Vol\.27,pp\.\.External Links:[Link](https://proceedings.neurips.cc/paper_files/paper/2014/file/f033ed80deb0234979a61f95710dbe25-Paper.pdf)Cited by:[§3\.4](https://arxiv.org/html/2610.19129#S3.SS4.p1.2)\.
- Ha and Schmidhuber \(2018\)D\. Ha and J\. SchmidhuberRecurrent world models facilitate policy evolution\.Advances in neural information processing systems31\.Cited by:[§1](https://arxiv.org/html/2610.19129#S1.p1.1),[§2](https://arxiv.org/html/2610.19129#S2.p1.1)\.
- Hafneret al\.\(2020\)D\. Hafner, T\. Lillicrap, J\. Ba, and M\. NorouziDream to control: learning behaviors by latent imagination\.External Links:1912\.01603,[Link](https://arxiv.org/abs/1912.01603)Cited by:[§2](https://arxiv.org/html/2610.19129#S2.p1.1)\.
- Hafneret al\.\(2025\)D\. Hafner, W\. Yan, and T\. LillicrapTraining agents inside of scalable world models\.arXiv preprint arXiv:2509\.24527\.Cited by:[§1](https://arxiv.org/html/2610.19129#S1.p1.1)\.
- Intelligenceet al\.\(2025\)P\. Intelligence, A\. Amin, R\. Aniceto, A\. Balakrishna, K\. Black, K\. Conley, G\. Connors, J\. Darpinian, K\. Dhabalia, J\. DiCarlo, D\. Driess, M\. Equi, A\. Esmail, Y\. Fang, C\. Finn, C\. Glossop, T\. Godden, I\. Goryachev, L\. Groom, H\. Hancock, K\. Hausman, G\. Hussein, B\. Ichter, S\. Jakubczak, R\. Jen, T\. Jones, B\. Katz, L\. Ke, C\. Kuchi, M\. Lamb, D\. LeBlanc, S\. Levine, A\. Li\-Bell, Y\. Lu, V\. Mano, M\. Mothukuri, S\. Nair, K\. Pertsch, A\. Z\. Ren, C\. Sharma, L\. X\. Shi, L\. Smith, J\. T\. Springenberg, K\. Stachowicz, W\. Stoeckle, A\. Swerdlow, J\. Tanner, M\. Torne, Q\. Vuong, A\. Walling, H\. Wang, B\. Williams, S\. Yoo, L\. Yu, U\. Zhilinsky, and Z\. Zhouπ0\.6∗\\pi^\{\*\}\_\{0\.6\}: A vla that learns from experience\.External Links:2511\.14759,[Link](https://arxiv.org/abs/2511.14759)Cited by:[§7](https://arxiv.org/html/2610.19129#S7.p1.1)\.
- Janneret al\.\(2022\)M\. Janner, Y\. Du, J\. B\. Tenenbaum, and S\. LevinePlanning with diffusion for flexible behavior synthesis\.External Links:2205\.09991,[Link](https://arxiv.org/abs/2205.09991)Cited by:[§2](https://arxiv.org/html/2610.19129#S2.p2.1)\.
- Janneret al\.\(2020\)M\. Janner, I\. Mordatch, and S\. Levineγ\\gamma\-Models: generative temporal difference learning for infinite\-horizon prediction\.InAdvances in Neural Information Processing Systems,Cited by:[§1](https://arxiv.org/html/2610.19129#S1.p2.1),[§2](https://arxiv.org/html/2610.19129#S2.p4.1),[§3\.4](https://arxiv.org/html/2610.19129#S3.SS4.p1.1),[§4](https://arxiv.org/html/2610.19129#S4.p1.1)\.
- Kaelbling \(1993\)L\. P\. KaelblingLearning to achieve goals\.InInternational Joint Conference on Artificial Intelligence,Cited by:[§2](https://arxiv.org/html/2610.19129#S2.p1.1)\.
- Kostrikovet al\.\(2022\)I\. Kostrikov, A\. Nair, and S\. LevineOffline reinforcement learning with implicit q\-learning\.InInternational Conference on Learning Representations,External Links:[Link](https://arxiv.org/abs/2110.06169)Cited by:[§6\.1](https://arxiv.org/html/2610.19129#S6.SS1.p4.1)\.
- Kostrikovet al\.\(2020\)I\. Kostrikov, D\. Yarats, and R\. FergusImage augmentation is all you need: regularizing deep reinforcement learning from pixels\.arXiv preprint arXiv:2004\.13649\.Cited by:[Table 6](https://arxiv.org/html/2610.19129#A1.T6)\.
- Lawsonet al\.\(2025\)D\. Lawson, A\. Hugessen, C\. Cloutier, G\. Berseth, and K\. KhetarpalSelf\-predictive representations for combinatorial generalization in behavioral cloning\.External Links:2506\.10137,[Link](https://arxiv.org/abs/2506.10137)Cited by:[§2](https://arxiv.org/html/2610.19129#S2.p3.1),[§3\.4](https://arxiv.org/html/2610.19129#S3.SS4.p2.3)\.
- Lee and Choi \(2025\)K\. Lee and J\. ChoiState\-covering trajectory stitching for diffusion planners\.External Links:2506\.00895,[Link](https://arxiv.org/abs/2506.00895)Cited by:[§2](https://arxiv.org/html/2610.19129#S2.p2.1)\.
- Lipmanet al\.\(2023\)Y\. Lipman, R\. T\. Q\. Chen, H\. Ben\-Hamu, M\. Nickel, and M\. LeFlow matching for generative modeling\.InThe Eleventh International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=PqvMRDCJT9t)Cited by:[§3\.2](https://arxiv.org/html/2610.19129#S3.SS2.p1.1),[§3\.2](https://arxiv.org/html/2610.19129#S3.SS2.p2.1)\.
- Lloyd \(1982\)S\. LloydLeast squares quantization in pcm\.IEEE transactions on information theory28\(2\),pp\. 129–137\.Cited by:[Table 7](https://arxiv.org/html/2610.19129#A1.T7)\.
- Luet al\.\(2023\)C\. Lu, P\. J\. Ball, Y\. W\. Teh, and J\. Parker\-HolderSynthetic experience replay\.InThirty\-seventh Conference on Neural Information Processing Systems,External Links:[Link](https://openreview.net/forum?id=6jNQ1AY1Uf)Cited by:[§2](https://arxiv.org/html/2610.19129#S2.p1.1)\.
- Luoet al\.\(2025\)Y\. Luo, U\. A\. Mishra, Y\. Du, and D\. XuGenerative trajectory stitching through diffusion composition\.External Links:2503\.05153,[Link](https://arxiv.org/abs/2503.05153)Cited by:[§2](https://arxiv.org/html/2610.19129#S2.p2.1)\.
- Mnihet al\.\(2015\)V\. Mnih, K\. Kavukcuoglu, D\. Silver, A\. A\. Rusu, J\. Veness, M\. G\. Bellemare, A\. Graves, M\. Riedmiller, A\. K\. Fidjeland, G\. Ostrovski, S\. Petersen, C\. Beattie, A\. Sadik, I\. Antonoglou, H\. King, D\. Kumaran, D\. Wierstra, S\. Legg, and D\. HassabisHuman\-level control through deep reinforcement learning\.Nature518\(7540\),pp\. 529–533\.External Links:ISSN 00280836,[Link](http://dx.doi.org/10.1038/nature14236)Cited by:[§3\.4](https://arxiv.org/html/2610.19129#S3.SS4.p1.4)\.
- Myerset al\.\(2025\)V\. Myers, B\. C\. Zheng, A\. Dragan, K\. Fang, and S\. LevineTemporal representation alignment: successor features enable emergent compositionality in robot instruction following\.External Links:2502\.05454,[Link](https://arxiv.org/abs/2502.05454)Cited by:[§2](https://arxiv.org/html/2610.19129#S2.p3.1)\.
- O’Neillet al\.\(2024\)A\. O’Neill, A\. Rehman, A\. Maddukuri, A\. Gupta, A\. Padalkar, A\. Lee, A\. Pooley, A\. Gupta, A\. Mandlekar, A\. Jain,et al\.Open x\-embodiment: robotic learning datasets and rt\-x models: open x\-embodiment collaboration 0\.In2024 IEEE International Conference on Robotics and Automation \(ICRA\),pp\. 6892–6903\.Cited by:[§7](https://arxiv.org/html/2610.19129#S7.p1.1)\.
- Parket al\.\(2025a\)S\. Park, K\. Frans, B\. Eysenbach, and S\. LevineOGBench: benchmarking offline goal\-conditioned rl\.InInternational Conference on Learning Representations \(ICLR\),External Links:[Link](https://arxiv.org/abs/2410.20092)Cited by:[Appendix E](https://arxiv.org/html/2610.19129#A5.p1.1),[Appendix E](https://arxiv.org/html/2610.19129#A5.p2.2),[§1](https://arxiv.org/html/2610.19129#S1.p4.1),[§6\.3](https://arxiv.org/html/2610.19129#S6.SS3.p1.1),[§6](https://arxiv.org/html/2610.19129#S6.p1.1)\.
- Parket al\.\(2026a\)S\. Park, K\. Frans, D\. Mann, B\. Eysenbach, A\. Kumar, and S\. LevineHorizon reduction makes rl scalable\.Advances in Neural Information Processing Systems38,pp\. 8350–8389\.Cited by:[§6\.3](https://arxiv.org/html/2610.19129#S6.SS3.p1.1)\.
- Parket al\.\(2023\)S\. Park, D\. Ghosh, B\. Eysenbach, and S\. LevineHiql: offline goal\-conditioned rl with latent states as actions\.Advances in Neural Information Processing Systems36,pp\. 34866–34891\.Cited by:[§6\.1](https://arxiv.org/html/2610.19129#S6.SS1.p4.1)\.
- Parket al\.\(2025b\)S\. Park, Q\. Li, and S\. LevineFlow q\-learning\.InForty\-second International Conference on Machine Learning,External Links:[Link](https://openreview.net/forum?id=KVf2SFL1pi)Cited by:[§6\.1](https://arxiv.org/html/2610.19129#S6.SS1.p4.1)\.
- Parket al\.\(2026b\)S\. Park, D\. Mann, and S\. LevineDual goal representations\.External Links:2510\.06714,[Link](https://arxiv.org/abs/2510.06714)Cited by:[§2](https://arxiv.org/html/2610.19129#S2.p3.1)\.
- Penget al\.\(2019\)X\. B\. Peng, A\. Kumar, G\. Zhang, and S\. LevineAdvantage\-weighted regression: simple and scalable off\-policy reinforcement learning\.arXiv preprint arXiv:1910\.00177\.Cited by:[§5\.2](https://arxiv.org/html/2610.19129#S5.SS2.p3.1)\.
- Peters and Schaal \(2007\)J\. Peters and S\. SchaalReinforcement learning by reward\-weighted regression for operational space control\.InProceedings of the 24th international conference on Machine learning,pp\. 745–750\.Cited by:[§5\.2](https://arxiv.org/html/2610.19129#S5.SS2.p3.1)\.
- Rezende and Mohamed \(2015\)D\. Rezende and S\. MohamedVariational inference with normalizing flows\.InProceedings of the 32nd International Conference on Machine Learning,F\. Bach and D\. Blei \(Eds\.\),Proceedings of Machine Learning Research, Vol\.37,Lille, France,pp\. 1530–1538\.External Links:[Link](https://proceedings.mlr.press/v37/rezende15.html)Cited by:[§3\.4](https://arxiv.org/html/2610.19129#S3.SS4.p1.2)\.
- Sutton \(1991\)R\. S\. SuttonDyna, an integrated architecture for learning, planning, and reacting\.ACM Sigart Bulletin2\(4\),pp\. 160–163\.Cited by:[§1](https://arxiv.org/html/2610.19129#S1.p1.1),[§2](https://arxiv.org/html/2610.19129#S2.p1.1)\.
- Thakooret al\.\(2022\)S\. Thakoor, M\. Rowland, D\. Borsa, W\. Dabney, R\. Munos, and A\. BarretoGeneralised policy improvement with geometric policy composition\.InProceedings of the 39th International Conference on Machine Learning,K\. Chaudhuri, S\. Jegelka, L\. Song, C\. Szepesvari, G\. Niu, and S\. Sabato \(Eds\.\),Proceedings of Machine Learning Research, Vol\.162,pp\. 21272–21307\.External Links:[Link](https://proceedings.mlr.press/v162/thakoor22a.html)Cited by:[§1](https://arxiv.org/html/2610.19129#S1.p2.1),[§2](https://arxiv.org/html/2610.19129#S2.p4.1),[§3\.4](https://arxiv.org/html/2610.19129#S3.SS4.p1.1)\.
- Tomaret al\.\(2024\)M\. Tomar, P\. Hansen\-Estruch, P\. Bachman, A\. Lamb, J\. Langford, M\. E\. Taylor, and S\. LevineVideo occupancy models\.External Links:2407\.09533,[Link](https://arxiv.org/abs/2407.09533)Cited by:[§2](https://arxiv.org/html/2610.19129#S2.p4.1)\.
- Wanget al\.\(2026\)K\. Wang, I\. Javali, M\. Bortkiewicz, T\. Trzcinski, and B\. Eysenbach1000 layer networks for self\-supervised rl: scaling depth can enable new goal\-reaching capabilities\.Advances in Neural Information Processing Systems38,pp\. 157643–157670\.Cited by:[Appendix E](https://arxiv.org/html/2610.19129#A5.p2.2)\.
- Wanget al\.\(2018\)Q\. Wang, J\. Xiong, L\. Han, H\. Liu, T\. Zhang,et al\.Exponentially weighted imitation learning for batched historical data\.Advances in Neural Information Processing Systems31\.Cited by:[§5\.2](https://arxiv.org/html/2610.19129#S5.SS2.p3.1)\.
- Wanget al\.\(2023\)T\. Wang, A\. Torralba, P\. Isola, and A\. ZhangOptimal goal\-reaching reinforcement learning via quasimetric learning\.InInternational Conference on Machine Learning,External Links:[Link](https://arxiv.org/abs/2304.01203)Cited by:[§6\.1](https://arxiv.org/html/2610.19129#S6.SS1.p4.1)\.
- Wanget al\.\(2020\)Z\. Wang, A\. Novikov, K\. Zolna, J\. S\. Merel, J\. T\. Springenberg, S\. E\. Reed, B\. Shahriari, N\. Siegel, C\. Gulcehre, N\. Heess,et al\.Critic regularized regression\.Advances in Neural Information Processing Systems33,pp\. 7768–7778\.Cited by:[§5\.2](https://arxiv.org/html/2610.19129#S5.SS2.p3.1)\.
- Wiltzeret al\.\(2024\)H\. Wiltzer, J\. Farebrother, A\. Gretton, Y\. Tang, A\. Barreto, W\. Dabney, M\. G\. Bellemare, and M\. RowlandA distributional analogue to the successor representation\.External Links:2402\.08530,[Link](https://arxiv.org/abs/2402.08530)Cited by:[§2](https://arxiv.org/html/2610.19129#S2.p4.1)\.
- Yaratset al\.\(2021\)D\. Yarats, A\. Zhang, I\. Kostrikov, B\. Amos, J\. Pineau, and R\. FergusImproving sample efficiency in model\-free reinforcement learning from images\.InProceedings of the aaai conference on artificial intelligence,Vol\.35,pp\. 10674–10681\.Cited by:[Table 6](https://arxiv.org/html/2610.19129#A1.T6)\.
- Zhenget al\.\(2026\)C\. Zheng, S\. Park, S\. Levine, and B\. EysenbachIntention\-conditioned flow occupancy models\.External Links:2506\.08902,[Link](https://arxiv.org/abs/2506.08902)Cited by:[§2](https://arxiv.org/html/2610.19129#S2.p4.1)\.
- Zhouet al\.\(2024\)Z\. Zhou, C\. Zhu, R\. Zhou, Q\. Cui, A\. Gupta, and S\. S\. DuFree from bellman completeness: trajectory stitching via model\-based return\-conditioned supervised learning\.InThe Twelfth International Conference on Learning Representations,External Links:[Link](https://arxiv.org/abs/2310.19308)Cited by:[§2](https://arxiv.org/html/2610.19129#S2.p1.1)\.

## Appendix ATraining Details

HyperparameterFlow MatchingODE SolverMidpointTrain ODEd​τ\\mathrm\{d\}\\tau0\.10\.1Sample ODEd​τ\\mathrm\{d\}\\tau0\.10\.1Network \(U\-Net\)tt\-Positional Embedding Dim256256tt\-Positional Embedding MLP\(256,256\)\(256,256\)Block Dimension\(512,512,512\)\(512,512,512\)Encoder MLP\(512,512,512\)\(512,512,512\)activationMishOptimizer \(AdamW\)AdamWβ1\\beta\_\{1\}0\.90\.9AdamWβ2\\beta\_\{2\}0\.9990\.999AdamWϵ\\epsilon10−410^\{\-4\}Learning Rate10−410^\{\-4\}Weight Decay10−310^\{\-3\}TrainingTarget Network EMA10−310^\{\-3\}Gradient StepsTDState,Visual3MManipulation1MMC, MCHC1MBatch Size1024Discount \(γ\\gamma\)MC, TD0\.990\.99MCHCState0\.9980\.998Visual0\.9950\.995Manipulation0\.990\.99Discount \(ω\\omega\)State,Visual0\.990\.99Manipulation0\.980\.98Table 4:Hyperparameters for TD\-Flow training, following similar architecture details to[Farebrother et al\. \(2025\)](https://arxiv.org/html/2610.19129#bib.bib43)\.Table 5:Hyperparameters for GCFBC training\.HyperparameterFlow MatchingODE SolverEulerSample ODEd​τ\\mathrm\{d\}\\tau0\.10\.1Network \(MLP\)Hidden Layers\(1024,1024,1024,1024\)\(1024,1024,1024,1024\)activationGELUOptimizer \(Adam\)Adamβ1\\beta\_\{1\}0\.90\.9Adamβ2\\beta\_\{2\}0\.9990\.999Learning RateState, Visual3×10−43\\times 10^\{\-4\}Manipulation3×10−53\\times 10^\{\-5\}Weight Decay00TrainingGradient StepsState1MManipulation2MVisual500kBatch SizeState1024Manipulation1024Visual256Augmentation \(only\)Batch Proportion \(α\\alpha\)0\.50\.5Table 6:Hyperparameters for visual autoencoder training\. For our autoencoder, we use a similar architecture to\([Kostrikov et al\., 2020](https://arxiv.org/html/2610.19129#bib.bib12);[Yarats et al\., 2021](https://arxiv.org/html/2610.19129#bib.bib13)\)\. For the encoder, we utilize two convolutional layers, with ReLU activations after each layer\. The output of the convolutional layers is flattened, and linearly mapped to the latent dimension\|z\|=50\|z\|=50, followed by LayerNorm and tanh activation\. The decoder utilizes a similar architecture, utilizing deconvolution layers, with the same kernel size, followed by ReLU \(except for the final layer\)\.HyperparameterArchitectureConvolution Layers22Channels3232Kernel\(3×3\)\(3\\times 3\)ActivationReLULatent Dimension,\|z\|\|z\|5050Optimizer \(Adam\)Adamβ1\\beta\_\{1\}0\.90\.9Adamβ2\\beta\_\{2\}0\.9990\.999Learning Rate10−410^\{\-4\}Weight Decay00TrainingGradient Steps100kBatch Size1024Table 7:Hyperparameters for clustering\-based goal\-augmentation\.We follow\([Ghugare et al\., 2024](https://arxiv.org/html/2610.19129#bib.bib37)\)and use k\-means clustering\([Lloyd, 1982](https://arxiv.org/html/2610.19129#bib.bib7)\), withkkset as below\.
## Appendix BAdditional Objectives

### B\.1TD\-HC

For reference, we provide the original Temporal Difference Horizon Consistency \(TDHC\) objective\([Farebrother et al\., 2026](https://arxiv.org/html/2610.19129#bib.bib21)\), which we modify to use MC samples in Section[4](https://arxiv.org/html/2610.19129#S4):

ℒTDHC​\(ϕ,ϕ¯\)=\(1−γ\)​ℒtd\-ss​\(ϕ,γ\)\+γ​1−γ1−ω​ℒtd\-bs​\(ϕ,ϕ¯,γ,ω\)\+γ​γ−ω1−ω​ℒtd\-hc\-bs​\(ϕ,ϕ¯,γ,ω\)\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{TDHC\}\}\(\\phi;\\bar\{\\phi\}\)=\(1\-\{\\color\[rgb\]\{0,0,1\}\\gamma\}\)\\mathcal\{L\}\_\{\\textrm\{td\-ss\}\}\(\\phi;\{\\color\[rgb\]\{0,0,1\}\\gamma\}\)\+\{\\color\[rgb\]\{0,0,1\}\\gamma\}\\frac\{1\-\{\\color\[rgb\]\{0,0,1\}\\gamma\}\}\{1\-\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\}\\mathcal\{L\}\_\{\\textrm\{td\-bs\}\}\(\\phi;\\bar\{\\phi\},\{\\color\[rgb\]\{0,0,1\}\\gamma\},\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\)\+\{\\color\[rgb\]\{0,0,1\}\\gamma\}\\frac\{\{\\color\[rgb\]\{0,0,1\}\\gamma\}\-\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\}\{1\-\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\}\\mathcal\{L\}\_\{\\textrm\{td\-hc\-bs\}\}\(\\phi;\\bar\{\\phi\},\{\\color\[rgb\]\{0,0,1\}\\gamma\},\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\)\(14\)ℒtd\-ss=𝔼τ∼U⁡\(0,1\),\(s,a,s′\)∼𝒟βx0∼p0,xτ=\(1−τ\)​x0\+τ​s′​\[‖vϕτ​\(xτ\|s,a,γ\)−\(s′−x0\)‖2\]\\displaystyle\\mathcal\{L\}\_\{\\textrm\{td\-ss\}\}=\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}\\tau\\sim U\(0,1\),\(s,a,s^\{\\prime\}\)\\sim\\mathcal\{D\}^\{\\beta\}\\\\ x\_\{0\}\\sim p\_\{0\},x\_\{\\tau\}=\(1\-\\tau\)x\_\{0\}\+\\tau s^\{\\prime\}\\end\{subarray\}\}\\left\[\\\|v^\{\\tau\}\_\{\\phi\}\(x\_\{\\tau\}\|s,a,\{\\color\[rgb\]\{0,0,1\}\\gamma\}\)\-\(s^\{\\prime\}\-x\_\{0\}\)\\\|^\{2\}\\right\]ℒtd\-bs=𝔼τ∼U⁡\(0,1\),\(s,a,s′\)∼𝒟β,a′∼π⁡\(s′\)x0∼p0,xτ=ψϕ¯τ​\(x0\|s′,a′,ω\)​\[‖vϕτ​\(xτ\|s,a,γ\)−vϕ¯τ​\(xτ\|s′,a′,ω\)‖2\]\\displaystyle\\mathcal\{L\}\_\{\\textrm\{td\-bs\}\}=\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}\\tau\\sim U\(0,1\),\(s,a,s^\{\\prime\}\)\\sim\\mathcal\{D\}^\{\\beta\},a^\{\\prime\}\\sim\\pi\(s^\{\\prime\}\)\\\\ x\_\{0\}\\sim p\_\{0\},x\_\{\\tau\}=\\psi^\{\\tau\}\_\{\\bar\{\\phi\}\}\(x\_\{0\}\|s^\{\\prime\},a^\{\\prime\},\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\)\\end\{subarray\}\}\\left\[\\\|v^\{\\tau\}\_\{\\phi\}\(x\_\{\\tau\}\|s,a,\{\\color\[rgb\]\{0,0,1\}\\gamma\}\)\-v^\{\\tau\}\_\{\\bar\{\\phi\}\}\(x\_\{\\tau\}\|s^\{\\prime\},a^\{\\prime\},\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\)\\\|^\{2\}\\right\]ℒtd\-hc\-bs=𝔼τ∼U⁡\(0,1\),\(s,a,s′\)∼𝒟β,a′∼π⁡\(s′\)x01,x02∼p0,s\+=ψϕ¯1​\(x01\|s′,a′,ω\),a\+∼π⁡\(s\+\)xτ=ψϕ¯τ​\(x02\|s\+,a\+,γ\)​\[‖vϕτ​\(xτ\|s,a,γ\)−vϕ¯τ​\(xτ\|s\+,a\+,γ\)‖2\]\\displaystyle\\mathcal\{L\}\_\{\\textrm\{td\-hc\-bs\}\}=\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}\\tau\\sim U\(0,1\),\(s,a,s^\{\\prime\}\)\\sim\\mathcal\{D\}^\{\\beta\},a^\{\\prime\}\\sim\\pi\(s^\{\\prime\}\)\\\\ x\_\{0\}^\{1\},x\_\{0\}^\{2\}\\sim p\_\{0\},s\_\{\+\}=\\psi^\{1\}\_\{\\bar\{\\phi\}\}\(x\_\{0\}^\{1\}\|s^\{\\prime\},a^\{\\prime\},\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\),a\_\{\+\}\\sim\\pi\(s\_\{\+\}\)\\\\ x\_\{\\tau\}=\\psi^\{\\tau\}\_\{\\bar\{\\phi\}\}\(x\_\{0\}^\{2\}\|s\_\{\+\},a\_\{\+\},\{\\color\[rgb\]\{0,0,1\}\\gamma\}\)\\end\{subarray\}\}\\left\[\\\|v^\{\\tau\}\_\{\\phi\}\(x\_\{\\tau\}\|s,a,\{\\color\[rgb\]\{0,0,1\}\\gamma\}\)\-v^\{\\tau\}\_\{\\bar\{\\phi\}\}\(x\_\{\\tau\}\|s\_\{\+\},a\_\{\+\},\{\\color\[rgb\]\{0,0,1\}\\gamma\}\)\\\|^\{2\}\\right\]\\Withγ=ω\{\\color\[rgb\]\{0,0,1\}\\gamma\}=\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}, the originalTD2\\mathrm\{TD\}^\{2\}objective equation[6](https://arxiv.org/html/2610.19129#S3.E6)is recovered\. To train with the TDHC objective, the flow is trained with discount conditioning, and in each batch, the discount is sampled uniformlyγ∈\[γmin,γmax\]\{\\color\[rgb\]\{0,0,1\}\\gamma\}\\in\[\\gamma\_\{\\mathrm\{min\}\},\\gamma\_\{\\mathrm\{max\}\}\]andω∈\[γmin,γ\]\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\\in\[\\gamma\_\{\\mathrm\{min\}\},\{\\color\[rgb\]\{0,0,1\}\\gamma\}\]\. Instead, the MCHC objective equation[9](https://arxiv.org/html/2610.19129#S4.E9)we use fixesγ=γmax\{\\color\[rgb\]\{0,0,1\}\\gamma\}=\\gamma\_\{\\mathrm\{max\}\}, and samples\(s\+,a\+\)∼Dβ\(s\_\{\+\},a\_\{\+\}\)\\sim D^\{\\beta\}, wheres\+=st\+ks\_\{\+\}=s\_\{t\+k\},k∼geom⁡\(1−ω\)k\\sim\\mathrm\{geom\}\(1\-\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\), replacing having to learn and sample withvϕτ\(⋅\|⋅,ω\)v\_\{\\phi\}^\{\\tau\}\(\\cdot\|\\cdot,\{\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\}\)\.

### B\.2GHMs for Manipulation

For manipulation tasks, we modify both GCFBC and the GHM to utilize action chunks of sizeHHas shown in equation[15](https://arxiv.org/html/2610.19129#A2.E15)\. For the GHM we also predict future states starting from after the action chunk, at increments matching the chunk size\. We also proportionally reduceγ,ω\\gamma,\\omega\. For example, withh=10h=10, we reduceγ\\gammafrom0\.990\.99to0\.90\.9, to match the same expected horizon\. Additionally, we bootstrap with samples from a learned prior,βθ\(⋅\|s\+\)\\beta\_\{\\theta\}\(\\cdot\|s\_\{\+\}\), predicting the next action chunk, rather than usinga\+a\_\{\+\}\. This is useful, as it can allow for bootstrapping from the end states in a trajectory, where an action chunk is not available\.

ℒMCHC\-​β​\(ϕ,ϕ¯\)=\(1−γ\)​ℒtd\-ss​\(ϕ\)\+γ​1−γ1−ω​ℒmc​\(ϕ\)\+γ​γ−ω1−ω​ℒtd\-mc\-bs​\(ϕ,ϕ¯\)\\displaystyle\\mathcal\{L\}\_\{\\text\{MCHC\-\}\\beta\}\(\\phi;\\bar\{\\phi\}\)=\(1\-\{\\color\[rgb\]\{0,0,1\}\\gamma\}\)\\mathcal\{L\}\_\{\\textrm\{td\-ss\}\}\(\\phi\)\+\{\\color\[rgb\]\{0,0,1\}\\gamma\}\\frac\{1\-\{\\color\[rgb\]\{0,0,1\}\\gamma\}\}\{1\-\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\}\\mathcal\{L\}\_\{\\textrm\{mc\}\}\(\\phi\)\+\{\\color\[rgb\]\{0,0,1\}\\gamma\}\\frac\{\{\\color\[rgb\]\{0,0,1\}\\gamma\}\-\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\}\{1\-\{\\color\[rgb\]\{1,0\.5,0\}\\omega\}\}\\mathcal\{L\}\_\{\\textrm\{td\-mc\-bs\}\}\(\\phi;\\bar\{\\phi\}\)\(15\)ℒtd\-ss\(ϕ\)=𝔼τ∼U\(0,1\),\(st,at:t\+h,st\+h\)∼𝒟βx0∼p0,xτ=\(1−τ\)​x0\+τ​st\+h\[∥vϕτ\(xτ\|s,at:t\+h\)−\(st\+h−x0\)∥2\]\\displaystyle\\mathcal\{L\}\_\{\\textrm\{td\-ss\}\}\(\\phi\)=\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}\\tau\\sim U\(0,1\),\(s\_\{t\},a\_\{t:t\+h\},s\_\{t\+h\}\)\\sim\\mathcal\{D\}^\{\\beta\}\\\\ x\_\{0\}\\sim p\_\{0\},x\_\{\\tau\}=\(1\-\\tau\)x\_\{0\}\+\\tau s\_\{t\+h\}\\end\{subarray\}\}\\left\[\\\|v^\{\\tau\}\_\{\\phi\}\(x\_\{\\tau\}\|s,a\_\{t:t\+h\}\)\-\(s\_\{t\+h\}\-x\_\{0\}\)\\\|^\{2\}\\right\]ℒmc\(ϕ\)=𝔼τ∼U\(0,1\),\(st,at:t\+h,s\+\)∼𝒟βx0∼p0,xτ=\(1−τ\)​x0\+τ​s\+\[∥vϕτ\(xτ\|s,at:t\+h\)−\(s\+−x0\)∥2\]\\displaystyle\\mathcal\{L\}\_\{\\textrm\{mc\}\}\(\\phi\)=\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}\\tau\\sim U\(0,1\),\(s\_\{t\},a\_\{t:t\+h\},s\_\{\+\}\)\\sim\\mathcal\{D\}^\{\\beta\}\\\\ x\_\{0\}\\sim p\_\{0\},x\_\{\\tau\}=\(1\-\\tau\)x\_\{0\}\+\\tau s\_\{\+\}\\end\{subarray\}\}\\left\[\\\|v^\{\\tau\}\_\{\\phi\}\(x\_\{\\tau\}\|s,a\_\{t:t\+h\}\)\-\(s\_\{\+\}\-x\_\{0\}\)\\\|^\{2\}\\right\]ℒtd\-mc\-bs\(ϕ\)=𝔼τ∼U\(0,1\),\(st,at:t\+h,s\+\)∼𝒟βx0∼p0,a\+∼βθ\(⋅\|s\+\),xτ=ψτϕ¯\(x0\|s\+,a\+\)\[∥vϕτ\(xτ\|st,at:t\+h\)−vϕ¯τ\(xτ\|s\+,a\+\)∥2\]\\displaystyle\\mathcal\{L\}\_\{\\textrm\{td\-mc\-bs\}\}\(\\phi\)=\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}\\tau\\sim U\(0,1\),\(s\_\{t\},a\_\{t:t\+h\},s\_\{\+\}\)\\sim\\mathcal\{D\}^\{\\beta\}\\\\ x\_\{0\}\\sim p\_\{0\},a\_\{\+\}\\sim\\beta\_\{\\theta\}\(\\cdot\|s\_\{\+\}\),x\_\{\\tau\}=\\psi^\{\\tau\}\_\{\\bar\{\\phi\}\}\(x\_\{0\}\|s\_\{\+\},a\_\{\+\}\)\\end\{subarray\}\}\\left\[\\\|v^\{\\tau\}\_\{\\phi\}\(x\_\{\\tau\}\|s\_\{t\},a\_\{t:t\+h\}\)\-v^\{\\tau\}\_\{\\bar\{\\phi\}\}\(x\_\{\\tau\}\|s\_\{\+\},a\_\{\+\}\)\\\|^\{2\}\\right\]\\where:s\+=st\+h⁡\(1\+k\),k∼geom\(1−ω\),anda\+also uses chunk sizeh=10\.\\displaystyle\\text\{where: \}s\_\{\+\}=s\_\{t\+h\(1\+k\)\},k\\sim\\mathrm\{geom\}\(1\-\{\\color\[rgb\]\{0,0,1\}\\omega\}\),\\text\{ and \}a\_\{\+\}\\text\{ also uses chunk size \}h=10\.

### B\.3One\-step Dynamics Model

As a baseline to geometric horizon models, we utilize auto\-regressive prediction with one\-step models to sample from the successor measure\. This follows from first sampling a random offsetk∼geom⁡\(1−γ\)k\\sim\\mathrm\{geom\}\(1\-\\gamma\)\. We then perform autoregressive prediction forkksteps, first samplingst\+1∼p\(⋅\|s,a\)s\_\{t\+1\}\\sim p\(\\cdot\|s,a\), and then samplingk−1k\-1times fromst\+i\+1∼pβ\(⋅\|st\+i\)s\_\{t\+i\+1\}\\sim p^\{\\beta\}\(\\cdot\|s\_\{t\+i\}\), fori=1i=1tok−1k\-1\. We parameterizep\(⋅\|s,a\)p\(\\cdot\|s,a\)andpβ\(⋅\|s\)p^\{\\beta\}\(\\cdot\|s\)via velocity fieldvϕτ​\(xτ\|s,a\)v^\{\\tau\}\_\{\\phi\}\(x\_\{\\tau\}\|s,a\), trained via flow matching:

ℒ1\-step​\(ϕ\)\\displaystyle\\mathcal\{L\}\_\{\\textrm\{1\-step\}\}\(\\phi\)=𝔼τ∼𝒰⁡\(0,1\),\(s,a,s′\)∼𝒟β,am∼mask⁡\(a\)x0∼p0,xτ=\(1−τ\)​x0\+τ​s′​\[‖vϕτ​\(xτ∣s,am\)−\(s′−x0\)‖22\]\\displaystyle=\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}\\tau\\sim\\mathcal\{U\}\(0,1\),\\,\(s,a,s^\{\\prime\}\)\\sim\\mathcal\{D\}^\{\\beta\},\\,a\_\{m\}\\sim\\mathrm\{mask\}\(a\)\\\\ x\_\{0\}\\sim p\_\{0\},\\,x\_\{\\tau\}=\(1\-\\tau\)x\_\{0\}\+\\tau s^\{\\prime\}\\end\{subarray\}\}\\left\[\\left\\\|v^\{\\tau\}\_\{\\phi\}\(x\_\{\\tau\}\\mid s,a\_\{m\}\)\-\(s^\{\\prime\}\-x\_\{0\}\)\\right\\\|\_\{2\}^\{2\}\\right\]\(16\)We utilize a standard 1\-step flow matching objective, however we condition onam∼mask⁡\(a\)a\_\{m\}\\sim\\mathrm\{mask\}\(a\), where with probability0\.50\.5, the action is replaced by a learned embedding \(shared for all actions\)\. Thus, this objective trains both forms ofpp\. We utilize the same architecture as GHM as described in Table[6](https://arxiv.org/html/2610.19129#A1.T6), and train for 1 million steps\.

While in theory, this objective can perfectly capture the successor measure, we found issues with compounding error during autoregressive prediction as shown in Figure[3](https://arxiv.org/html/2610.19129#S6.F3)\. This is particularly due topβ\(⋅\|s\)p^\{\\beta\}\(\\cdot\|s\)being memory\-less, i\.e\. after the first action, there is no conditioning information inpβ\(⋅\|s\)p^\{\\beta\}\(\\cdot\|s\)that captures an intention \(such as running to the right\)\. Thus, we find a stronger performing baseline incorporates history, leading to:

ℒhistory 1\-step​\(ϕ\)\\displaystyle\\mathcal\{L\}\_\{\\text\{history 1\-step\}\}\(\\phi\)=𝔼τ∼𝒰\(0,1\),\(st−h\+1:t,a,st\+1\)∼𝒟βx0∼p0,xτ=\(1−τ\)​x0\+τ​st\+1am∼mask⁡\(a\)\[‖vϕτ\(xτ∣st−h\+1:t,am\)−\(st\+1−x0\)‖22\]\\displaystyle=\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}\\tau\\sim\\mathcal\{U\}\(0,1\),\\,\(s\_\{t\-h\+1:t\},a,s\_\{t\+1\}\)\\sim\\mathcal\{D\}^\{\\beta\}\\\\ x\_\{0\}\\sim p\_\{0\},\\,x\_\{\\tau\}=\(1\-\\tau\)x\_\{0\}\+\\tau s\_\{t\+1\}\\\\ a\_\{m\}\\sim\\mathrm\{mask\}\(a\)\\end\{subarray\}\}\\left\[\\left\\\|v^\{\\tau\}\_\{\\phi\}\(x\_\{\\tau\}\\mid s\_\{t\-h\+1:t\},a\_\{m\}\)\-\(s\_\{t\+1\}\-x\_\{0\}\)\\right\\\|\_\{2\}^\{2\}\\right\]\(17\)Where we utilize a history length ofh=4h=4, as used in Figures[3](https://arxiv.org/html/2610.19129#S6.F3),[7](https://arxiv.org/html/2610.19129#A4.F7)\. When sampling with a specific discountγ\\gamma, with expected horizonH=11−γH=\\frac\{1\}\{1\-\\gamma\}, we capkkfor rollouts at2​H2H\. For example, ifγ=0\.99\\gamma=0\.99, we capk=min⁡\(geom⁡\(1−0\.99\),200\)k=\\min\(\\mathrm\{geom\}\(1\-0\.99\),200\)\.

As GHMs can sample in O\(1\) steps, when training a policy, we can efficiently sample a fresh batch of goals,s\+i∼Mγβ\(⋅\|st,at\)s\_\{\+\}^\{i\}\\sim M^\{\\beta\}\_\{\\gamma\}\(\\cdot\|s\_\{t\},a\_\{t\}\), for all batch indices, every time we form a batch for the policy update\. However, this is not realistic with one\-step models atγ=0\.99\\gamma=0\.99\. Thus, we instead cache a2​H=2002H=200sized trajectory from each state from the dataset, storing\(N,200\)\(N,200\)states, whereNNis the original number of transitions \(e\.g\.1,000,0001,000,000\)\. Then, while training the policy, we geometrically sample future states from this window\. We utilize this procedure when performing augmentation with one\-step dynamics models \(Dyn\-Aug\) as in Table[1](https://arxiv.org/html/2610.19129#S6.T1)\.

## Appendix CRelation to 1\-step policy improvement

###### Proof\.

𝔼s∼dβ​\(s\)a∼β⁡\(a∣s\)s\+∼dβ​\(s\+\)​\[Qβ​\(s,a,s\+\)dβ​\(s\+\)​log⁡πθ​\(a∣s,g=s\+\)\]\\displaystyle\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}s\\sim d^\{\\beta\}\(s\)\\\\ a\\sim\\beta\(a\\mid s\)\\\\ s\_\{\+\}\\sim d^\{\\beta\}\(s\_\{\+\}\)\\end\{subarray\}\}\\left\[\\frac\{Q^\{\\beta\}\(s,a,s\_\{\+\}\)\}\{d^\{\\beta\}\(s\_\{\+\}\)\}\\log\\pi\_\{\\theta\}\(a\\mid s,g=s\_\{\+\}\)\\right\]\(18\)=11−γ​𝔼s∼dβ​\(s\)a∼β⁡\(a∣s\)s\+∼dβ​\(s\+\)​\[mβ​\(s\+∣s,a\)dβ​\(s\+\)​log⁡πθ​\(a∣s,g=s\+\)\]\\displaystyle=\\frac\{1\}\{1\-\\gamma\}\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}s\\sim d^\{\\beta\}\(s\)\\\\ a\\sim\\beta\(a\\mid s\)\\\\ s\_\{\+\}\\sim d^\{\\beta\}\(s\_\{\+\}\)\\end\{subarray\}\}\\left\[\\frac\{m^\{\\beta\}\(s\_\{\+\}\\mid s,a\)\}\{d^\{\\beta\}\(s\_\{\+\}\)\}\\log\\pi\_\{\\theta\}\(a\\mid s,g=s\_\{\+\}\)\\right\]\(19\)=11−γ​𝔼s∼dβ​\(s\)a∼β⁡\(a∣s\)s\+∼Mβ​\(s\+∣s,a\)​\[log⁡πθ​\(a∣s,g=s\+\)\]\.\\displaystyle=\\frac\{1\}\{1\-\\gamma\}\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}s\\sim d^\{\\beta\}\(s\)\\\\ a\\sim\\beta\(a\\mid s\)\\\\ s\_\{\+\}\\sim M^\{\\beta\}\(s\_\{\+\}\\mid s,a\)\\end\{subarray\}\}\\left\[\\log\\pi\_\{\\theta\}\(a\\mid s,g=s\_\{\+\}\)\\right\]\.\(20\)The second equality uses the identityQβ​\(s,a,s\+\)=\(1−γ\)−1​mβ​\(s\+∣s,a\)Q^\{\\beta\}\(s,a,s\_\{\+\}\)=\(\{1\-\\gamma\}\)^\{\-1\}m^\{\\beta\}\(s\_\{\+\}\\mid s,a\)in the goal\-conditioned setting, and the final change corresponds to importance sampling\. ∎

## Appendix DSamples

![Refer to caption](https://arxiv.org/html/2610.19129v1/pointmaze_terminal_train_eval_grid.png)Figure 6:Example ground\-truth training and evaluation trajectories onpointmaze\.![Refer to caption](https://arxiv.org/html/2610.19129v1/pointmaze_large_collected_local098_history4_cell3_6_left_seed0_directed_empirical.png)\(\(a\)\)
![Refer to caption](https://arxiv.org/html/2610.19129v1/pointmaze_large_collected_local098_history4_cell3_6_up_seed0_directed_empirical.png)\(\(b\)\)

Figure 7:Generated samples inpointmaze\-large with varying horizon\.Samples are conditioned on different sources displayed with, and ground\-truth action, whose effect can be inferred by the ground\-truth trajectory shown inblack\. Across each row, we vary discountγ\\gamma, withH=\(1−γ\)−1H=\(1\-\\gamma\)^\{\-1\}\.![Refer to caption](https://arxiv.org/html/2610.19129v1/pointmaze_medium_terminal_reload_sample_grid.png)\(\(a\)\)pointmaze\-medium
![Refer to caption](https://arxiv.org/html/2610.19129v1/pointmaze_large_terminal_reload_sample_grid.png)\(\(b\)\)pointmaze\-large

Figure 8:Generated samples inpointmazeenvironments\.Samples are conditioned on the source displayed with, and ground\-truth action, whose effect can be inferred by the ground\-truth trajectory shown inblack\. These samples reflect the settings used in Table[1](https://arxiv.org/html/2610.19129#S6.T1), whereMCHC\-​β\\text\{MCHC\-\}\\betausesγ=0\.998\\gamma=0\.998whileTD2​\-​β\\text\{TD\}^\{2\}\\text\{\-\}\\betaand MC useγ=0\.99\\gamma=0\.99\.![Refer to caption](https://arxiv.org/html/2610.19129v1/pointmaze_teleport_terminal_reload_sample_grid.png)\(\(a\)\)pointmaze\-teleport
![Refer to caption](https://arxiv.org/html/2610.19129v1/pointmaze_giant_terminal_reload_sample_grid.png)\(\(b\)\)pointmaze\-giant

Figure 9:Generated samples inpointmazeenvironments, continued\.![Refer to caption](https://arxiv.org/html/2610.19129v1/cube_triple_mchc_ghm_seed0_step1000000_empirical_mc0_full_history_sample_columns_generated_overlay.png)Figure 10:Samples oncube\-triple\-stitch\. The first row shows initial states, and the second shows overlaid empirical futures from the same trajectory\. The third row shows overlaid generated samples, while the following rows show individual generated samples\.Table 8:TD\-Augwith Contrastive RL\. We utilizeTD\-Augfor improving value learning with CRL\. We calculate success rate averaged over the last 3 checkpoints, and then report the mean and standard deviation over 3 seeds\.
## Appendix EExtension to Offline RL with Contrastive RL

In this section, we apply our data augmentation strategy to Contrastive RL \(CRL\)\([Eysenbach et al\., 2022](https://arxiv.org/html/2610.19129#bib.bib35)\)\. In particular, instead of augmenting goals for training the policy, we use augmentation to learn better goal\-conditioned \(on\-policy\) value functions\. CRL in OGBench[Park et al\. \(2025a\)](https://arxiv.org/html/2610.19129#bib.bib36)uses the binary NCE objective to train a criticf⁡\(s,a,g\)f\(s,a,g\):

ℒCRL=−𝔼\(s,a\)∼𝒟β,s−∼𝒟βs\+∼Pgoal\(⋅\|s,a\)\[logσ\(f\(s,a,s\+\)−log\(1−σ\(f\(s,a,s−\)\)\]\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{CRL\}\}=\-\\mathbb\{E\}\_\{\\begin\{subarray\}\{c\}\(s,a\)\\sim\\mathcal\{D\}^\{\\beta\},s\_\{\-\}\\sim\\mathcal\{D\}^\{\\beta\}\\\\ s\_\{\+\}\\sim P\_\{\\mathrm\{goal\}\(\\cdot\|s,a\)\}\\end\{subarray\}\}\\left\[\\log\\sigma\(f\(s,a,s\_\{\+\}\)\-\\log\(1\-\\sigma\(f\(s,a,s\_\{\-\}\)\)\\right\]\(21\)Here,s\+=s→k,k∼geom⁡\(1−γ\)s\_\{\+\}=s\_\{\\rightarrow k\},k\\sim\\mathrm\{geom\}\(1\-\\gamma\), an MC sample of future state,σ\\sigmais the sigmoid function, and the critic is parameterized byf⁡\(s,a,g\)=ϕ​\(s,a\)T​ψ​\(g\)/df\(s,a,g\)=\\phi\(s,a\)^\{T\}\\psi\(g\)/\\sqrt\{d\}\. With full rollouts from the mixture policy, the optimal critic estimates a log density ratio with the successor measuref∗​\(s,a,g\)=log⁡\(mγβ​\(g\|s,a\)/p⁡\(g\)\)f^\{\*\}\(s,a,g\)=\\log\(m^\{\\beta\}\_\{\\gamma\}\(g\|s,a\)/p\(g\)\), and thus,Qγβ​\(s,a,g\)∝exp⁡\(f∗​\(s,a,g\)\)Q^\{\\beta\}\_\{\\gamma\}\(s,a,g\)\\propto\\exp\(f^\{\*\}\(s,a,g\)\)for a fixedgg\. However, in practice, standard CRL is subject to the same compositional generalization gap as BC, where we only have truncated trajectories from individual policies, not the mixture policy\. Thus, if we can approximately sample goals from the successor measureMγβ\(⋅\|s,a\)M^\{\\beta\}\_\{\\gamma\}\(\\cdot\|s,a\)withTD\-Aug, we can improve the performance of value learning with CRL by replacingPgoal\(⋅\|s,a\)P\_\{\\mathrm\{goal\}\}\(\\cdot\|s,a\), withPaug\(⋅\|s,a\)P\_\{\\mathrm\{aug\}\}\(\\cdot\|s,a\)as in equation[12](https://arxiv.org/html/2610.19129#S5.E12)\. Withα=0\\alpha=0, we maintain the original CRL objective, and with0<α≤10<\\alpha\\leq 1, we mix in generated goals with proportionα\\alpha\.

Isolating the contribution ofTD\-Augsolely tovalue learning, we train the policy on goals sampled without augmentation for DDPG\+BC\([Fujimoto and Gu, 2021](https://arxiv.org/html/2610.19129#bib.bib2)\):

argmaxθ𝔼\(s,a\)∼𝒟β,g∼Pmixed\(⋅\|s,a\)\[λlogπθ\(a\|s,g\)\+f\(s,πθμ\(s,g\),g\)\]\\displaystyle\\arg\\max\_\{\\theta\}\\mathbb\{E\}\_\{\(s,a\)\\sim\\mathcal\{D\}^\{\\beta\},g\\sim P\_\{\\mathrm\{mixed\}\}\(\\cdot\|s,a\)\}\\left\[\\lambda\\log\\pi\_\{\\theta\}\(a\|s,g\)\+f\(s,\\pi^\{\\mu\}\_\{\\theta\}\(s,g\),g\)\\right\]\(22\)WherePmixed\(⋅\|s,a\)P\_\{\\mathrm\{mixed\}\}\(\\cdot\|s,a\)denotes a proportion of0\.50\.5for random states in the dataset, and0\.50\.5for future MC samples, thus only using ground\-truth samples\. In Table[8](https://arxiv.org/html/2610.19129#A4.T8), we utilize the standard OGBench\([Park et al\., 2025a](https://arxiv.org/html/2610.19129#bib.bib36)\)configuration for CRL, withα=0\.25\\alpha=0\.25for the augmentation probability when training the contrastive critic\. We can see that augmenting goals withTD\-Augleads to improved overall performance \(4→364\\rightarrow 36\)\. We also apply our visual setup as in Section[6\.2](https://arxiv.org/html/2610.19129#S6.SS2)also utilizingα=0\.25\\alpha=0\.25, except we reconstruct generated frames with the decoder so that we can perform CRL without modification\. Here, we see smaller gains \(OPEN56→61\)56\\rightarrow 61\)\. Further work could study if synthetic samples or other data augmentation may allow for improved scaling of CRL in the offline setting, which has been successful previously in the online setting\([Wang et al\., 2026](https://arxiv.org/html/2610.19129#bib.bib8)\)\.

Table 9:TD\-Augwith Contrastive RL in visual environments\.