Contrastive Representations for Temporal Reasoning

TL;DR

This paper introduces CRTR, a novel contrastive learning method that captures temporal structure in representations, enabling efficient reasoning and solving complex puzzles like Rubik's Cube without external search.

Contribution

CRTR is the first method to effectively remove spurious features in temporal contrastive learning, improving generalization and reasoning over complex temporal domains.

Findings

CRTR outperforms standard contrastive learning in capturing temporal structure.

CRTR enables solving Rubik's Cube from arbitrary states with fewer search steps.

CRTR generalizes across all initial states of the Cube.

Abstract

In classical AI, perception relies on learning state-based representations, while planning, which can be thought of as temporal reasoning over action sequences, is typically achieved through search. We study whether such reasoning can instead emerge from representations that capture both perceptual and temporal structure. We show that standard temporal contrastive learning, despite its popularity, often fails to capture temporal structure due to its reliance on spurious features. To address this, we introduce Combinatorial Representations for Temporal Reasoning (CRTR), a method that uses a negative sampling scheme to provably remove these spurious features and facilitate temporal reasoning. CRTR achieves strong results on domains with complex temporal structure, such as Sokoban and Rubik's Cube. In particular, for the Rubik's Cube, CRTR learns representations that generalize across all initial states and allow it to solve the puzzle using fewer search steps than BestFS, though with longer solutions. To our knowledge, this is the first method that efficiently solves arbitrary Cube states using only learned representations, without relying on an external search algorithm.