Driving · 2026-10-15
Imagining Goals with Geometric Horizon Models for Policy Generalization
Daniel Lawson, Dhruv Sreenivas, Adriana Hugessen, Khimya Khetarpal, Glen Berseth
AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.
TL;DR
The authors propose TD-Aug, which trains a geometric horizon model (GHM) using Temporal Difference Flows (TD-Flows) to sample counterfactual future outcomes from an offline mixture policy. They then use these generated outcomes as goals to augment goal-conditioned behavior cloning (BC). TD-Aug improves zero-shot compositional goal generalization on stitching-based OGBench tasks.
Why it matters
Zero-shot goal generalization for BC is presented as challenging, especially when success requires stitching across trajectories. By training an offline GHM to approximate the successor measure of a behavior-mixture policy and using its samples for goal augmentation, TD-Aug targets the compositional generalization gap that standard BC does not cover. The approach is evaluated across low-dimensional, high-dimensional, and manipulation tasks in OGBench.
Method
- Train a geometric horizon model (GHM) to approximate the successor measure of a behavior-mixture policy, using TD learning with TD-Flows.
- Use the GHM to sample future outcomes and perform goal augmentation for goal-conditioned BC (train on a mix of ground-truth and generated goals).
- Introduce a TD-Flow objective that combines MC and TD sampling (MCHC) to improve long-horizon generation.
Limitation
The authors do not state a specific limitation in the provided text.
Abstract (from 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.
Related papers
- Robot Generalization Beyond Volume: The Diversity Path via Zero-Shot Capabilities
- World-Model Policy Arbiter for Goal-Conditioned Reinforcement Learning
- AnyGoal: On World Model Generalization Beyond Expert Rollouts
- Rethinking World-Action Model for Compositional and In-Context Robotic Manipulation
- The Planning Limits of Latent World Models
Publication record
Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.
First public release: 2026-10-13 · Indexed by Pub-AI: 2026-10-15 UTC
Imagining Goals with Geometric Horizon Models for Policy Generalization
Preprint posted: 2026-10-13 · Original record ↗
- arxiv
- 2610.19129
Explicit source versions (1)
- v1 · 2026-10-13
- Preprint available
Source date: 2026-10-13 · First detected: 2026-10-25 UTC · Historical baseline
Source evidence (1)
- Official evidence ↗ · atom:published (official metadata)
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.