# Imagining Goals with Geometric Horizon Models for Policy Generalization

arXiv: https://arxiv.org/abs/2610.19129 · HTML: https://mlsanity.com/paper/2610.19129/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 40 of its 49 references that are in the map.
- First posted on arXiv 13 Oct 2026.

## Builds on (most cited first, 30 of 40)

- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Hindsight Experience Replay](https://mlsanity.com/paper/1707.01495/) (2017)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [A Minimalist Approach to Offline Reinforcement Learning](https://mlsanity.com/paper/2106.06860/) (2021)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [Open X-Embodiment: Robotic Learning Datasets and RT-X Models](https://mlsanity.com/paper/2310.08864/) (2023)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1910.00177/) (2019)
- [Genie: Generative Interactive Environments](https://mlsanity.com/paper/2402.15391/) (2024)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Critic Regularized Regression](https://mlsanity.com/paper/2006.15134/) (2020)
- [Improving Sample Efficiency in Model-Free Reinforcement Learning from Images](https://mlsanity.com/paper/1910.01741/) (2019)
- [Contrastive Learning as Goal-Conditioned Reinforcement Learning](https://mlsanity.com/paper/2206.07568/) (2022)
- [OGBench: Benchmarking Offline Goal-Conditioned RL](https://mlsanity.com/paper/2410.20092/) (2024)
- [Flow Q-Learning](https://mlsanity.com/paper/2502.02538/) (2025)
- [Flow Matching in Latent Space](https://mlsanity.com/paper/2307.08698/) (2023)
- [Model-Based Value Estimation for Efficient Model-Free Reinforcement Learning](https://mlsanity.com/paper/1803.00101/) (2018)
- [HIQL: Offline Goal-Conditioned RL with Latent States as Actions](https://mlsanity.com/paper/2307.11949/) (2023)
- [Training Agents Inside of Scalable World Models](https://mlsanity.com/paper/2509.24527/) (2025)
- [Optimal Goal-Reaching Reinforcement Learning via Quasimetric Learning](https://mlsanity.com/paper/2304.01203/) (2023)
- [Synthetic Experience Replay](https://mlsanity.com/paper/2303.06614/) (2023)
- [Horizon Reduction Makes RL Scalable](https://mlsanity.com/paper/2506.04168/) (2025)
- [Generative Temporal Difference Learning for Infinite-Horizon Prediction](https://mlsanity.com/paper/2010.14496/) (2020)
- [1000 Layer Networks for Self-Supervised RL: Scaling Depth Can Enable New Goal-Reaching Capabilities](https://mlsanity.com/paper/2503.14858/) (2025)
- [Generative Trajectory Stitching through Diffusion Composition](https://mlsanity.com/paper/2503.05153/) (2025)
- [Closing the Gap between TD Learning and Supervised Learning -- A Generalisation Point of View](https://mlsanity.com/paper/2401.11237/) (2024)
- [Temporal Difference Flows](https://mlsanity.com/paper/2503.09817/) (2025)

## Abstract (arXiv)

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.