Decision Transformer Integration in Hugging Face Transformers

TL;DR

Hugging Face has integrated the Decision Transformer into the transformers library and the Hugging Face Hub. This integration allows researchers to apply offline reinforcement learning (RL) by treating decision-making as a sequence modeling problem, removing the need for real-time environment interaction during training.

Offline Reinforcement Learning vs. Online RL

Offline reinforcement learning enables agents to learn optimal policies using static datasets collected from other agents or human demonstrations, rather than interacting with an environment through trial and error.

In contrast, online RL requires the agent to gather data directly from the environment. This approach often necessitates a high-fidelity simulator or direct real-world interaction, both of which can be expensive, complex to build, or insecure if the agent exploits simulator flaws.

While offline RL avoids these risks, it faces the "counterfactual queries problem," where the agent may encounter states or decide on actions for which no data exists in the training set.

The Decision Transformer Architecture

The Decision Transformer abstracts reinforcement learning as a conditional-sequence modeling problem. Instead of using traditional RL methods like fitting a value function to maximize cumulative reward (return), it uses a Transformer architecture to generate future actions based on a desired return, past states, and past actions.

Technical Workflow

  1. Input Sequence: The model takes the last K timesteps as input, consisting of three components: Return-to-go, State, and Action.
  2. Embedding: These inputs are processed via linear layers (for vector states) or CNN encoders (for image frames).
  3. Autoregressive Prediction: A GPT-2 based model processes these tokens to predict future actions autoregressively, effectively modeling the joint distribution of states, actions, and rewards.

By shifting the paradigm from reward maximization to generative trajectory modeling, the Decision Transformer generates a series of actions designed to achieve a specific target return.

Implementation in 🤗 Transformers

The Decision Transformer is now available in the transformers library, accompanied by nine pre-trained model checkpoints for continuous control tasks within the Gym environment (including Hopper, Walker2D, and Halfcheetah).

Model Loading and Usage

Users can load a pre-trained model using the DecisionTransformerModel class:

from transformers import DecisionTransformerModel

model_name = "edbeeching/decision-transformer-gym-hopper-expert"
model = DecisionTransformerModel.from_pretrained(model_name)

Autoregressive Action Prediction

Because the model is autoregressive, predictions at time-step t are conditioned on outputs from previous time-steps. The implementation requires preparing inputs (states, actions, returns-to-go, and timesteps) and padding them to the model's max_length before passing them through the model to predict the next action.

Performance Control via Target Return

One of the primary advantages of return-conditioned offline RL is the ability to control policy performance by adjusting the target return. This allows for the dynamic adjustment of agent difficulty, which is particularly useful for creating opponent bots in multiplayer settings.

Future Roadmap for Deep RL at Hugging Face

Hugging Face plans to further expand its support for the Deep Reinforcement Learning ecosystem with the following initiatives:

  • Integration of RL-baselines3-zoo.
  • Uploading a collection of pre-trained RL agents from rl-trained-agents (using stable-baselines3) to the Hub.
  • Integration of additional Deep RL libraries.
  • Implementation of Convolutional Decision Transformers specifically for Atari environments.

Sources