Policy generalization improved by imagining future goals with geometric horizon models
Imagining Goals with Geometric Horizon Models for Policy Generalization
Machine LearningArtificial Intelligence
Summary
Training robot policies using only fixed offline data can limit how well they handle new, unseen tasks. The authors focus on helping policies better combine what they already know to solve new problems without additional training. They introduce a technique called TD-Aug that imagines new possible future outcomes by sampling from a special geometric model to create novel goals. Using these imagined goals to train policies improves their ability to generalize and complete tasks requiring combining skills.
What this means in practice
- •For robotics engineers: Develop control policies that perform well on new robot tasks without additional training by training on imagined future goals.
- •For automation system developers: Improve automation policies to combine known behaviors for novel goals using imagined horizon outcomes from existing datasets.
Authors
Daniel Lawson, Dhruv Sreenivas, Adriana Hugessen, Khimya Khetarpal, Glen Berseth
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.