Imagining Goals with Geometric Horizon Models for Policy Generalization

Daniel Lawson; Dhruv Sreenivas; Adriana Hugessen; Khimya Khetarpal; Glen Berseth

cs.LGOct 13, 2026

Imagining Goals with Geometric Horizon Models for Policy Generalization

Authors: Daniel Lawson, Dhruv Sreenivas, Adriana Hugessen, Khimya Khetarpal, Glen Berseth

Organizations: Mila · Université de Montréal · Google DeepMind

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.

Figures & tables

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.

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 .

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 .

pointmaze	GCFBC	+Cluster-Aug	+Dyn-Aug	+ MC-Aug	+ $\text{TD-Aug}_{2}$	+ TD-Aug	GCIQL	QRL	CRL	HIQL
medium-stitch	$55\,{\scriptstyle{\pm\,10}}$	$89\,{\scriptstyle{\pm\,5}}$	$94\,{\scriptstyle{\pm\,6}}$	$70\,{\scriptstyle{\pm\,2}}$	$96\,{\scriptstyle{\pm\,1}}$	$\mathbf{99}\,{\scriptstyle{\pm\,0}}$	$74\,{\scriptstyle{\pm\,1}}$	$88\,{\scriptstyle{\pm\,5}}$	$4\,{\scriptstyle{\pm\,2}}$	$71\,{\scriptstyle{\pm\,8}}$
large-stitch	$26\,{\scriptstyle{\pm\,6}}$	$62\,{\scriptstyle{\pm\,8}}$	$51\,{\scriptstyle{\pm\,10}}$	$28\,{\scriptstyle{\pm\,17}}$	$40\,{\scriptstyle{\pm\,3}}$	$\mathbf{97}\,{\scriptstyle{\pm\,3}}$	$40\,{\scriptstyle{\pm\,1}}$	$89\,{\scriptstyle{\pm\,15}}$	$0\,{\scriptstyle{\pm\,0}}$	$18\,{\scriptstyle{\pm\,4}}$
giant-stitch	$0\,{\scriptstyle{\pm\,0}}$	$3\,{\scriptstyle{\pm\,6}}$	$0\,{\scriptstyle{\pm\,0}}$	$0\,{\scriptstyle{\pm\,0}}$	$0\,{\scriptstyle{\pm\,0}}$	$\mathbf{54}\,{\scriptstyle{\pm\,17}}$	$0\,{\scriptstyle{\pm\,0}}$	$\textbf{54}\,{\scriptstyle{\pm\,9}}$	$0\,{\scriptstyle{\pm\,0}}$	$0\,{\scriptstyle{\pm\,0}}$
teleport-stitch	$44\,{\scriptstyle{\pm\,7}}$	$47\,{\scriptstyle{\pm\,2}}$	$40\,{\scriptstyle{\pm\,7}}$	$48\,{\scriptstyle{\pm\,3}}$	$45\,{\scriptstyle{\pm\,4}}$	$\mathbf{66}\,{\scriptstyle{\pm\,3}}$	$27\,{\scriptstyle{\pm\,1}}$	$2\,{\scriptstyle{\pm\,4}}$	$10\,{\scriptstyle{\pm\,5}}$	$30\,{\scriptstyle{\pm\,3}}$
average	31	50	46	37	45	79	35	58	4	30

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.

Figure 4: Extension of TD-Aug to visual environments.

visual-antmaze	AE-GCFBC	+ MC-Aug	+ $\text{TD-Aug}_{2}$	+ TD-Aug	GCBC	GCIQL	QRL	CRL	HIQL
medium-stitch	$45\,{\scriptstyle{\pm\,4}}$	$60\,{\scriptstyle{\pm\,2}}$	$56\,{\scriptstyle{\pm\,1}}$	${\color[rgb]{0.75,0.5,0.25}\textbf{84}}\,{\scriptstyle{\pm\,7}}$	$69\,{\scriptstyle{\pm\,4}}$	$9\,{\scriptstyle{\pm\,0}}$	$0\,{\scriptstyle{\pm\,0}}$	$82\,{\scriptstyle{\pm\,2}}$	$\mathbf{92}\,{\scriptstyle{\pm\,1}}$
large-stitch	$12\,{\scriptstyle{\pm\,1}}$	$19\,{\scriptstyle{\pm\,1}}$	$9\,{\scriptstyle{\pm\,0}}$	${\color[rgb]{0.75,0.5,0.25}\textbf{28}}\,{\scriptstyle{\pm\,1}}$	$21\,{\scriptstyle{\pm\,6}}$	$1\,{\scriptstyle{\pm\,0}}$	$1\,{\scriptstyle{\pm\,0}}$	$29\,{\scriptstyle{\pm\,6}}$	$\mathbf{46}\,{\scriptstyle{\pm\,2}}$
average	29	40	33	56	45	5	1	56	69

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.

