Paper deep dive
Imagining Goals with Geometric Horizon Models for Policy Generalization
Daniel Lawson, Dhruv Sreenivas, Adriana Hugessen, Khimya Khetarpal, Glen Berseth
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 92%
Last extracted: 10/15/2026, 5:05:11 AM
Summary
The paper introduces TD-Aug, a method for improving zero-shot goal generalization in behavior cloning policies by using Geometric Horizon Models (GHMs) trained with Temporal Difference Flows (TD-Flows). TD-Aug generates synthetic future states (goals) by sampling from a geometric horizon model, allowing the policy to learn compositional tasks (stitching) that were not present in the static offline dataset. The authors propose a new objective, MCHC-beta, which combines Monte Carlo and Temporal Difference sampling to improve long-horizon generation quality.
Entities (7)
Relation Signals (5)
TD-Aug → uses → Geometric Horizon Model
confidence 95% · we propose TD-Aug, which samples from a geometric horizon model
TD-Aug → evaluatedon → OGBench
confidence 93% · We demonstrate that training on these future outcomes ... significantly improves generalization in stitching-based OGBench tasks.
Geometric Horizon Model → trainedwith → Temporal Difference Flows
confidence 92% · generative model we consider is a geometric horizon model (GHM) ... trained using Temporal Difference Flows (TD-Flows)
TD-Aug → improves → Zero-shot goal generalization
confidence 90% · training on these future outcomes as goals for goal-conditioned BC policies significantly improves generalization
MCHC-beta → ispartof → TD-Aug
confidence 85% · Train M_phi_beta ... using L_MCHC-beta
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2610.19129v1
- Canonical: https://arxiv.org/abs/2610.19129v1
Trouble viewing inline? Open PDF directly →
Full Text
77,718 characters extracted from source content.
Expand or collapse full text
Imagining Goals with Geometric Horizon Models for Policy Generalization Daniel Lawson Affiliation: Mila Dhruv Sreenivas Affiliation: Mila Adriana Hugessen Affiliation: Mila Khimya Khetarpal Affiliation: Mila Affiliation: Université de Montréal,Google DeepMind, †Equal Advising Glen Berseth Affiliation: 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 †footnotetext: Correspondence: daniel.lawson@mila.quebec 1 Introduction 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) present an interesting strategy for improving policies from offline data (Hafner et al., 2025). These models can be used to generate synthetic experience (Sutton, 1991; Bruce et al., 2024), 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). 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, 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; Thakoor et al., 2022) trained using Temporal Difference Flows (TD-Flows) (Farebrother et al., 2025). 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, and refer to our method as TD-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; Farebrother et al., 2026) and geometric Monte Carlo (MC) (Eysenbach et al., 2022) sampling, demonstrating that TD methods are crucial for generalization. Using horizon consistency (Farebrother et al., 2026), 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, we first validate TD-Aug on low-dimensional navigation tasks in OGBench (Park et al., 2025a) that require stitching for success, substantially improving behavior cloning generalization, from 31%31\% to 79%79\% average success. Then, we demonstrate that our approach can also improve generalization on high-dimensional, image-based tasks (29%29\% to 56%56\%) in Section 6.2 and complicated manipulation tasks (13%13\% to 52%52\%) in Section 6.3. Figure 1: Overview of TD-Aug. We a) 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 1 TD-Aug Representation Learning (Optional) Use either states s, or learned representations z=ϕ(s)z=φ(s), as the space for GHM and policy (Section 6.2). Geometric Horizon Model Training Train MϕβM^β_φ to approximate the successor measure of the behavior mixture policy, using ℒMCHC-βL_MCHC-β (Equation 9). Policy Training Train πθ _θ by BC on a mixture of generated goals from MϕβM^β_φ with proportion α, as well as real samples (Equation 11). 2 Related Work Data augmentation with 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; Ha and Schmidhuber, 2018; Feinberg et al., 2018; Hafner et al., 2020; Lu et al., 2023; Zhou et al., 2024). In the goal-conditioned setting, Hindsight Experience Replay (HER) (Andrychowicz et al., 2017) allows re-using transitions with different achieved outcomes as goals (Kaelbling, 1993). When training policies with BC, sampling hindsight goals is beneficial, but does not cover compositional generalization (Ghugare et al., 2024), 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), which could be chained together (Luo et al., 2025; Lee and Choi, 2025), 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 in O(1)O(1) steps through amortization during GHM training. Representation learning can be employed to improve compositional generalization of 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; Lawson et al., 2025; Park et al., 2026b). 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γ -models (Janner et al., 2020), propose a generative variant of temporal difference learning, allowing for approximate sampling from the successor measure (Thakoor et al., 2022; Tomar et al., 2024; Wiltzer et al., 2024; Farebrother et al., 2025), a continuous generalization of successor representations (Dayan, 1993) of a policy. Specifically, we consider Temporal Difference Flows (Farebrother et al., 2025) (TD-Flows), which use flow-matching objectives for learning such models. Most recent work with TD-Flows (Farebrother et al., 2025; Farebrother et al., 2026) has considered using GHMs for decision-time planning. GHMs can also be used for policy improvement (Zheng et al., 2026), or model-based value expansion (Feinberg et al., 2018), where future synthetic states can be used to estimate bootstrap targets for value estimation (Janner et al., 2020). Recent work has used MVE for offline RL using GHMs and Universal Horizon Models (UHMs) (Chung et al., 2026), 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. Figure 2: Ground-truth and generated samples. Left, on pointmaze-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, 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. 3 Background 3.1 Goal-Conditioned Reinforcement Learning We study offline learning in goal-conditioned Markov Decision Processes (MDPs), defined as the tuple (,,,P,r,γ,μ)(S,A,G,P,r,γ,μ), where S is the state space, A is the action space, G is the goal space, P:×→Δ()P:S×A→ (S) is the transition kernel, r is the reward function, γ∈[0,1)γ∈[0,1) is the discount, and μ∈Δ()μ∈ (S) is an initial state distribution. We consider that goals coincide with states (=G=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)= _g(s_t+1). We aim to learn a policy π:×→Δ()π:S×G→ (A) using a dataset 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)] J(π)=E_g p(g),s_0 μ,a_t π(· s_t,g),s_t+1 P(· s_t,a_t) [ _t=0^∞γ^tr_g(s_t,a_t,s_t+1) ] For the subset X⊆X , we define the normalized successor measure of a policy π as: Mγπ(X∣s,a)≜(1−γ)∑k≥0γkPr(sk+1∈X∣s0=s,a0=a,π) M^π_γ(X s,a) (1-γ) _k≥ 0γ^kPr(s_k+1∈ X s_0=s,a_0=a,π) (1) where PrPr denotes the time-specific probability that we are in X starting from (s,a)(s,a) when following π. We write the density of the successor measure mγπm^π_γ via Mγπ(ds+∣s,a)=mγπ(s+∣s,a)ds+ M^π_γ(ds_+ s,a)=m^π_γ(s_+ s,a)\,ds_+, and the action-marginalized successor measure is defined as Mγπ(⋅∣s)≜a∼π(⋅∣s)[Mγπ(⋅∣s,a)] M^π_γ(· s) E_a π(· s)[M^π_γ(· s,a)]. We also define the discounted state occupancy density as the marginal of mγπ m^π_γ over the initial state distribution: dπ(s)≜s0∼μ,a0∼π(⋅∣s0)[mγπ(s∣s0,a0)] d^π(s) E_s_0 μ,a_0 π(· s_0)[m^π_γ(s s_0,a_0)]. With the goal reaching reward, the successor measure’s density relates to the goal-conditioned value function via Qγπ(s,a,g)=(1−γ)−1mγπ(g∣s,a)Q^π_γ(s,a,g)=(1-γ)^-1m^π_γ(g s,a). 3.2 Flow Matching Following success in modeling policies and geometric horizon models (Farebrother et al., 2025), we utilize flow matching (Lipman et al., 2023). With a source distribution p0∈Δ(ℝd)p_0∈ (R^d) and a target distribution p1∈Δ(ℝd)p_1∈ (R^d), flow matching learns a velocity field vτ:ℝd→ℝdv^τ:R^d ^d such that for x0∼p0x_0 p_0, the sample x1=x0+∫01vτ(xτ)τx_1=x_0+ _0^1v^τ(x_τ)dτ matches the target distribution (x1∼p1x_1 p_1). In practice, we use conditional flow matching (Lipman et al., 2023): ℒCFM(θ)=x1∼,τ∼U(0,1),xτ∼pτ|1(⋅∣x1)[∥vθτ(xτ)−uτ(xτ∣x1)∥22]. _ CFM(θ)=E_x_1 ,τ U(0,1),x_τ p_τ 1(· x_1) [ v^τ_θ(x_τ)-u_τ(x_τ x_1) _2^2 ]. We use the conditional optimal transport (CondOT) coupling, drawing a straight line between x0x_0 and x1x_1, xτ=(1−τ)x0+τx1x_τ=(1-τ)x_0+τ x_1. This gives the true conditional velocity field target as uτ(xτ∣x1)=x1−x0u_τ(x_τ x_1)=x_1-x_0. Our source distribution p0p_0 is the standard Gaussian, (0,I)N(0,I). 3.3 Training Policies We train policies on datasets =(s0i,a0i,⋯,sTi,aTi)i=1ND=\(s_0^i,a_0^i,·s,s_T^i,a_T^i)\_i=1^N comprised of trajectories collected by a set of unknown policies βj(a|s)\ _j(a|s)\. These policies can be marginalized by a mixture policy β(a|s)=∑jβj(a|s)p(βj|s)β(a~|~s)= _j _j(a~|~s)p( _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)] d^β(s)=E_p( _j) [d _j(s) ] (Ghugare et al., 2024). 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^β(·|s,a) _ _j p( _j|s,a) [M _j(·|s,a) ]. We consider the standard goal-conditioned BC objective, which samples a current state from the dataset (OPENdβ)d^β), 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 policy MβjM _j: ℒBC(θ)=−s∼dβ,a∼β(⋅|s)βj∼p(βj|s,a),s+∼Mβj(⋅|s,a)[logπθ(a|s,g=s+)] _BC(θ)=-E_ subarraycs d^β,a β(·~|~s)\\ _j p( _j|s,a),~s_+ M _j(·|s,a) subarray [ _θ(a|s,g=s_+) ] (2) We perform goal-conditioned behavior cloning using flow matching, representing the policy by a conditional velocity field vθτ(aτ∣s,g)v^τ_θ(a_τ s,g): ℒBC(θ)=(s,a,s+),τ∼U(0,1),a0∼p0,aτ=(1−τ)a0+τa[‖vθτ(aτ∣s,g=s+)−(a−a0)‖22] _BC(θ)=E_(s,a,s_+),τ U(0,1),a_0 p_0,a_τ=(1-τ)a_0+τ a [\|v^τ_θ(a_τ s,g=s_+)-(a-a_0)\|_2^2 ] (3) To sample from πθ(⋅|s,g) _θ(·|s,g), using a0∼p0a_0 p_0, we integrate with vθτv_θ^τ, via a1=a0+∫01vθτ(aτ|s,g)τa_1=a_0+ _0^1v^τ_θ(a_τ|s,g)dτ. 3.4 Temporal 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; Thakoor et al., 2022). These models can be trained via TD learning, as for all (s,a)(s,a), the successor measure MγπM^π_γ satisfies: Mγπ(⋅|s,a)=(1−γ)P(⋅|s,a)+γs′∼P(⋅|s,a),a′∼π(⋅|s′)[Mγπ(⋅|s′,a′)] M^π_γ(·|s,a)=(1-γ)P(·|s,a)+ _s P(·|s,a),a π(·|s )[M^π_γ(·|s ,a )] (4) Optimizing this objective is challenging with prior generative modeling techniques such as GANs (Goodfellow et al., 2014) or normalizing flows (Rezende and Mohamed, 2015), especially at large discounts γ. This led to the proposal to use flow matching to learn GHMs (Farebrother et al., 2025), which showed strong success in simulated control domains. TD-Flows parameterize the GHM as a velocity field vϕτ:ℝd×→ℝd v^τ_φ:R^d×S×A ^d with parameters ϕφ, such that with source x0∼p0∈Δ(ℝd)x_0 p_0∈ (R^d), the velocity field induces the flow ψτ:ℝd×→ℝdψ^τ:R^d×S×A ^d: ψϕτ(x0|s,a)=x0+∫0τvϕi(ψϕi(x0|s,a)|s,a)i. ψ^τ_φ(x_0|s,a)=x_0+ _0^τv^i_φ(ψ^i_φ(x_0|s,a)|s,a)di. (5) We aim to learn vϕτv^τ_φ, such that for x0∼p0x_0 p_0, the samples ψϕ1(x0|s,a)ψ^1_φ(x_0|s,a) are distributed according to Mγπ(⋅|s,a) M^π_γ(·|s,a). Thus, we will use the notation Mϕπ(⋅|s,a) M^π_φ(·|s,a) for our approximate successor measure parameterized by vϕτ v^τ_φ. Exploiting the Bellman structure of the successor measure, Farebrother et al. (2025) show that a flow that models said measure can be learned by minimizing the TD2-CFMTD^2-CFM loss: ℒTD2(ϕ;ϕ¯)=(1−γ)ℒtd-s(ϕ)+γℒtd-bs(ϕ;ϕ¯), where: _TD^2(φ; φ)=(1-γ)L_ td-s(φ)+ _ td-bs(φ; φ), where: (6) ℒtd-s(ϕ)=τ∼U(0,1),(s,a,s′)∼β,x0∼p0,xτ=(1−τ)x0+τs′[‖vϕτ(xτ|s,a)−(s′−x0)‖2] _ td-s(φ)=E_τ U(0,1),(s,a,s ) ^β,x_0 p_0,x_τ=(1-τ)x_0+τ s [\|v^τ_φ(x_τ|s,a)-(s -x_0)\|^2 ] ℒ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] _ td-bs(φ)=E_τ U(0,1),(s,a,s ) ^β,a π(·|s ),x_0 p_0,x_τ=ψ^τ_ φ(x_0|s ,a ) [\|v^τ_φ(x_τ|s,a)-v^τ_ φ(x_τ|s ,a )\|^2 ] where ϕ¯ φ are target flow parameters, as commonly used in value-based RL (Mnih et al., 2015). Recalling equation 4, ℒtd-sL_ td-s targets the single-step transition kernel, while ℒtd-bsL_ td-bs bootstraps using the successor measure at the next state. Together, this leads to approximating Mγπ M^π_γ. We can also approximate the successor measure Mβ(⋅|s,a) M^β(·|s,a) of the mixture policy β, rather than an arbitrary policy π. While this requires sampling a′∼β(⋅|s′)a β(·|s ) 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]L_ td-bs(φ)=E_τ U(0,1),(s,a,s ,a ) ^β,x_0 p_0,x_τ=ψ^τ_ φ(x_0|s ,a ) [\|v^τ_φ(x_τ|s,a)-v^τ_ φ(x_τ|s ,a )\|^2 ] We refer to this objective as ℒTD2-βL_TD^2-β. This is a simpler objective than the fully off-policy objective in equation 6, as we bootstrap over fully in-support actions. If one does not need to stitch across trajectories, we can utilize a simpler approach to model MγβM^β_γ 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]. _ MC(φ)=E_ subarrayc(s,a) ^β,k (1-γ),s_+=s_→ k\\ τ U(0,1),x_0 p_0,x_τ=(1-τ)x_0+τ s_+ subarray [ v^τ_φ(x_τ s,a)-(s_+-x_0) _2^2 ]. (7) Here, s→ks_→ k denotes the state that is k steps further from s along the sampled trajectory in the dataset. If we had infinitely long trajectories, collected directly by the mixture policy, this objective would approximate Mβ(⋅|s,a)M^β(·|s,a). However, in practice, we have datasets with truncated horizon, and collected from a set of data-collecting policies. This leads ℒMC(ϕ)L_ MC(φ) to instead approximate a mixture of successor measures βj∼p(βj|s,a)Mγβj(⋅|s,a) E_ _j p( _j|s,a)M _j_γ(·|s,a) (Lawson et al., 2025), preferring the use of ℒTD2-βL_TD^2-β. 4 Improving 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-βL_TD^2-β) to address scaling to long horizon prediction, which is affected by the discount γ used 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), which proposes how to use a "local horizon" GHM (e.g. a GHM conditioned on a discount ω<γ [rgb]1,0.5,0ω<γ) to obtain an estimate of the successor measure to approximate a higher discount GHM (γ [rgb]0,0,1γ) (Janner et al., 2020; Farebrother et al., 2026), i.e. ω<γ [rgb]1,0.5,0ω< [rgb]0,0,1γ through a Bellman-like relation: Mγπ(⋅|s,a)≜ M^π_ [rgb]0,0,1γ(·|s,a) (1−γ)P(⋅|s,a)+γ1−γ1−ωs′∼P(⋅|s,a),a′∼π(⋅|s′)[Mωπ(⋅|s′,a′)] (1- [rgb]0,0,1γ)P(·|s,a)+ [rgb]0,0,1γ 1- [rgb]0,0,1γ1- [rgb]1,0.5,0ωE_s P(·|s,a),a π(·|s )[M^π_ [rgb]1,0.5,0ω(·|s ,a )] + + γ−ω1−ωs′∼P(⋅|s,a),a′∼π(⋅|s′),s+∼Mπω(⋅|s′,a′),a+∼π(⋅|s+)[Mγπ(⋅|s+,a+)] [rgb]0,0,1γ [rgb]0,0,1γ- [rgb]1,0.5,0ω1- [rgb]1,0.5,0ωE_s P(·|s,a),a π(·|s ),s_+ M^π_ [rgb]1,0.5,0ω(·|s ,a ),a_+ π(·|s_+)[M^π_ [rgb]0,0,1γ(·|s_+,a_+)] (8) This can be converted to a TD-Flow loss, ℒTDHCL_TDHC, by modifying equation 6, shown in Appendix B.1. We propose a variant of ℒTDHCL_TDHC adapted for the on-policy setting to approximate MγβM^β_ [rgb]0,0,1γ. While equation 7 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 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 a horizon-consistent objective utilizing MC samples from the dataset. While we will learn MγβM^β_ [rgb]0,0,1γ, we will not learn MωβM^β_ [rgb]1,0.5,0ω, and instead utilize MC samples from dataset trajectories. This leads to the objective: ℒMCHC-β(ϕ;ϕ¯)=(1−γ)ℒtd-s(ϕ)+γ1−γ1−ωℒmc(ϕ)+γ−ω1−ωℒtd-mc-bs(ϕ;ϕ¯), where: _MCHC-β(φ; φ)=(1- [rgb]0,0,1γ)L_ td-s(φ)+ [rgb]0,0,1γ 1- [rgb]0,0,1γ1- [rgb]1,0.5,0ωL_ mc(φ)+ [rgb]0,0,1γ [rgb]0,0,1γ- [rgb]1,0.5,0ω1- [rgb]1,0.5,0ωL_ td-mc-bs(φ; φ), where: (9) ℒtd-s(ϕ)=τ∼U(0,1),(s,a,s′)∼β,x0∼p0,xτ=(1−τ)x0+τs′[‖vϕτ(xτ|s,a)−(s′−x0)‖2] _ td-s(φ)=E_τ U(0,1),(s,a,s ) ^β,x_0 p_0,x_τ=(1-τ)x_0+τ s [\|v^τ_φ(x_τ|s,a)-(s -x_0)\|^2 ] ℒ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] _ mc(φ)=E_ subarraycτ U(0,1),(s,a) ^β,k (1- [rgb]1,0.5,0ω)\\ s_+=s_→ k+1,x_0 p_0,x_τ=(1-τ)x_0+τ s_+ subarray [\|v^τ_φ(x_τ|s,a)-(s_+-x_0)\|^2 ] ℒ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] _ td-mc-bs(φ)=E_ subarraycτ U(0,1),(s,a) ^β,k (1- [rgb]1,0.5,0ω)\\ s_+,a_+=s_→ k+1,a_→ k+1,x_0 p_0,x_τ=ψ^τ_ φ(x_0|s_+,a_+) subarray [\|v^τ_φ(x_τ|s,a)-v^τ_ φ(x_τ|s_+,a_+)\|^2 ]\ This loss is not just supervised by the next ground truth states, but also future MC samples. Additionally, instead of bootstrapping from the next state st+1s_t+1, we will bootstrap from this MC sample, which is obtained as a future state after st+1s_t+1, i.e. s+=st+k+1s_+=s_t+k+1, where k∼geom(1−ω)k (1- [rgb]1,0.5,0ω). As we utilize MC samples, which are both truncated, and are collected by a single policy, they do not correspond to the true MωβM^β_ [rgb]1,0.5,0ω. Thus, this objective is a biased approximation of MγβM^β_ [rgb]0,0,1γ as ω [rgb]1,0.5,0ω increases, but allows increasing γ [rgb]0,0,1γ while still training on higher proportions of ground-truth data. 5 Imagining Goals for Policy Generalization 5.1 Generalization Problem We consider a generalization problem for BC policies. Following the setting from Section 3.3, we train on trajectories collected by unknown policies βj(a|s)\ _j(a|s)\, with mixture policy β(a|s)β(a|s). We will now define training and test distributions for πθ _θ equation 2, following Ghugare et al. (2024): 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 R^π_id= E_ subarraycs d^β, _j p( _j|s),s_g M _j(·|s) subarray [R _θ(s,s_g) ]_ in-distribution training tasks ,R^π_gen= E_ subarraycs d^β,s_g M^β(·|s) subarray [R _θ(s,s_g) ]_tasks requiring compositional generalization (10) With Rθπ(s,sg)R^π_θ(s,s_g) defining the success rate of the policy to reach sgs_g from s, RidπR^π_id averages performance over in-distribution tasks and RgenπR^π_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 _j, but still have sufficient data coverage to complete through composing or stitching trajectories collected by different policies, i.e. reachable with β. While equation 2 corresponds to improving on RidπR^π_id, the BC objective does not cover RgenπR^π_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 β during training, rather than that of individual policies, our policy could learn to reach goals required for better RgenπR^π_gen. 5.2 Method 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-βL_TD^2-β or ℒMCHC-βL_MCHC-β, rather than ℒMC(ϕ)L_ MC(φ). Thus, we train on-policy TD-Flows, which will yield an approximation Mϕβ(⋅|s,a) M^β_φ(·|s,a) for the successor measure of the mixture policy. After obtaining Mϕβ M^β_φ, we perform goal augmentation where 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+)] _ _θE_ subarraycs d^β,a β(a|s),~s_+ P_aug(·|s,a) subarray [ _θ(a|s,g=s_+) ] (11) wherePaug(⋅|s,a)≜αMβϕ(⋅|s,a)⏟imagined compositional samples+(1−α)p(βj|s,a)Mβj(⋅|s,a)⏟ground truth future states ~P_aug(·|s,a) αM^β_φ(·|s,a)_imagined compositional samples+ (1-α)E_p( _j|s,a)M _j(·|s,a)_ground truth future states (12) We generally refer our goal augmentation procedure as TD-Aug, and by default use ℒMCHC-βL_MCHC-β. We can also instantiate our method with ℒTD2-βL_TD^2-β, which we denote later with TD-Aug2TD-Aug_2. We provide an overview of our method in Algorithm 1 and Figure 1. 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). We also perform O(1)O(1) sampling steps per sample to cover an expected horizon H=(1−γ)−1H=(1-γ)^-1, while one-step models need O(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; Wang et al., 2018; Peng et al., 2019; Wang et al., 2020), for a brief proof, see Appendix C: Theorem 1 (Equivalence to 1-step improvement). With α=1α=1, and exact samples from Mβ(⋅|s,a)M^β(·|s,a), our augmentation strategy in equation 11 is equivalent to value-weighted regression used for one-step improvement over behavior policy β: θ∗∈argmaxθs∼dβ,s+∼dβa∼β(⋅∣s)[Qβ(s,a,g=s+)dβ(s+)logπθ(a∣s,g=s+)] θ^*∈ _θ\;E_ subarraycs d^β,s_+ d^β\\ a β(· s) subarray [ Q^β(s,a,g=s_+)d^β(s_+)\, _θ(a s,g=s_+) ] (13) 6 Experiments 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 for pointmaze-large-stitch. In this section, we show that TD-Aug can improve combinatorial generalization across a variety of stitch tasks in the OGBench (Park et al., 2025a) benchmark suite. In Section 6.1, we (a) evaluate the quality of generated samples as horizon scales for different GHM objectives on low-dimensional navigation tasks, and (b) show that TD-Aug is 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, and Section 6.3, we demonstrate that we can also improve stitching on manipulation tasks. 6.1 State-based Environments We first evaluate our method on pointmaze-stitch navigation 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 first compare the quality of generative models. Qualitatively, we provide visualizations on pointmaze-medium,large,teleport,giant-stitch in Figure 8. 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 in teleport, 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-βTD^2-β,MCHC-β), can generate samples that require stitching, capturing states reachable beyond than the training trajectory horizon. However, the TD2-βTD^2-β samples are more diffuse, or may go backwards from the direction implied by the ground truth samples. In contrast, the MCHC-βMCHC-β samples balance implicit stitching with the quality of MC samples. Next, we perform a quantitative comparison between model samples on a held-out evaluation dataset on pointmaze-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. 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 horizons H=(1−γ)−1H=(1-γ)^-1, measuring against ground truth samples at the corresponding horizon, as shown in Figure 3. In addition to the GHM based models, we also generate samples from autoregressive one-step prediction, which we further discuss in Appendix B.3. While some models have similar EMD at the initial horizon (100), we can see that the errors of TD2-βTD^2-β, 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-βMCHC-β has 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. We bold the top performing methods. pointmaze GCFBC +Cluster-Aug +Dyn-Aug +MC-Aug +TD-Aug2TD-Aug_2 +TD-Aug GCIQL QRL CRL HIQL medium-stitch 55± 1055\, ±\,10 89± 589\, ±\,5 94± 694\, ±\,6 70± 270\, ±\,2 96± 196\, ±\,1 ± 099\, ±\,0 74± 174\, ±\,1 88± 588\, ±\,5 4± 24\, ±\,2 71± 871\, ±\,8 large-stitch 26± 626\, ±\,6 62± 862\, ±\,8 51± 1051\, ±\,10 28± 1728\, ±\,17 40± 340\, ±\,3 ± 397\, ±\,3 40± 140\, ±\,1 89± 1589\, ±\,15 0± 00\, ±\,0 18± 418\, ±\,4 giant-stitch 0± 00\, ±\,0 3± 63\, ±\,6 0± 00\, ±\,0 0± 00\, ±\,0 0± 00\, ±\,0 ± 1754\, ±\,17 0± 00\, ±\,0 54± 954\, ±\,9 0± 00\, ±\,0 0± 00\, ±\,0 teleport-stitch 44± 744\, ±\,7 47± 247\, ±\,2 40± 740\, ±\,7 48± 348\, ±\,3 45± 445\, ±\,4 ± 366\, ±\,3 27± 127\, ±\,1 2± 42\, ±\,4 10± 510\, ±\,5 30± 330\, ±\,3 average 31 50 46 37 45 79 35 58 4 30 Figure 4: Extension of TD-Aug to visual environments. Next, we study how geometric horizon models can empirically improve a policy’s combinatorial generalization capabilities in Table 1. Without any augmentation, a goal-conditioned flow-based behavior cloning (GCFBC) baseline (Park et al., 2025b) 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 to TD-Aug using ℒMCHC-βL_MCHC-β, we compare to TD-Aug2TD-Aug_2 using ℒTD2-βL_TD^2-β, as well as MC-Aug using ℒMCL_MC. We compare to clustering-based goal augmentation (Cluster-Aug) (Ghugare et al., 2024), 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. For these methods, we list additional training details in Appendix A, including hyperparameters for BC in Table 5, and for TD-Flows in Table 4. Although the task horizon varies, we use the same hyperparameters for each environment in pointmaze for TD-Aug. This includes the effective horizon, which we set to 500500 (γ=0.998γ=0.998), while for MC-Aug and TD-Aug2TD-Aug_2 we use 100100 (γ=0.99γ=0.99), as they perform worse at higher horizons. After training GHMs, all policies are trained for 10610^6 steps, and each method is trained with batches using goals that are of a proportion 1−α1-α of future states in the same trajectory (i.e. hindsight goal sampling), and with proportion α with generated goals. The BC baseline uses α=0α=0, and augmentation uses α=0.5α=0.5. We can see that TD-Aug significantly improves generalization, including on long-horizon tasks, such as pointmaze-giant-stitch. In contrast, MC-Aug maintains close to original performance, demonstrating that generating additional synthetic short-horizon samples does not improve generalization. We also see that with TD-Aug, GCFBC outperforms prior offline GCRL methods in pointmaze, including GCIQL (Kostrikov et al., 2022), QRL (Wang et al., 2023), Contrastive RL (CRL) (Eysenbach et al., 2022), and HIQL (Park et al., 2023). Observing the weaker performance of CRL on stitch tasks, which also uses MC samples like GCBC, we show contrastive learning on generated samples leads to improved performance in Appendix E. 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. We bold the highest performing method overall, and highlight in brown the top BC method. visual-antmaze AE-GCFBC +MC-Aug +TD-Aug2TD-Aug_2 +TD-Aug GCBC GCIQL QRL CRL HIQL medium-stitch 45± 445\, ±\,4 60± 260\, ±\,2 56± 156\, ±\,1 84± 7 [rgb]0.75,0.5,0.2584\, ±\,7 69± 469\, ±\,4 9± 09\, ±\,0 0± 00\, ±\,0 82± 282\, ±\,2 ± 192\, ±\,1 large-stitch 12± 112\, ±\,1 19± 119\, ±\,1 9± 09\, ±\,0 28± 1 [rgb]0.75,0.5,0.2528\, ±\,1 21± 621\, ±\,6 1± 01\, ±\,0 1± 01\, ±\,0 29± 629\, ±\,6 ± 246\, ±\,2 average 29 40 33 56 45 5 1 56 69 Figure 5: Samples for visual-antmaze-medium. We visualize samples generated with MCHC-βMCHC-β at γ=0.995γ=0.995. We display both real, geometrically sampled, states, as well as reconstructions from the GHM’s predicted latents using the AE decoder. 6.2 Composition in Visual Navigation In this section, we extend TD-Aug to visual environments with visual-antmaze-stitch. Instead of learning the GHM in pixel space, we first train an autoencoder (AE), obtaining a low-dimensional latent space Z. We provide details for training the AE in Table 6. Then, we utilize the same architectures for the GHM, and policy, as performed in the non-visual setting (Section 6.1), but now operate over latent space, as shown in Figure 4. As a baseline, we report both standard BC in OGBench (GCBC), using early state-goal fusion, and an autoencoder BC baseline, labeled AE-GCFBC. The standard baseline concatenates state and goal images, and passes them through a small IMPALA-based (Espeholt et al., 2018) CNN. In contrast, the autoencoder BC baseline separately encodes goals and states into 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) and feed them to the policy without decoding. With this setup, in Table 2, we report adding augmentation to AE-GCFBC. We can see that without augmentation, this architecture under-performs GCBC. While we improve with MC-Aug, we can see that ℒTD2-βL_TD^2-β, used in TD-Aug2TD-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 (Z) space. Thus, we find the combination of MC and TD to be the most effective, where TD-Aug improves 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 the oracle goal representation. We use the original play dataset, and then introduce a new stitch dataset, consisting of only short trajectories. Then, we add TD-Aug to close the generalization gap from stitch. We calculate average success on the last 3 checkpoints, and report the mean, std over 3 seeds. We highlight the top method on stitch. dataset GCFBC +TD-Aug2TD-Aug_2 +TD-Aug GCIQL HIQL play stitch stitch stitch stitch stitch cube-double 75± 175\, ±\,1 20± 020\, ±\,0 78± 278\, ±\,2 84± 0 [rgb]0.75,0.5,0.2584\, ±\,0 63± 163\, ±\,1 6± 26\, ±\,2 cube-triple 35± 535\, ±\,5 20± 020\, ±\,0 36± 5 [rgb]0.75,0.5,0.2536\, ±\,5 31± 231\, ±\,2 8± 28\, ±\,2 0± 00\, ±\,0 puzzle-4x4 23± 223\, ±\,2 0± 00\, ±\,0 18± 018\, ±\,0 40± 2 [rgb]0.75,0.5,0.2540\, ±\,2 16± 516\, ±\,5 0± 00\, ±\,0 average 44 13 44 52 29 2 6.3 Composition 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 consider cube-double,triple-play, puzzle-4x4-play tasks with oracle goal representation (Park et al., 2025a; Park et al., 2026a). In addition to the original play datasets, we construct a new challenging stitch dataset, 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)π(a_t:t+h|s,g), obtains good average success (44%)(44\%) on the play tasks, it struggles when training on the stitch (13%)(13\%). Thus, we apply TD-Aug to close the gap. We also train GHMs, with action chunking, Mϕβ(⋅|s,at:t+h) M^β_φ(·|s,a_t:t+h), and predict h-sizedh-sized jumps in the full state space, with additional details in Appendix B.2. On stitch, with TD-Aug, GCFBC’s performance substantially increases from 13→5213→ 52. We provide visualizations of the GHM’s samples in Figure 10. 7 Conclusion 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; Black et al., 2024; Intelligence et al., 2025). 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 Andrychowicz et al. (2017) M. Andrychowicz, F. Wolski, A. Ray, J. Schneider, R. Fong, P. Welinder, B. McGrew, J. Tobin, P. Abbeel, and W. Zaremba Hindsight experience replay. Advances in neural information processing systems 30. Cited by: §2. Black et 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 _0: A vision-language-action flow model for general robot control. External Links: 2410.24164, Link Cited by: §7. Bruce et 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. In Forty-first International Conference on Machine Learning, Cited by: §1. Chung et al. (2026) H. Chung, J. Lee, and S. Oh Offline reinforcement learning with universal horizon models. External Links: 2605.15603, Link Cited by: §2. Dao et al. (2023) Q. Dao, H. Phung, B. Nguyen, and A. Tran Flow matching in latent space. arXiv preprint arXiv:2307.08698. Cited by: §6.2. Dayan (1993) P. Dayan Improving generalization for temporal difference learning: the successor representation. Neural computation 5 (4), p. 613–624. Cited by: §2. Espeholt et 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. In International conference on machine learning, p. 1407–1416. Cited by: §6.2. Eysenbach et al. (2022) B. Eysenbach, T. Zhang, S. Levine, and R. R. Salakhutdinov Contrastive learning as goal-conditioned reinforcement learning. Advances in Neural Information Processing Systems 35, p. 35603–35620. Cited by: Appendix E, §1, §6.1. Farebrother et al. (2026) J. Farebrother, M. Pirotta, A. Tirinzoni, M. G. Bellemare, A. Lazaric, and A. Touati Compositional planning with jumpy world models. External Links: 2602.19634, Link Cited by: §B.1, §1, §2, §4. Farebrother et al. (2025) J. Farebrother, M. Pirotta, A. Tirinzoni, R. Munos, A. Lazaric, and A. Touati Temporal difference flows. In Forty-second International Conference on Machine Learning, External Links: Link Cited by: Table 4, §1, §1, §2, §3.2, §3.4, §3.4. Feinberg et al. (2018) V. Feinberg, A. Wan, I. Stoica, M. I. Jordan, J. E. Gonzalez, and S. Levine Model-based value estimation for efficient model-free reinforcement learning. arXiv preprint arXiv:1803.00101. Cited by: §2, §2. Fujimoto and Gu (2021) S. Fujimoto and S. S. Gu A minimalist approach to offline reinforcement learning. Advances in neural information processing systems 34, p. 20132–20145. Cited by: Appendix E. Ghugare et al. (2024) R. Ghugare, M. Geist, G. Berseth, and B. Eysenbach Closing the gap between TD learning and supervised learning - a generalisation point of view.. In The Twelfth International Conference on Learning Representations, External Links: Link Cited by: Table 7, §1, §2, §3.3, §5.1, §5.2, §6.1. Goodfellow et al. (2014) I. J. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio Generative adversarial nets. In Advances in Neural Information Processing Systems, Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K.Q. Weinberger (Eds.), Vol. 27, p. . External Links: Link Cited by: §3.4. Ha and Schmidhuber (2018) D. Ha and J. Schmidhuber Recurrent world models facilitate policy evolution. Advances in neural information processing systems 31. Cited by: §1, §2. Hafner et al. (2020) D. Hafner, T. Lillicrap, J. Ba, and M. Norouzi Dream to control: learning behaviors by latent imagination. External Links: 1912.01603, Link Cited by: §2. Hafner et al. (2025) D. Hafner, W. Yan, and T. Lillicrap Training agents inside of scalable world models. arXiv preprint arXiv:2509.24527. Cited by: §1. Intelligence et 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∗π^*_0.6: A vla that learns from experience. External Links: 2511.14759, Link Cited by: §7. Janner et al. (2022) M. Janner, Y. Du, J. B. Tenenbaum, and S. Levine Planning with diffusion for flexible behavior synthesis. External Links: 2205.09991, Link Cited by: §2. Janner et al. (2020) M. Janner, I. Mordatch, and S. Levine γ-Models: generative temporal difference learning for infinite-horizon prediction. In Advances in Neural Information Processing Systems, Cited by: §1, §2, §3.4, §4. Kaelbling (1993) L. P. Kaelbling Learning to achieve goals. In International Joint Conference on Artificial Intelligence, Cited by: §2. Kostrikov et al. (2022) I. Kostrikov, A. Nair, and S. Levine Offline reinforcement learning with implicit q-learning. In International Conference on Learning Representations, External Links: Link Cited by: §6.1. Kostrikov et al. (2020) I. Kostrikov, D. Yarats, and R. Fergus Image augmentation is all you need: regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004.13649. Cited by: Table 6. Lawson et al. (2025) D. Lawson, A. Hugessen, C. Cloutier, G. Berseth, and K. Khetarpal Self-predictive representations for combinatorial generalization in behavioral cloning. External Links: 2506.10137, Link Cited by: §2, §3.4. Lee and Choi (2025) K. Lee and J. Choi State-covering trajectory stitching for diffusion planners. External Links: 2506.00895, Link Cited by: §2. Lipman et al. (2023) Y. Lipman, R. T. Q. Chen, H. Ben-Hamu, M. Nickel, and M. Le Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, External Links: Link Cited by: §3.2, §3.2. Lloyd (1982) S. Lloyd Least squares quantization in pcm. IEEE transactions on information theory 28 (2), p. 129–137. Cited by: Table 7. Lu et al. (2023) C. Lu, P. J. Ball, Y. W. Teh, and J. Parker-Holder Synthetic experience replay. In Thirty-seventh Conference on Neural Information Processing Systems, External Links: Link Cited by: §2. Luo et al. (2025) Y. Luo, U. A. Mishra, Y. Du, and D. Xu Generative trajectory stitching through diffusion composition. External Links: 2503.05153, Link Cited by: §2. Mnih et 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. Hassabis Human-level control through deep reinforcement learning. Nature 518 (7540), p. 529–533. External Links: ISSN 00280836, Link Cited by: §3.4. Myers et al. (2025) V. Myers, B. C. Zheng, A. Dragan, K. Fang, and S. Levine Temporal representation alignment: successor features enable emergent compositionality in robot instruction following. External Links: 2502.05454, Link Cited by: §2. O’Neill et 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. In 2024 IEEE International Conference on Robotics and Automation (ICRA), p. 6892–6903. Cited by: §7. Park et al. (2025a) S. Park, K. Frans, B. Eysenbach, and S. Levine OGBench: benchmarking offline goal-conditioned rl. In International Conference on Learning Representations (ICLR), External Links: Link Cited by: Appendix E, Appendix E, §1, §6.3, §6. Park et al. (2026a) S. Park, K. Frans, D. Mann, B. Eysenbach, A. Kumar, and S. Levine Horizon reduction makes rl scalable. Advances in Neural Information Processing Systems 38, p. 8350–8389. Cited by: §6.3. Park et al. (2023) S. Park, D. Ghosh, B. Eysenbach, and S. Levine Hiql: offline goal-conditioned rl with latent states as actions. Advances in Neural Information Processing Systems 36, p. 34866–34891. Cited by: §6.1. Park et al. (2025b) S. Park, Q. Li, and S. Levine Flow q-learning. In Forty-second International Conference on Machine Learning, External Links: Link Cited by: §6.1. Park et al. (2026b) S. Park, D. Mann, and S. Levine Dual goal representations. External Links: 2510.06714, Link Cited by: §2. Peng et al. (2019) X. B. Peng, A. Kumar, G. Zhang, and S. Levine Advantage-weighted regression: simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00177. Cited by: §5.2. Peters and Schaal (2007) J. Peters and S. Schaal Reinforcement learning by reward-weighted regression for operational space control. In Proceedings of the 24th international conference on Machine learning, p. 745–750. Cited by: §5.2. Rezende and Mohamed (2015) D. Rezende and S. Mohamed Variational inference with normalizing flows. In Proceedings of the 32nd International Conference on Machine Learning, F. Bach and D. Blei (Eds.), Proceedings of Machine Learning Research, Vol. 37, Lille, France, p. 1530–1538. External Links: Link Cited by: §3.4. Sutton (1991) R. S. Sutton Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin 2 (4), p. 160–163. Cited by: §1, §2. Thakoor et al. (2022) S. Thakoor, M. Rowland, D. Borsa, W. Dabney, R. Munos, and A. Barreto Generalised policy improvement with geometric policy composition. In Proceedings 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, p. 21272–21307. External Links: Link Cited by: §1, §2, §3.4. Tomar et al. (2024) M. Tomar, P. Hansen-Estruch, P. Bachman, A. Lamb, J. Langford, M. E. Taylor, and S. Levine Video occupancy models. External Links: 2407.09533, Link Cited by: §2. Wang et al. (2026) K. Wang, I. Javali, M. Bortkiewicz, T. Trzcinski, and B. Eysenbach 1000 layer networks for self-supervised rl: scaling depth can enable new goal-reaching capabilities. Advances in Neural Information Processing Systems 38, p. 157643–157670. Cited by: Appendix E. Wang et 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 Systems 31. Cited by: §5.2. Wang et al. (2023) T. Wang, A. Torralba, P. Isola, and A. Zhang Optimal goal-reaching reinforcement learning via quasimetric learning. In International Conference on Machine Learning, External Links: Link Cited by: §6.1. Wang et 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 Systems 33, p. 7768–7778. Cited by: §5.2. Wiltzer et al. (2024) H. Wiltzer, J. Farebrother, A. Gretton, Y. Tang, A. Barreto, W. Dabney, M. G. Bellemare, and M. Rowland A distributional analogue to the successor representation. External Links: 2402.08530, Link Cited by: §2. Yarats et al. (2021) D. Yarats, A. Zhang, I. Kostrikov, B. Amos, J. Pineau, and R. Fergus Improving sample efficiency in model-free reinforcement learning from images. In Proceedings of the aaai conference on artificial intelligence, Vol. 35, p. 10674–10681. Cited by: Table 6. Zheng et al. (2026) C. Zheng, S. Park, S. Levine, and B. Eysenbach Intention-conditioned flow occupancy models. External Links: 2506.08902, Link Cited by: §2. Zhou et al. (2024) Z. Zhou, C. Zhu, R. Zhou, Q. Cui, A. Gupta, and S. S. Du Free from bellman completeness: trajectory stitching via model-based return-conditioned supervised learning. In The Twelfth International Conference on Learning Representations, External Links: Link Cited by: §2. Appendix A Training Details Hyperparameter Flow Matching ODE Solver Midpoint Train ODE dτdτ 0.10.1 Sample ODE dτdτ 0.10.1 Network (U-Net) t-Positional Embedding Dim 256256 t-Positional Embedding MLP (256,256)(256,256) Block Dimension (512,512,512)(512,512,512) Encoder MLP (512,512,512)(512,512,512) activation Mish Optimizer (AdamW) AdamW β1 _1 0.90.9 AdamW β2 _2 0.9990.999 AdamW ϵε 10−410^-4 Learning Rate 10−410^-4 Weight Decay 10−310^-3 Training Target Network EMA 10−310^-3 Gradient Steps TD State,Visual 3M Manipulation 1M MC, MCHC 1M Batch Size 1024 Discount (γ) MC, TD 0.990.99 MCHC State 0.9980.998 Visual 0.9950.995 Manipulation 0.990.99 Discount (ω) State,Visual 0.990.99 Manipulation 0.980.98 Table 4: Hyperparameters for TD-Flow training, following similar architecture details to Farebrother et al. (2025). Table 5: Hyperparameters for GCFBC training. Hyperparameter Flow Matching ODE Solver Euler Sample ODE dτdτ 0.10.1 Network (MLP) Hidden Layers (1024,1024,1024,1024)(1024,1024,1024,1024) activation GELU Optimizer (Adam) Adam β1 _1 0.90.9 Adam β2 _2 0.9990.999 Learning Rate State, Visual 3×10−43× 10^-4 Manipulation 3×10−53× 10^-5 Weight Decay 00 Training Gradient Steps State 1M Manipulation 2M Visual 500k Batch Size State 1024 Manipulation 1024 Visual 256 Augmentation (only) Batch Proportion (α) 0.50.5 Table 6: Hyperparameters for visual autoencoder training. For our autoencoder, we use a similar architecture to (Kostrikov et al., 2020; Yarats et al., 2021). 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). Hyperparameter Architecture Convolution Layers 22 Channels 3232 Kernel (3×3)(3× 3) Activation ReLU Latent Dimension, |z||z| 5050 Optimizer (Adam) Adam β1 _1 0.90.9 Adam β2 _2 0.9990.999 Learning Rate 10−410^-4 Weight Decay 00 Training Gradient Steps 100k Batch Size 1024 Table 7: Hyperparameters for clustering-based goal-augmentation. We follow (Ghugare et al., 2024) and use k-means clustering (Lloyd, 1982), with k set as below. Hyperparameter Clusters medium 40 large 40 teleport 40 giant 80 Appendix B Additional Objectives B.1 TD-HC For reference, we provide the original Temporal Difference Horizon Consistency (TDHC) objective (Farebrother et al., 2026), which we modify to use MC samples in Section 4: ℒTDHC(ϕ,ϕ¯)=(1−γ)ℒtd-s(ϕ,γ)+γ1−γ1−ωℒtd-bs(ϕ,ϕ¯,γ,ω)+γγ−ω1−ωℒtd-hc-bs(ϕ,ϕ¯,γ,ω) _TDHC(φ; φ)=(1- [rgb]0,0,1γ)L_ td-s(φ; [rgb]0,0,1γ)+ [rgb]0,0,1γ 1- [rgb]0,0,1γ1- [rgb]1,0.5,0ωL_ td-bs(φ; φ, [rgb]0,0,1γ, [rgb]1,0.5,0ω)+ [rgb]0,0,1γ [rgb]0,0,1γ- [rgb]1,0.5,0ω1- [rgb]1,0.5,0ωL_ td-hc-bs(φ; φ, [rgb]0,0,1γ, [rgb]1,0.5,0ω) (14) ℒtd-s=τ∼U(0,1),(s,a,s′)∼βx0∼p0,xτ=(1−τ)x0+τs′[‖vϕτ(xτ|s,a,γ)−(s′−x0)‖2] _ td-s=E_ subarraycτ U(0,1),(s,a,s ) ^β\\ x_0 p_0,x_τ=(1-τ)x_0+τ s subarray [\|v^τ_φ(x_τ|s,a, [rgb]0,0,1γ)-(s -x_0)\|^2 ] ℒ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] _ td-bs=E_ subarraycτ U(0,1),(s,a,s ) ^β,a π(s )\\ x_0 p_0,x_τ=ψ^τ_ φ(x_0|s ,a , [rgb]1,0.5,0ω) subarray [\|v^τ_φ(x_τ|s,a, [rgb]0,0,1γ)-v^τ_ φ(x_τ|s ,a , [rgb]1,0.5,0ω)\|^2 ] ℒ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] _ td-hc-bs=E_ subarraycτ U(0,1),(s,a,s ) ^β,a π(s )\\ x_0^1,x_0^2 p_0,s_+=ψ^1_ φ(x_0^1|s ,a , [rgb]1,0.5,0ω),a_+ π(s_+)\\ x_τ=ψ^τ_ φ(x_0^2|s_+,a_+, [rgb]0,0,1γ) subarray [\|v^τ_φ(x_τ|s,a, [rgb]0,0,1γ)-v^τ_ φ(x_τ|s_+,a_+, [rgb]0,0,1γ)\|^2 ]\ With γ=ω [rgb]0,0,1γ= [rgb]1,0.5,0ω, the original TD2TD^2 objective equation 6 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] [rgb]0,0,1γ∈[ _min, _max] and ω∈[γmin,γ] [rgb]1,0.5,0ω∈[ _min, [rgb]0,0,1γ]. Instead, the MCHC objective equation 9 we use fixes γ=γmax [rgb]0,0,1γ= _max, and samples (s+,a+)∼Dβ(s_+,a_+) D^β, where s+=st+ks_+=s_t+k, k∼geom(1−ω)k (1- [rgb]1,0.5,0ω), replacing having to learn and sample with vϕτ(⋅|⋅,ω)v_φ^τ(·|·, [rgb]1,0.5,0ω). B.2 GHMs for Manipulation For manipulation tasks, we modify both GCFBC and the GHM to utilize action chunks of size H as shown in equation 15. For the GHM we also predict future states starting from after the action chunk, at increments matching the chunk size. We also proportionally reduce γ,ωγ,ω. For example, with h=10h=10, we reduce γ from 0.990.99 to 0.90.9, to match the same expected horizon. Additionally, we bootstrap with samples from a learned prior, βθ(⋅|s+) _θ(·|s_+), predicting the next action chunk, rather than using a+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-s(ϕ)+γ1−γ1−ωℒmc(ϕ)+γγ−ω1−ωℒtd-mc-bs(ϕ,ϕ¯) _MCHC-β(φ; φ)=(1- [rgb]0,0,1γ)L_ td-s(φ)+ [rgb]0,0,1γ 1- [rgb]0,0,1γ1- [rgb]1,0.5,0ωL_ mc(φ)+ [rgb]0,0,1γ [rgb]0,0,1γ- [rgb]1,0.5,0ω1- [rgb]1,0.5,0ωL_ td-mc-bs(φ; φ) (15) ℒtd-s(ϕ)=τ∼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] _ td-s(φ)=E_ subarraycτ U(0,1),(s_t,a_t:t+h,s_t+h) ^β\\ x_0 p_0,x_τ=(1-τ)x_0+τ s_t+h subarray [\|v^τ_φ(x_τ|s,a_t:t+h)-(s_t+h-x_0)\|^2 ] ℒmc(ϕ)=τ∼U(0,1),(st,at:t+h,s+)∼βx0∼p0,xτ=(1−τ)x0+τs+[∥vϕτ(xτ|s,at:t+h)−(s+−x0)∥2] _ mc(φ)=E_ subarraycτ U(0,1),(s_t,a_t:t+h,s_+) ^β\\ x_0 p_0,x_τ=(1-τ)x_0+τ s_+ subarray [\|v^τ_φ(x_τ|s,a_t:t+h)-(s_+-x_0)\|^2 ] ℒ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] _ td-mc-bs(φ)=E_ subarraycτ U(0,1),(s_t,a_t:t+h,s_+) ^β\\ x_0 p_0,a_+ _θ(·|s_+),x_τ=ψ^τ_ φ(x_0|s_+,a_+) subarray [\|v^τ_φ(x_τ|s_t,a_t:t+h)-v^τ_ φ(x_τ|s_+,a_+)\|^2 ]\ where: s+=st+h(1+k),k∼geom(1−ω), and a+ also uses chunk size h=10. : s_+=s_t+h(1+k),k (1- [rgb]0,0,1ω), and a_+ also uses chunk size h=10. B.3 One-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 offset k∼geom(1−γ)k (1-γ). We then perform autoregressive prediction for k steps, first sampling st+1∼p(⋅|s,a)s_t+1 p(·|s,a), and then sampling k−1k-1 times from st+i+1∼pβ(⋅|st+i)s_t+i+1 p^β(·|s_t+i), for i=1i=1 to k−1k-1. We parameterize p(⋅|s,a)p(·|s,a) and pβ(⋅|s)p^β(·|s) via velocity field vϕτ(xτ|s,a)v^τ_φ(x_τ|s,a), trained via flow matching: ℒ1-step(ϕ) _ 1-step(φ) =τ∼(0,1),(s,a,s′)∼β,am∼mask(a)x0∼p0,xτ=(1−τ)x0+τs′[‖vϕτ(xτ∣s,am)−(s′−x0)‖22] =E_ subarraycτ (0,1),\,(s,a,s ) ^β,\,a_m (a)\\ x_0 p_0,\,x_τ=(1-τ)x_0+τ s subarray [ \|v^τ_φ(x_τ s,a_m)-(s -x_0) \|_2^2 ] (16) We utilize a standard 1-step flow matching objective, however we condition on am∼mask(a)a_m (a), where with probability 0.50.5, the action is replaced by a learned embedding (shared for all actions). Thus, this objective trains both forms of p. We utilize the same architecture as GHM as described in Table 6, 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. This is particularly due to pβ(⋅|s)p^β(·|s) being memory-less, i.e. after the first action, there is no conditioning information in pβ(⋅|s)p^β(·|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(ϕ) _history 1-step(φ) =τ∼(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] =E_ subarraycτ (0,1),\,(s_t-h+1:t,a,s_t+1) ^β\\ x_0 p_0,\,x_τ=(1-τ)x_0+τ s_t+1\\ a_m (a) subarray [ \|v^τ_φ(x_τ s_t-h+1:t,a_m)-(s_t+1-x_0) \|_2^2 ] (17) Where we utilize a history length of h=4h=4, as used in Figures 3, 7. When sampling with a specific discount γ, with expected horizon H=11−γH= 11-γ, we cap k for rollouts at 2H2H. For example, if γ=0.99γ=0.99, we cap k=min(geom(1−0.99),200)k= (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 M^β_γ(·|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γ=0.99. Thus, we instead cache a 2H=2002H=200 sized trajectory from each state from the dataset, storing (N,200)(N,200) states, where N is 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. Appendix C Relation 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+)] _ subarraycs d^β(s)\\ a β(a s)\\ s_+ d^β(s_+) subarray [ Q^β(s,a,s_+)d^β(s_+) _θ(a s,g=s_+) ] (18) =11−γs∼dβ(s)a∼β(a∣s)s+∼dβ(s+)[mβ(s+∣s,a)dβ(s+)logπθ(a∣s,g=s+)] = 11-γE_ subarraycs d^β(s)\\ a β(a s)\\ s_+ d^β(s_+) subarray [ m^β(s_+ s,a)d^β(s_+) _θ(a s,g=s_+) ] (19) =11−γs∼dβ(s)a∼β(a∣s)s+∼Mβ(s+∣s,a)[logπθ(a∣s,g=s+)]. = 11-γE_ subarraycs d^β(s)\\ a β(a s)\\ s_+ M^β(s_+ s,a) subarray [ _θ(a s,g=s_+) ]. (20) The second equality uses the identity Qβ(s,a,s+)=(1−γ)−1mβ(s+∣s,a)Q^β(s,a,s_+)=(1-γ)^-1m^β(s_+ s,a) in the goal-conditioned setting, and the final change corresponds to importance sampling. ∎ Appendix D Samples Figure 6: Example ground-truth training and evaluation trajectories on pointmaze. ((a)) ((b)) Figure 7: Generated samples in pointmaze-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 in black. Across each row, we vary discount γ, with H=(1−γ)−1H=(1-γ)^-1. ((a)) pointmaze-medium ((b)) pointmaze-large Figure 8: Generated samples in pointmaze environments. Samples are conditioned on the source displayed with , and ground-truth action, whose effect can be inferred by the ground-truth trajectory shown in black. These samples reflect the settings used in Table 1, where MCHC-βMCHC-β uses γ=0.998γ=0.998 while TD2-βTD^2-β and MC use γ=0.99γ=0.99. ((a)) pointmaze-teleport ((b)) pointmaze-giant Figure 9: Generated samples in pointmaze environments, continued. Figure 10: Samples on cube-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-Aug with Contrastive RL. We utilize TD-Aug for 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. pointmaze CRL +MC-Aug +TD-Aug2TD-Aug_2 +TD-Aug medium-stitch 4± 24\, ±\,2 2± 02\, ±\,0 24± 624\, ±\,6 ± 670\, ±\,6 large-stitch 0± 00\, ±\,0 0± 00\, ±\,0 4± 54\, ±\,5 ± 1746\, ±\,17 giant-stitch 0± 00\, ±\,0 0± 00\, ±\,0 0± 00\, ±\,0 ± 01\, ±\,0 teleport-stitch 10± 510\, ±\,5 6± 36\, ±\,3 6± 66\, ±\,6 ± 1026\, ±\,10 average 4 2 9 36 Appendix E Extension to Offline RL with Contrastive RL In this section, we apply our data augmentation strategy to Contrastive RL (CRL) (Eysenbach et al., 2022). 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) uses the binary NCE objective to train a critic f(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−))] _CRL=-E_ subarrayc(s,a) ^β,s_- ^β\\ s_+ P_goal(·|s,a) subarray [ σ(f(s,a,s_+)- (1-σ(f(s,a,s_-)) ] (21) Here, s+=s→k,k∼geom(1−γ)s_+=s_→ k,k (1-γ), an MC sample of future state, σ is the sigmoid function, and the critic is parameterized by f(s,a,g)=ϕ(s,a)Tψ(g)/df(s,a,g)=φ(s,a)^Tψ(g)/ d. With full rollouts from the mixture policy, the optimal critic estimates a log density ratio with the successor measure f∗(s,a,g)=log(mγβ(g|s,a)/p(g))f^*(s,a,g)= (m^β_γ(g|s,a)/p(g)), and thus, Qγβ(s,a,g)∝exp(f∗(s,a,g))Q^β_γ(s,a,g) (f^*(s,a,g)) for a fixed g. 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 measure Mγβ(⋅|s,a)M^β_γ(·|s,a) with TD-Aug, we can improve the performance of value learning with CRL by replacing Pgoal(⋅|s,a)P_goal(·|s,a), with Paug(⋅|s,a)P_aug(·|s,a) as in equation 12. With α=0α=0, we maintain the original CRL objective, and with 0<α≤10<α≤ 1, we mix in generated goals with proportion α. Isolating the contribution of TD-Aug solely to value learning, we train the policy on goals sampled without augmentation for DDPG+BC (Fujimoto and Gu, 2021): argmaxθ(s,a)∼β,g∼Pmixed(⋅|s,a)[λlogπθ(a|s,g)+f(s,πθμ(s,g),g)] _θE_(s,a) ^β,g P_mixed(·|s,a) [λ _θ(a|s,g)+f(s,π^μ_θ(s,g),g) ] (22) Where Pmixed(⋅|s,a)P_mixed(·|s,a) denotes a proportion of 0.50.5 for random states in the dataset, and 0.50.5 for future MC samples, thus only using ground-truth samples. In Table 8, we utilize the standard OGBench (Park et al., 2025a) configuration for CRL, with α=0.25α=0.25 for the augmentation probability when training the contrastive critic. We can see that augmenting goals with TD-Aug leads to improved overall performance (4→364→ 36). We also apply our visual setup as in Section 6.2 also utilizing α=0.25α=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→ 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). Table 9: TD-Aug with Contrastive RL in visual environments. visual-antmaze CRL +TD-Aug medium-stitch 82± 282\, ±\,2 ± 186\, ±\,1 large-stitch 29± 629\, ±\,6 ± 436\, ±\,4 average 56 61