Imagining Goals with Geometric Horizon Models for Policy Generalization

1 INTRODUCTION

• TD-Aug trains a geometric horizon model on offline data to generate outcomes after mixing behaviors, then uses generated goals to train policies for compositional tasks.The method targets zero-shot evaluation on harder tasks composed from knowledge learned on a subset of tasks.Source: §1 · p. 1, Figure 1 · p. 2
• The proposed TD-Flow objective combines Monte Carlo and temporal-difference sampling using horizon consistency, improving long-horizon generation.The paper reports that TD methods are crucial for generalization.Source: §1 · p. 1
• TD-Aug raises behavior-cloning average success from 31% to 79% on stitching-based OGBench navigation tasks, from 29% to 56% on image-based tasks, and from 13% to 52% on manipulation tasks.The reported evaluations cover low-dimensional navigation, image-based tasks, and complicated manipulation tasks.Source: §1 · p. 2

2 RELATED WORK

• Unlike hindsight goals, which do not cover compositional generalization, clustering-based stitching covers only 1-step joins and relies on heuristics such as state-distance metrics.TD-Aug instead uses generative models to capture a future-state distribution that stitches across trajectories.Source: §2 · p. 2
• Long-horizon prediction methods can chain jointly sampled sub-trajectories, whereas TD-Aug amortizes generation to produce long-horizon samples in O(1) steps without intermediate denoising or heuristic chaining.The compared chaining approaches are cited as Janner et al. (2022), Luo et al. (2025), and Lee & Choi (2025).Source: §2 · p. 2
• Representation learning can encode temporal dependencies to support compositional generalization without state prediction, but stability and long-horizon generalization remain challenges.Generative modeling and representation learning can also be used together for high-dimensional state prediction.Source: §2 · p. 2
• Figure 2 illustrates generated reachable future states that may require stitching several trajectories and extend beyond the horizon of short training trajectories.The example is conditioned on an initial state and action in pointmaze-giant-stitch.Source: Figure 2 · p. 3
• Geometric horizon models approximate samples from a policy’s successor measure; this work studies them for goal augmentation and zero-shot goal generalization.Prior applications include decision-time planning, policy improvement, and model-based value expansion.Source: §2 · p. 3

3.1 GOAL-CONDITIONED REINFORCEMENT LEARNING

• Goal-conditioned offline learning trains a policy to reach state-valued goals using a dataset, with a goal-reaching reward defined by a Dirac delta.The successor measure describes discounted future-state probabilities under a policy, and its density relates to the goal-conditioned value function.Source: §3.1 · p. 3

3.2 FLOW MATCHING

• Conditional flow matching uses a straight-line interpolation between source and target samples, with the conditional velocity target given by their difference.The source distribution is standard Gaussian.Source: §3.2 · p. 4

3.3 TRAINING POLICIES

• Standard goal-conditioned BC samples goals from future states in each trajectory, but datasets collected by multiple policies do not generally yield the mixture policy’s conditional successor measure.The training data can be represented by a Markovian mixture policy whose occupancy equals the expectation under the individual policies.Source: §3.3 · p. 4

3.4 TEMPORAL DIFFERENCE FLOWS

• TD-Flows learn geometric-horizon successor measures through flow matching, using TD bootstrapping to model future states beyond observed transitions.LTD2-β bootstraps with the next dataset action; by contrast, Monte Carlo fitting approximates a mixture of individual-policy successor measures when trajectories are truncated and separately collected.Source: §3.4 · pp. 4–5

4 IMPROVING LONG-HORIZON, ON-POLICY GENERATION BY COMBINING MC AND TD SAMPLING

• TD-HC uses a local-horizon GHM conditioned on ω < γ to estimate a higher-discount successor measure through a Bellman-like relation.The section adapts this approach to the on-policy setting.Source: §4 · p. 5
• The proposed horizon-consistent objective combines Monte Carlo samples with TD bootstrapping to improve long-horizon generation while retaining higher proportions of ground-truth data as γ increases.It bootstraps from a future Monte Carlo sample rather than only the next state.Source: §4 · pp. 5–6
• The objective’s Monte Carlo samples are truncated and collected by a single policy, so they do not correspond to the true mixture-policy successor measure as ω increases.Source: §4 · p. 6

5.1 GENERALIZATION PROBLEM

• Compositional generalization evaluates success on out-of-distribution tasks that may be unreachable by any one data-collecting policy but reachable by stitching their trajectories.The stated motivation is that mixture-policy successor-measure goals could teach a policy to reach these tasks.Source: §5.1 · p. 6

5.2 METHOD

• TD-Aug approximates and samples the mixture policy’s successor measure with on-policy TD-Flows, then trains on original goals mixed with generated goals.TD learning is used because full mixture-policy rollouts are unavailable; the default model objective is LMCHC-β, with LTD2-β also supported as TD-Aug2.Source: §5.2 · p. 6
• TD-Aug is not limited to a fixed number of stitches or heuristic sampling, and uses O(1) sampling steps per sample for expected horizon H = (1 − γ)^−1.The section also states that with α = 1 and exact mixture-policy successor-measure samples, the augmentation strategy is equivalent to value-weighted regression for one-step improvement over behavior policy β.Source: §5.2 · p. 6

6.1 STATE-BASED ENVIRONMENTS

• The state-based evaluation uses short trajectories of up to four maze cells and tests longer tasks spanning the maze.The stitch datasets were recollected to terminate at sub-goals because the original datasets contained transitions idling at goals.Source: §6.1 · p. 7
• For pointmaze-large-stitch, MCHC-β has the lowest EMD as the evaluated horizon increases.The comparison measures the distance between generated samples and empirical held-out samples at corresponding horizons.Source: §6.1 · p. 7, Figure 3 · p. 7
• On state-based stitch tasks, TD-Aug improves long-horizon generalization, whereas MC-Aug remains close to the original GCFBC performance.The comparison includes augmentation methods and reports success rates averaged over the last three checkpoints and three seeds.Source: §6.1 · p. 8, Table 1 · p. 8