Figure 5: Samples for visual-antmaze-medium . We visualize samples generated with $\text{MCHC-}\beta$ at $\gamma=0.995$ . We display both real, geometrically sampled, states, as well as reconstructions from the GHM’s predicted latents using the AE decoder.

dataset	GCFBC		+ $\text{TD-Aug}_{2}$	+ TD-Aug	GCIQL	HIQL
dataset	play	stitch	stitch	stitch	stitch	stitch
cube-double	$75\,{\scriptstyle{\pm\,1}}$	$20\,{\scriptstyle{\pm\,0}}$	$78\,{\scriptstyle{\pm\,2}}$	${\color[rgb]{0.75,0.5,0.25}\textbf{84}}\,{\scriptstyle{\pm\,0}}$	$63\,{\scriptstyle{\pm\,1}}$	$6\,{\scriptstyle{\pm\,2}}$
cube-triple	$35\,{\scriptstyle{\pm\,5}}$	$20\,{\scriptstyle{\pm\,0}}$	${\color[rgb]{0.75,0.5,0.25}\textbf{36}}\,{\scriptstyle{\pm\,5}}$	$31\,{\scriptstyle{\pm\,2}}$	$8\,{\scriptstyle{\pm\,2}}$	$0\,{\scriptstyle{\pm\,0}}$
puzzle-4x4	$23\,{\scriptstyle{\pm\,2}}$	$0\,{\scriptstyle{\pm\,0}}$	$18\,{\scriptstyle{\pm\,0}}$	${\color[rgb]{0.75,0.5,0.25}\textbf{40}}\,{\scriptstyle{\pm\,2}}$	$16\,{\scriptstyle{\pm\,5}}$	$0\,{\scriptstyle{\pm\,0}}$
average	44	13	44	52	29	2

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 .

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

	Hyperparameter
Flow Matching	ODE Solver	Midpoint
	Train ODE $\mathrm{d}\tau$	$0.1$
	Sample ODE $\mathrm{d}\tau$	$0.1$
Network (U-Net)	$t$ -Positional Embedding Dim	$256$
	$t$ -Positional Embedding MLP	$(256,256)$
	Block Dimension	$(512,512,512)$

Appendix

Table 4: Hyperparameters for TD-Flow training, following similar architecture details to Farebrother et al. (2025) .

	Hyperparameter
Flow Matching	ODE Solver	Euler
	Sample ODE $\mathrm{d}\tau$	$0.1$
Network (MLP)	Hidden Layers	$(1024,1024,1024,1024)$
	activation	GELU
Optimizer (Adam)	Adam $\beta_{1}$	$0.9$
	Adam $\beta_{2}$	$0.999$

Appendix

Table 5: Hyperparameters for GCFBC training.

	Hyperparameter
Architecture	Convolution Layers	$2$
	Channels	$32$
	Kernel	$(3\times 3)$
	Activation	ReLU
	Latent Dimension, $\|z\|$	$50$
Optimizer (Adam)	Adam $\beta_{1}$	$0.9$

Appendix

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

, 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
Clusters
medium	40
large	40
teleport	40
giant	80

Appendix

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.

Figure 6: Example ground-truth training and evaluation trajectories on pointmaze .

Panel of Figure 7 — 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 $\gamma$ , with $H=(1-\gamma)^{-1}$ .

Panel of Figure 8 — 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 $\text{MCHC-}\beta$ uses $\gamma=0.998$ while $\text{TD}^{2}\text{-}\beta$ and MC use $\gamma=0.99$ .

Panel of Figure 9 — 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.

