Contrastive Retrospection: honing in on critical steps for rapid learning and generalization in RL

TL;DR

Contrastive Retrospection (ConSpec) is a novel reinforcement learning algorithm that uses offline contrastive learning to identify critical steps in a task, improving credit assignment, rapid learning, and generalization.

Contribution

ConSpec introduces a contrastive learning-based method to identify critical steps in RL, enhancing interpretability and out-of-distribution generalization.

Findings

Significantly improves learning speed in diverse RL tasks.

Enables rapid and interpretable credit assignment.

Facilitates out-of-distribution generalization.

Abstract

In real life, success is often contingent upon multiple critical steps that are distant in time from each other and from the final reward. These critical steps are challenging to identify with traditional reinforcement learning (RL) methods that rely on the Bellman equation for credit assignment. Here, we present a new RL algorithm that uses offline contrastive learning to hone in on these critical steps. This algorithm, which we call Contrastive Retrospection (ConSpec), can be added to any existing RL algorithm. ConSpec learns a set of prototypes for the critical steps in a task by a novel contrastive loss and delivers an intrinsic reward when the current state matches one of the prototypes. The prototypes in ConSpec provide two key benefits for credit assignment: (i) They enable rapid identification of all the critical steps. (ii) They do so in a readily interpretable manner, enabling out-of-distribution generalization when sensory features are altered. Distinct from other contemporary RL approaches to credit assignment, ConSpec takes advantage of the fact that it is easier to retrospectively identify the small set of steps that success is contingent upon (and ignoring other states) than it is to prospectively predict reward at every taken step. ConSpec greatly improves learning in a diverse set of RL tasks. The code is available at the link: https://github.com/sunchipsster1/ConSpec