Closing the Gap between TD Learning and Supervised Learning -- A Generalisation Point of View
Raj Ghugare, Matthieu Geist, Glen Berseth, Benjamin Eysenbach
TL;DR
This work formalises stitching in goal-conditioned RL as a form of combinatorial generalisation across state–goal pairs that are not jointly seen in training data, and argues that standard OCBC methods based on supervised learning do not inherently possess this property. It provides a theoretical framework showing that the BC policy derived from a mixture of data-collecting policies need not match the test-time state–goal distribution, thereby lacking guaranteed combinatorial generalisation. To address this, the authors introduce temporal data augmentation that samples new goals from trajectories via nearby waypoint states, and prove that, under Lipschitz smoothness, this augmentation approximates one-step stitching distributions; empirically, augmentation improves stitching performance on both low- and high-dimensional tasks and reveals that data scale alone does not guarantee combinatorial generalisation. The work also entails new offline datasets designed to explicitly test stitching, demonstrating that purely SL-based RL methods struggle on unseen state–goal pairs while augmentation can recover substantial performance gains, with limitations tied to distance metrics and spurious correlations. Overall, the paper highlights the central role of combinatorial generalisation for data-efficient time-series learning and points to data augmentation as a simple, effective remedy with potential applicability beyond RL to other sequential domains like audio, video, and text.
Abstract
Some reinforcement learning (RL) algorithms can stitch pieces of experience to solve a task never seen before during training. This oft-sought property is one of the few ways in which RL methods based on dynamic-programming differ from RL methods based on supervised-learning (SL). Yet, certain RL methods based on off-the-shelf SL algorithms achieve excellent results without an explicit mechanism for stitching; it remains unclear whether those methods forgo this important stitching property. This paper studies this question for the problems of achieving a target goal state and achieving a target return value. Our main result is to show that the stitching property corresponds to a form of combinatorial generalization: after training on a distribution of (state, goal) pairs, one would like to evaluate on (state, goal) pairs not seen together in the training data. Our analysis shows that this sort of generalization is different from i.i.d. generalization. This connection between stitching and generalisation reveals why we should not expect SL-based RL methods to perform stitching, even in the limit of large datasets and models. Based on this analysis, we construct new datasets to explicitly test for this property, revealing that SL-based methods lack this stitching property and hence fail to perform combinatorial generalization. Nonetheless, the connection between stitching and combinatorial generalisation also suggests a simple remedy for improving generalisation in SL: data augmentation. We propose a temporal data augmentation and demonstrate that adding it to SL-based methods enables them to successfully complete tasks not seen together during training. On a high level, this connection illustrates the importance of combinatorial generalization for data efficiency in time-series data beyond tasks beyond RL, like audio, video, or text.