pointmaze	CRL	+ MC-Aug	+ $\text{TD-Aug}_{2}$	+ TD-Aug
medium-stitch	$4\,{\scriptstyle{\pm\,2}}$	$2\,{\scriptstyle{\pm\,0}}$	$24\,{\scriptstyle{\pm\,6}}$	$\mathbf{70}\,{\scriptstyle{\pm\,6}}$
large-stitch	$0\,{\scriptstyle{\pm\,0}}$	$0\,{\scriptstyle{\pm\,0}}$	$4\,{\scriptstyle{\pm\,5}}$	$\mathbf{46}\,{\scriptstyle{\pm\,17}}$
giant-stitch	$0\,{\scriptstyle{\pm\,0}}$	$0\,{\scriptstyle{\pm\,0}}$	$0\,{\scriptstyle{\pm\,0}}$	$\mathbf{1}\,{\scriptstyle{\pm\,0}}$
teleport-stitch	$10\,{\scriptstyle{\pm\,5}}$	$6\,{\scriptstyle{\pm\,3}}$	$6\,{\scriptstyle{\pm\,6}}$	$\mathbf{26}\,{\scriptstyle{\pm\,10}}$
average	4	2	9	36

Appendix

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.

visual-antmaze	CRL	+ TD-Aug
medium-stitch	$82\,{\scriptstyle{\pm\,2}}$	$\mathbf{86}\,{\scriptstyle{\pm\,1}}$
large-stitch	$29\,{\scriptstyle{\pm\,6}}$	$\mathbf{36}\,{\scriptstyle{\pm\,4}}$
average	56	61

Appendix

Table 9: TD-Aug with Contrastive RL in visual environments.

Robot Policy Generalization

Robot policy generalization studies how robots can reuse learned behaviors across new tasks, objects, environments, viewpoints, and embodiments.

See all 179 papers

An Informational Curse of Horizon in Goal-Conditioned Policy Learning
Oct 7, 2026John L. Zhou, Yuxuan Dong, Jonathan C. KaoGoal-Conditioned RL Generalization in RL
University of California, Los Angeles
AnyGoal: On World Model Generalization Beyond Expert Rollouts
Oct 19, 2026Heejeong Nam, Divyat Mahajan, Quentin Le Lidec +2World Models World Model-Based Planning
Brown University · Cold Spring Harbor Laboratory · Mila, Université de Montréal +1
World Action Planner: Generalizable Decision-Making with Action-Conditioned World Models
Jul 30, 2026Xiangcheng Zhang, Yilun DuWorld Models for Robotics World Model-Based Planning
Harvard University

Oct 7, 2026J/K move · Enter open · S save

An Informational Curse of Horizon in Goal-Conditioned Policy Learning

John L. Zhou, Yuxuan Dong, Jonathan C. Kao

University of California, Los Angeles

Goal-Conditioned RL Generalization in RL Behavior Cloning Offline RL

Open paper AskPDF arXiv

The difficulty of learning goal-reaching policies is often attributed to a "curse of horizon" that manifests as bias accumulation in temporal-difference backups and noisy advantage estimates. In this work, we identify an additional informational curse of horizon in goal-conditioned policy learning, where increasing the goal relabeling horizon can significantly reduce policy generalization and performance. Through a series of controlled experiments with oracle planners, we decouple the goal horizons sampled during training from those that the policy is asked to reach at test time. Even when evaluated only on a sequence of nearby subgoals, goal-conditioned behavioral cloning (BC) policies suffer from severe, training horizon-dependent performance degradation that is mitigated by reinforcement learning (RL) objectives. We explain this phenomenon as a horizon-dependent decrease in the conditional mutual information between actions and hindsight-relabeled goals, and find empirically that both BC and RL policies trained on longer-horizon goals exhibit a shift in sensitivity from goal to state information, as measured by the policy's input Jacobians. Motivated by this observation, we find that distilling the input Jacobians of short-horizon policies into long-horizon policies yields significant performance gains, especially in combinatorial manipulation tasks. Taken together, our results highlight goal relabeling horizon as an important consideration when learning generalist policies from offline data.

Abstract

Robot Policy Generalization

Explore similar work

An Informational Curse of Horizon in Goal-Conditioned Policy Learning

AnyGoal: On World Model Generalization Beyond Expert Rollouts

World Action Planner: Generalizable Decision-Making with Action-Conditioned World Models