Recurrent Action Transformer with Memory

TL;DR

This paper introduces RATE, a transformer-based offline RL architecture with recurrent memory, enhancing decision-making in memory-dependent, partially observable environments while maintaining competitiveness on standard benchmarks.

Contribution

We propose RATE, a novel transformer architecture with integrated recurrent memory, improving long-term information retention in offline RL tasks.

Findings

RATE outperforms baselines in memory-intensive environments

RATE maintains competitive performance on standard benchmarks

Memory mechanisms are crucial for effective offline RL in POMDPs

Abstract

Transformers have become increasingly popular in offline reinforcement learning (RL) due to their ability to treat agent trajectories as sequences, reframing policy learning as a sequence modeling task. However, in partially observable environments (POMDPs), effective decision-making depends on retaining information about past events -- something that standard transformers struggle with due to the quadratic complexity of self-attention, which limits their context length. One solution to this problem is to extend transformers with memory mechanisms. We propose the Recurrent Action Transformer with Memory (RATE), a novel transformer-based architecture for offline RL that incorporates a recurrent memory mechanism designed to regulate information retention. We evaluate RATE across a diverse set of environments: memory-intensive tasks (ViZDoom-Two-Colors, T-Maze, Memory Maze, Minigrid-Memory, and POPGym), as well as standard Atari and MuJoCo benchmarks. Our comprehensive experiments demonstrate that RATE significantly improves performance in memory-dependent settings while remaining competitive on standard tasks across a broad range of baselines. These findings underscore the pivotal role of integrated memory mechanisms in offline RL and establish RATE as a unified, high-capacity architecture for effective decision-making over extended horizons. Code: https://sites.google.com/view/rate-model/.