Offline Reinforcement Learning as One Big Sequence Modeling Problem

TL;DR

This paper proposes viewing reinforcement learning as a sequence modeling problem using Transformer architectures, enabling flexible planning and achieving state-of-the-art results in long-horizon, sparse-reward tasks.

Contribution

It introduces a novel approach that applies sequence modeling techniques, like Transformers, to RL, simplifying design and improving performance in complex tasks.

Findings

Effective long-horizon dynamics prediction

State-of-the-art offline RL planning results

Versatility across multiple RL settings

Abstract

Reinforcement learning (RL) is typically concerned with estimating stationary policies or single-step models, leveraging the Markov property to factorize problems in time. However, we can also view RL as a generic sequence modeling problem, with the goal being to produce a sequence of actions that leads to a sequence of high rewards. Viewed in this way, it is tempting to consider whether high-capacity sequence prediction models that work well in other domains, such as natural-language processing, can also provide effective solutions to the RL problem. To this end, we explore how RL can be tackled with the tools of sequence modeling, using a Transformer architecture to model distributions over trajectories and repurposing beam search as a planning algorithm. Framing RL as sequence modeling problem simplifies a range of design decisions, allowing us to dispense with many of the components common in offline RL algorithms. We demonstrate the flexibility of this approach across long-horizon dynamics prediction, imitation learning, goal-conditioned RL, and offline RL. Further, we show that this approach can be combined with existing model-free algorithms to yield a state-of-the-art planner in sparse-reward, long-horizon tasks.