6.2 COMPOSITION IN VISUAL NAVIGATION

• Visual experiments extend the GHM and policy to an autoencoder latent space rather than modeling images directly.Figure 4 depicts the visual-environment extension.Source: §6.2 · p. 8, Figure 4 · p. 8
• In visual navigation, TD-Aug improves latent-space AE-GCFBC and outperforms the standard GCBC baseline.The authors report that combining MC and TD is most effective; they attribute LTD2-β’s weaker performance to bootstrapping difficulty in approximate latent space.Source: Table 2 · p. 8, §6.2 · p. 9

6.3 COMPOSITION IN MANIPULATION

• On manipulation stitch tasks, TD-Aug raises GCFBC average success from 13 to 52.The evaluation uses action chunking and oracle goal representations.Source: §6.3 · p. 9, Table 3 · p. 9

7 CONCLUSION

• TD-based geometric horizon models enhance behavior-cloning generalization through goal augmentation.The conclusion presents this as the paper’s demonstrated result.Source: §7 · p. 9
• The conclusion points to real-world robotics as a setting where behavior cloning has scaled as a pre-training algorithm.It cites O’Neill et al. (2024), Black et al. (2024), and Intelligence et al. (2025).Source: §7 · p. 9
• Efficient and stable TD-based GHM training for high-dimensional, partially observable robotics tasks remains an open question.Source: §7 · p. 9

A TRAINING DETAILS

• The appendix provides separate hyperparameter tables for TD-Flow, GCFBC, visual autoencoder, and clustering-based goal augmentation.The autoencoder table specifies a 50-dimensional latent representation.Source: Table 4 · p. 14, Table 5 · p. 15, Table 6 · p. 15

B.1 TD-HC

• TDHC samples discount values during training, while MCHC fixes γ at γmax and uses geometrically sampled Monte Carlo future states.With γ = ω, the original TD2 objective is recovered.Source: §B.1 · p. 16

B.2 GHMS FOR MANIPULATION

• For manipulation, GHMs predict future states in action-chunk-sized jumps and reduce γ and ω proportionally to preserve the expected horizon.With chunk size h = 10, γ is reduced from 0.99 to 0.9; bootstrapping uses a learned prior to predict the next action chunk.Source: §B.2 · pp. 16–17

B.3 ONE-STEP DYNAMICS MODEL

• The baseline samples k ∼ geom(1 − γ), predicts the first next state conditioned on state and action, then predicts k − 1 states under the behavior mixture.Source: §B.3 · p. 17
• The model uses masked actions, replacing the action with a shared learned embedding with probability 0.5, and is trained for 1 million steps.Source: §B.3 · p. 17
• Autoregressive prediction compounds error because memoryless pβ(·|s) loses information about behavior intent after the first action.Source: §B.3 · p. 17
• Rollouts use history length h = 4 and cap k at 2H; at γ = 0.99, the cap is 200.Source: §B.3 · p. 17
• Because one-step models cannot efficiently resample fresh long-horizon goals, Dyn-Aug caches 200 future states per dataset state and geometrically samples from each window.Source: §B.3 · p. 17

C RELATION TO 1-STEP POLICY IMPROVEMENT

• The relation to one-step policy improvement uses Qβ(s, a, s+) = (1 − γ)−1mβ(s+ | s, a) and then applies importance sampling.Source: C RELATION TO 1-STEP POLICY IMPROVEMENT · p. 18

• Figure 6 presents example ground-truth training and evaluation trajectories on pointmaze.Source: Figure 6 · p. 19
• Figure 7 varies the geometric horizon across pointmaze-large generated samples, with ground-truth trajectories shown in black.Source: Figure 7 · p. 20
• Figure 8 shows pointmaze generated samples under the Table 1 settings: MCHC-β uses γ = 0.998, while TD2-β and MC use γ = 0.99.Source: Figure 8 · pp. 21–22
• Figure 9 contrasts initial states, empirical futures, overlaid generated samples, and individual generated samples for cube-triple-stitch.Source: Figure 9 · p. 23
• Table 8 reports success rates averaged over the last 3 checkpoints, with mean and standard deviation across 3 seeds.Source: Table 8 · p. 24

• TD-Aug is applied to CRL by augmenting goals for learning goal-conditioned on-policy value functions rather than policy training.Source: E EXTENSION TO OFFLINE RL WITH… · p. 24
• With full mixture-policy rollouts, the optimal critic estimates a log density ratio involving the successor measure.Source: E EXTENSION TO OFFLINE RL WITH… · p. 24
• Standard CRL faces the compositional generalization gap because offline data contains truncated trajectories from individual policies, not mixture-policy rollouts.Source: E EXTENSION TO OFFLINE RL WITH… · p. 24
• Generated goals replace the original goal distribution in the critic objective, with α = 0 retaining CRL and 0 < α ≤ 1 mixing in generated goals.Source: E EXTENSION TO OFFLINE RL WITH… · p. 24
• For DDPG+BC, policy goals are sampled without augmentation to isolate TD-Aug's contribution to value learning.Source: E EXTENSION TO OFFLINE RL WITH… · p. 24
• TD-Aug goal augmentation improves overall CRL performance from 4 to 36, with smaller gains from 56 to 61 in the visual setup.Source: E EXTENSION TO OFFLINE RL WITH… · p. 24, Table 9 · p. 24