Contrastive Difference Predictive Coding

TL;DR

The paper tackles learning long-horizon temporal structure in time-series data for goal-conditioned RL by introducing TD InfoNCE, a temporal-difference version of the InfoNCE objective that estimates the discounted state occupancy measure $p^{\pi}(s_{t+} \mid s,a)$ in a data-efficient, off-policy manner. By deriving a TD-style loss and connecting it to a nonparametric successor representation, the authors develop a goal-conditioned RL algorithm that can stitch together disparate data and perform off-policy reasoning. Empirically, TD InfoNCE achieves strong performance on online and offline benchmarks, demonstrating improved sample efficiency (up to ~1500x over Monte Carlo InfoNCE in tabular settings) and robustness to stochastic dynamics, while also enabling data stitching and short-cut discovery in skewed datasets. The work significantly advances representation learning for temporally extended tasks, enabling more data-efficient planning and off-policy reasoning in goal-conditioned reinforcement learning.

Abstract

Predicting and reasoning about the future lie at the heart of many time-series questions. For example, goal-conditioned reinforcement learning can be viewed as learning representations to predict which states are likely to be visited in the future. While prior methods have used contrastive predictive coding to model time series data, learning representations that encode long-term dependencies usually requires large amounts of data. In this paper, we introduce a temporal difference version of contrastive predictive coding that stitches together pieces of different time series data to decrease the amount of data required to learn predictions of future events. We apply this representation learning method to derive an off-policy algorithm for goal-conditioned RL. Experiments demonstrate that, compared with prior RL methods, ours achieves $2 \times$ median improvement in success rates and can better cope with stochastic environments. In tabular settings, we show that our method is about $20 \times$ more sample efficient than the successor representation and $1500 \times$ more sample efficient than the standard (Monte Carlo) version of contrastive predictive coding.

Figures (14)

• Figure 1: TD InfoNCE is a nonparametric version of the successor representation. (Top) The distances between learned representations indicate the probability of transitioning to the next state and a set of randomly-sampled states. (Bottom) We update these representations so they assign high likelihood to (a) the next state and (b) states likely to be visited after the next state. See Sec. \ref{['sec:method']} for details.
• Figure 2: Evaluation on online GCRL benchmarks.(Left) TD InfoNCE performs similarly to or outperforms all baselines on both state-based and image-based tasks. (Right) On stochastic versions of the state-based tasks, TD InfoNCE outperforms the strongest baseline (QRL). Appendix Fig. \ref{['fig:online-eval']} shows the learning curves.
• Figure 3: Estimating the discounted state occupancy measure in a tabular setting.(Left) Temporal difference methods have lower errors than the Monte Carlo method. Also note that our TD InfoNCE converges as fast as the best baseline (successor representation). (Right) TD InfoNCE is more data efficient than other methods. Using a dataset of size 10M, TD InfoNCE achieves an error rate $25\%$ lower than the best baseline; TD InfoNCE also matches the performance of C-learning with $130\times$ less data.
• Figure 4: Stitching trajectories in a dataset. The behavioral policy collects "Z" style trajectories. Unlike the Monte Carlo method (contrastive RL) , our TD InfoNCE successfully "stitches" these trajectories together, navigating between pairs of (start ✖, goal ★) states unseen in the training trajectories. Appendix Fig. \ref{['fig:stitching-property-more']} shows additional examples.
• Figure 5: Searching for shortcuts in skewed datasets.(Left) Conditioned on different initial states ✖ and goals ★, we collect datasets with $95\%$ long paths (dark) and $5\%$ short paths (light). (Center) TD InfoNCE infers the shortest path, (Right) while contrastive RL fails to find this path. Appendix Fig. \ref{['fig:searching-shortcut-more']} shows additional examples.