Transformers in Reinforcement Learning: A Survey

Abstract.

Transformers have significantly impacted domains like natural language processing, computer vision, and robotics, where they improve performance compared to other neural networks. This survey explores how transformers are used in reinforcement learning (RL), where they are seen as a promising solution for addressing challenges such as unstable training, credit assignment, lack of interpretability, and partial observability. We begin by providing a brief domain overview of RL, followed by a discussion on the challenges of classical RL algorithms. Next, we delve into the properties of the transformer and its variants and discuss the characteristics that make them well-suited to address the challenges inherent in RL. We examine the application of transformers to various aspects of RL, including representation learning, transition and reward function modeling, and policy optimization. We also discuss recent research that aims to enhance the interpretability and efficiency of transformers in RL, using visualization techniques and efficient training strategies. Often, the transformer architecture must be tailored to the specific needs of a given application. We present a broad overview of how transformers have been adapted for several applications, including robotics, medicine, language modeling, cloud computing, and combinatorial optimization. We conclude by discussing the limitations of using transformers in RL and assess their potential for catalyzing future breakthroughs in this field.

1. Introduction

Reinforcement learning (RL) is a learning paradigm that enables sequential decision-making by learning from feedback obtained through trial and error. It is usually formalized in terms of an Markov decision process (MDP), which provides a mathematical framework for modeling the interaction between an agent and its environment.

Most RL algorithms optimize the agent’s policy to select actions that maximize the expected cumulative reward. In deep RL, neural networks are used as function approximators for mapping the current state of the environment to the next action and for estimating future returns. This approach is beneficial when dealing with large or continuous state spaces that make tabular methods computationally expensive (Sutton and Barto, 1998) and has been successful in challenging applications (Arulkumaran et al., 2017; Nguyen et al., 2020; Latif et al., 2023). However, standard neural network architectures like convolutional neural networks (CNNs) and recurrent neural networks (RNNs) struggle with long-standing problems in RL. These problems include partial observability (Esslinger et al., 2022), inability to handle high-dimensional state and action spaces (Barto and Mahadevan, 2003), and difficulty in handling long-term dependencies (Chen et al., 2022a).

Partial observability is a challenge in RL (Liu et al., 2022a); in the absence of complete information, the agent may be unable to make optimal decisions. A typical way to address this problem is to integrate the agent’s input (Shao et al., 2019a) over time using CNNs and RNNs. However, RNNs tend to forget information (Pascanu et al., 2013), while CNNs are limited in the number of past-time steps they can process (Karpathy et al., 2014). Various strategies have been proposed to overcome this limitation, including gating mechanisms, gradient clipping, non-saturating activation functions, and manipulating gradient propagation paths (Ribeiro et al., 2020). Sometimes different data modalities, such as text, audio, and images are combined to provide additional information to the agent (Lathuilière et al., 2019; Song et al., 2021; Carta et al., 2020). However, integrating encoders for different modalities increases the model’s architectural complexity. With CNNs and RNNs, it is also difficult to determine which past actions contributed to current rewards (Ma et al., 2021). This is known as the credit assignment problem. These challenges and others, such as training instability, limit the scope of most RL applications to unrealistic virtual environments.

The transformer was first introduced in 2017 (Vaswani et al., 2017) and has rapidly impacted the field of deep learning (Lin et al., 2022), improving the state-of-the-art in natural language processing (NLP) and computer vision (CV) tasks (Tunstall et al., 2022; Khan et al., 2021; Devlin et al., 2019; Petit et al., 2021; Zhong et al., 2021). The key idea behind this neural network architecture is to use a self-attention mechanism to capture long-range relationships within the data. This ability to model large-scale context across sequences initially made transformers well-suited for machine translation tasks. Transformers have since been adapted to tackle more complex tasks like image segmentation (Petit et al., 2021), visual question answering (Zhong et al., 2021), and speech recognition (Dong et al., 2018).

This document surveys the use of transformers in RL. We begin by providing a concise overview of RL (Sec. 2.1) and transformers (Sec. 2.3) that is accessible to readers with a general background in machine learning. We highlight challenges that classical RL approaches face and how transformers can help deal with these challenges (Sec. 2.2 and 2.4). Transformers can be applied to RL in different ways (Fig. 1). We discuss how they can be used to learn representations (Sec. 3), model transition functions (Sec. 4), learn reward functions (Sec. 5) and learn policies (Sec. 6). In Sec.  7 and Sec.  8, we discuss different training and interpretation strategies, and in Sec. 9, we provide an overview of RL applications that use transformers, including robotics, medicine, language modeling, edge-cloud computing, combinatorial optimization, environmental sciences, scheduling, trading, and hyper-parameter optimization. Finally, we discuss limitations and open questions for future research (Sec. 10). With this work, we aim to inspire further research and facilitate the development of RL approaches for real-world applications.

Figure 1. This survey presents a comprehensive overview of the use of transformers in RL. Modeling an RL policy may involve representation learning, modeling the transition function, reward function learning, and policy learning. Transformers can be used across all of these tasks.

2. Background

This section introduces the fundamental concepts of RL and discusses associated challenges. We also provide an overview of transformers and their potential advantages in RL.

2.1. Reinforcement Learning

Reinforcement learning (RL) is a reward-based learning paradigm that enables agents to learn from their experience and improve their performance over time. This is commonly formulated in terms of Markov decision processes (MDPs), in which the agent chooses an action $(𝐚\in 𝐀)$ based on the state $(𝐬\in 𝐒)$ of the environment and receives feedback in the form of rewards $(r\in ℝ)$. The MDP framework assumes that the environment satisfies the Markov property, which asserts that the next state is independent of the past states, given the present state and the most recent action. This allows the agent to make decisions based only on the current state without tracking the history of previous states and actions. After taking action, the agent receives feedback from the environment in the form of a reward. The state is updated to a new value $({𝐬}^{\prime }\in 𝐒)$, as determined by the transition function $p({𝐬}^{\prime }\mid 𝐬,𝐚)$ describing how the environment responds to the agent’s action. The overall goal of RL is to learn how to solve a multi-step problem by maximizing the return $g$, which is the total discounted reward:

(1) $$g=\sum _{t=0}^T{\gamma }^t\cdot r_t,$$

where $r_t$ is the reward received at each time step $t$, $T$ is the total number of time steps in the given episode, and $\gamma \in [0,1]$ is a discount factor, affecting the importance given to immediate rewards versus future rewards in a given episode.

There are multiple categories of RL algorithms, each with advantages and disadvantages (AlMahamid and Grolinger, 2021). Choosing the appropriate category relies on factors such as the problem’s complexity, the size of the state and action spaces, and available computational resources. We now briefly review these categories.

Model-Based RL. In model-based RL, a transition function is learned using transitions ($𝐬,𝐚,r,{𝐬}^{\prime }$) generated by environmental interaction. This transition function models the probability distribution $p({𝐬}^{\prime },r\mid 𝐬,𝐚)$ over the subsequent state ${𝐬}^{\prime }$ and reward $r$ given the current state $𝐬$ and action $𝐚$. By leveraging this learned model, the agent plans and selects actions that maximize the expected return $g$. However, this approach can be computationally expensive and may suffer from inaccuracies in the learned model, leading to sub-optimal performance (Moerland et al., 2020). Despite these drawbacks, this approach is sample efficient. In other words, it can achieve good performance using relatively few interactions with the environment compared to other methods (Mohanty et al., 2020; Wang et al., 2022c).

Model-Free RL. In model-free RL, optimal actions are learned by direct interaction with the environment. Methods from this category cannot model state-transition dynamics to plan actions, which can result in slower convergence and lower sample efficiency compared to model-based RL (Yarats et al., 2021). However, model-free RL is more adaptable to environmental changes, making it more robust in complex or noisy environments (Mnih et al., 2013; Lillicrap et al., 2016; Schulman et al., 2017). Additionally, it is less computationally expensive as it does not need to learn a model of the environment.

Further, RL methods can be categorized as on-policy, off-policy, and offline RL depending on how the data collection relates to the policy being learned.

On-Policy RL. On-policy RL approaches use the current policy to gather transitions for updating the value function. For instance in SARSA (Sutton and Barto, 1998), the current policy is used to collect a tuple $(𝐬,𝐚,r,{𝐬}^{\prime },{𝐚}^{\prime })$, consisting of the current state-action pair $𝐬,𝐚$, the immediate reward $r$, and the next state-action pair ${𝐬}^{\prime },{𝐚}^{\prime }$. This is then used for estimating the return of the current state-action pair $Q_{\text{target}}(𝐬,𝐚)$ and the next state-action pair $Q_{\text{target}}({𝐬}^{\prime },{𝐚}^{\prime })$. The value function $Q$ is updated using the following temporal difference (TD) learning rule:

(2) $$Q_{\text{target}}(𝐬,𝐚)\leftarrow Q_{\text{target}}(𝐬,𝐚)+\alpha \cdot \left[r+\gamma \cdot Q_{\text{target}}({𝐬}^{\prime },{𝐚}^{\prime })-Q_{\text{target}}(𝐬,𝐚)\right],$$

where $\alpha \in (0,1]$ represents the learning rate. Although on-policy methods are comparatively easy to implement, they have several drawbacks. They tend to be sample inefficient (Larsen et al., 2021), requiring significant interaction with the environment to achieve good performance. Additionally, they can be susceptible to policy oscillation and instability (Young and Sutton, 2020), and they lack the flexibility to explore, slowing down the learning process and resulting in sub-optimal policies.

Off-Policy RL. Off-policy RL strategies use two policies — a behavior policy and a target policy. The behavior policy collects data that is subsequently used to estimate the expected return of an action under the given target policy. Since the behavior policy is used for data collection, it can explore different states and actions without affecting the current target policy. Thus, off-policy methods are well-suited for understanding the value of a given state and action. Usually, the target policy is updated using the behavior policy via importance sampling (IS). This adjusts the value estimates of the target policy based on the IS ratio between the behavior $b(𝐚|𝐬)$ and target $\pi (𝐚|𝐬)$ policies:

(3) $$\text{IS}=\frac{\pi (𝐚|𝐬)}{b(𝐚|𝐬)},$$

and the final value estimate is given as:

(4) $$Q_{\text{target}}(𝐬,𝐚)\leftarrow Q_{\text{target}}(𝐬,𝐚)+\alpha \cdot \left[\text{IS}\cdot (r+\gamma \cdot Q_{\text{target}}({𝐬}^{\prime },{𝐚}^{\prime }))-Q_{\text{target}}(𝐬,𝐚)\right].$$

Offline-RL. Offline RL or batch RL (Levine et al., 2020a) uses a static dataset of transitions, denoted by $𝐃={\{{𝐬}_{𝐭},{𝐚}_{𝐭},r_t,{𝐬}_{𝐭}^{{}^{\prime }}\}}_{t=1}^T$, collected using a behavior policy, and so does not require interaction with the environment to collect trajectories. Offline-RL updates the state-action value function $Q_{\text{target}}$ as:

(5) $$Q_{\text{target}}(𝐬,𝐚)\leftarrow Q_{\text{target}}(𝐬,𝐚)+\alpha \cdot \left[r+\gamma \cdot \underset{{𝐚}^{\prime }}{\max }(Q_{\text{target}}({𝐬}^{\prime },{𝐚}^{\prime }))-Q_{\text{target}}(𝐬,𝐚)\right],$$

where the max operator estimates the maximum expected return over all possible actions in the next state. Offline RL is a more practical strategy for safety-critical applications as it does not require interactions with the environment (Shi et al., 2021; Killian et al., 2023). However, the static nature of the dataset does not allow the agent to explore and adapt to new information, potentially limiting performance (Lu et al., 2022).

Multi-Agent Reinforcement Learning. Online, offline, and off-policy learning setups can be used to facilitate adaptive decision-making in dynamic environments with multiple interacting agents (Gronauer and Diepold, 2022). Each of the $I$ agents has its own policy ${\pi }_i({𝐚}_i|{𝐬}_i)$, state space ${𝐒}_i$, and action space ${𝐀}_i$. The agents interact with each other and the environment, and their actions can affect the outcomes of other agents. The goal in multi-agent reinforcement learning (MARL) is to learn a joint policy $\pi ({𝐚}_1,{𝐚}_2,\mathrm{\dots },{𝐚}_I|{𝐬}_1,{𝐬}_2,\mathrm{\dots },{𝐬}_I)$ that maximizes the collective reward of all agents. Formally, the objective in MARL can be expressed as maximizing the expected sum of discounted rewards of all the agents:

(6) $$𝔼\left[\sum _{t=0}^{\mathrm{\infty }}{\gamma }^t\sum _{i=1}^I{r}_i^{(t)}\right],$$

where $\gamma$ is the discount factor and $r_i^{(t)}$ is the reward received by agent $i$ at time $t$. In general, there are two main approaches to MARL: decentralized policies (Zhang et al., 2018a), and centralized training (Sharma et al., 2021). Decentralized policies involve each agent independently learning its policy without explicit coordination or communication with other agents. In contrast, centralized training uses a shared value function that considers joint states and actions, enabling agents to communicate and coordinate their actions through a communication protocol. Communication protocols (Foerster et al., 2016) facilitate information exchange and collaboration among agents.

Upside Down RL. Classical RL usually involves optimizing policies by estimating the expected future return. Upside-down RL (Schmidhuber, 2019; Arulkumaran et al., 2022) flips the traditional RL paradigm and uses the desired return $g$, the horizon $h$ (i.e., the time remaining until the end of the current trial), and the state as inputs. This input acts as a command which is mapped to action probabilities. Upside-down RL offers improved stability compared to classical RL as it avoids the need to estimate the value function, which can introduce instabilities in traditional RL algorithms (Chen et al., 2021a; Sutton and Barto, 1998). The loss function of upside-down RL can be defined as:

(7) $$ℒ(𝜽)=\sum _{t=0}{\left[{𝐚}_t-f({𝐬}_t,g_t,h_t,𝜽)\right]}^2,$$

where $𝜽$ contains the model parameters. The term ${𝐚}_t$ is the action at time step $t$, and $f({𝐬}_t,g_t,h_t,𝜽)$ is the predicted action when conditioned on state ${𝐬}_t$, expected future return $g_t$, and horizon $h_t$.

2.2. Challenges in Reinforcement Learning

In this section, we discuss the different challenges of classical RL algorithms.

Curse of Dimensionality. Real-world applications often involve high-dimensional state spaces, which makes it hard for classical RL algorithms to learn optimal policies (Barto and Mahadevan, 2003). This is because the required training data grows exponentially as the data dimensionality increases (Wang et al., 2020a). One way to mitigate this problem is to encode high-dimensional states into a lower-dimensional space. RL policies perform better when trained on encoded low-dimensional data (Yarats et al., 2021).

Partially Observable Environment. A partially observable environment presents a challenge for RL algorithms, as the agent cannot access observations that contain complete information about the environmental state at each time-step (Vinyals et al., 2019; OpenAI et al., 2019). Without complete information, the algorithm may struggle to make the best decision, leading to uncertainty and compromises in performance. To address this, the policy must maintain an internal representation of the state, often in the form of memory, from which the actual state can be estimated (Icarte et al., 2020). Historically, this has often been done with RNNs, but these cannot efficiently model long contexts (Pascanu et al., 2013; Hochreiter, 1998).

Credit Assignment. The term credit assignment refers to the problem of associating the actions taken by an agent with the reward it receives (Mesnard et al., 2021). This is challenging for two reasons: First, the reward may be delayed; the agent may not be able to observe the consequences of its actions until several time steps into the future. Second, other factors or multiple actions may influence the received reward, making it challenging to identify which action led to that reward. Inaccurate credit assignment can lead to slower training and sub-optimal policies (Ausin et al., 2021). Moreover, when the reward is sparse (i.e., when the agent receives little feedback for its actions), the credit-assignment problem becomes even more difficult (Seo et al., 2019). One potential solution is to use models that integrate information across all time steps, which may be better suited for solving this issue (Chen et al., 2021a).

Recent studies have exploited transformers to tackle these three key challenges. Transformers have demonstrated success in modeling long-term dependencies in sequential data while showing promising results in promoting generalization and faster learning in domains such as NLP and CV. We now provide a brief overview of transformers and explore the various ways in which they have been applied to learning optimal RL policies.

2.3. Transformers

Figure 2. a) The dot-product attention mechanism (for 4 embeddings of size 5 each). Input features ($𝐗$) are processed using the key ($𝐊$), query ($𝐐$), and value ($𝐕$) tensors. Each query undergoes a dot product with every key, and the result is normalized using the softmax operator to compute the attention. This is then used to weigh the values to produce the output. b) The (multi-head) self-attention mechanism is just one component in the transformer block, which also contains residual links, layer normalization, and parallel MLPs that process each input embedding separately.

Transformers are a class of neural network architectures consisting of multiple layers, each containing a multi-head self-attention mechanism, parallel fully-connected networks, residual connections, and layer normalization. Given a sequence of $N$ input embeddings, transformers produce a sequence of $N$ output embeddings, each of which represents the relationship between the corresponding input embedding and the rest of the input sequence. In NLP, the input embeddings may represent words from a given sentence, while in RL, they may represent different states.

The self-attention mechanism allows each input embedding (each row of $𝐗$) to simultaneously attend to all the other embeddings in the input sequence. It computes one attention score for each pair of inputs. This is done by projecting each input into a query $𝐐={\mathrm{𝐗𝐖}}_q$ and a key $𝐊={\mathrm{𝐗𝐖}}_k$ tensor. The attention scores are then computed by taking the dot product of each query vector (row of the query tensor) with every key vector (row of the key tensor), followed by a softmax operation that normalizes the resulting scores such that they add up to one for each query. The attention scores are then used to compute a weighted sum of value $𝐕={\mathrm{𝐗𝐖}}_v$ tensors (see Fig. 2):

(8) $$\text{Attention}(𝐐,𝐊,𝐕)=\mathrm{Softmax}\left(\frac{{\mathrm{𝐐𝐊}}^{\top }}{\sqrt{d_q}}\right)\cdot 𝐕.$$

To help stabilize gradients during training, the dot product is scaled by a factor of $\sqrt{d_q}$, where $d_q$ is the dimension of the query tensor.

Transformers often compute multiple sets of attention scores in parallel (each with a different set of learned parameters ${𝐖}_k,{𝐖}_q,{𝐖}_v$). This allows the model to attend to multiple aspects of the input sequence simultaneously and is known as multi-head self-attention. Each attention “head” output is concatenated and linearly transformed to produce the final output representation. For applications where the order of the inputs is important, a position encoding is added that allows the network to establish the position of each input.

Figure 3. a) Self-attention computes attention scores for a given segment by considering all input across all time steps. b) Masked self-attention considers the current and past time steps while disregarding the future. c) Transformer XL (Dai et al., 2019) establishes a connection between past segments, incorporating information from multiple previous segments. It calculates attention scores by combining the hidden state of the previous segment, enabling the modeling of longer-term dependencies.

In a transformer block (Fig. 2b), a residual connection is placed around the multi-head self-attention mechanism. This improves training stability by allowing the gradient to flow easily through the network. The output is then processed using layer normalization, which normalizes the activations of each layer across the feature dimension. Each output is processed in parallel by the same MLP. Once more, these are bypassed by a residual connection, and a second layer norm is subsequently added.

Architectural Variations. The bidirectional encoder representations from transformers (BERT) model (Devlin et al., 2019) and the generative pre-trained transformer (GPT) (Radford et al., 2018) are two popular variants of the transformer architecture. BERT (Fig. 4) is a transformer encoder in which each output receives information from every input in the self-attention mechanism (Fig. 3a). The goal is to process the incoming data to generate a latent representation that integrates contextual information. This can be particularly useful in RL, as it enables the agent to make informed decisions based on a more comprehensive understanding of the environment (Banino et al., 2022; Wang et al., 2023).

Figure 4. The BERT architecture. The input tokens are converted to word embeddings and passed through transformer layers to produce latent representations. Some input tokens are replaced with mask embedding as part of the formulation for the pre-training task. The output embeddings predict the missing (masked) word using a softmax and multi-class classification loss. In each loss term corresponding to the prediction of a masked word, BERT uses context from before and after the word.

Conversely, GPT (Fig. 5) uses a decoder architecture to auto-regressively generate a sequence of output tokens, considering only the past tokens. The use of masked self-attention (Fig. 3b) prevents it from cheating by looking ahead to tokens that it should not know yet during training by clamping the associated attention values to zero. In RL, this autoregressive nature can be used to implement an RL policy that is conditioned on a sequence of past states and actions (Hernandez et al., 2021; Chen et al., 2021a; Janner et al., 2021). GPT uses multiple blocks, each containing a multi-head attention mechanism. The original transformer (Vaswani et al., 2017) combined these two approaches in an encoder-decoder architecture for machine translation. An encoder architecture processes the incoming sentence, and the decoder auto-regressively produces the output sentence. In doing so, the decoder also considers the attention over the encoder’s latent representation using a “cross-attention” block.

Figure 5. The GPT architecture. The tokens are mapped to embeddings through transformer layers with masked self-attention, such that each position attends only to earlier tokens. The goal is to predict the next token at each position correctly. This approach is efficient to train as every word contributes to the final loss.

Vision Transformers. Inspired by the success of transformer-based architectures like BERT and GPT in NLP, Dosovitskiy et al. (2021) proposed the vision transformer (ViT) architecture for processing images. ViT architecture (Dosovitskiy et al., 2021) is suited to a wide range of RL tasks where images must be used to learn a policy (Tao et al., 2022; Kargar and Kyrki, 2021; Goulão and Oliveira, 2022). The ViT architecture is a transformer encoder that processes patches of the image (Fig. 6). Each patch is combined with a positional encoding that provides knowledge about its original image location.

Transformer-XL. The computational complexity of the self-attention mechanism increases quadratically with input sequence length due to an exponential increase in pairwise comparison. Hence, the transformer architectures discussed so far typically partition long input sequences into shorter sequences to reduce memory demands. While this approach helps minimize memory usage, it makes capturing global context challenging. Moreover, the conventional transformer model is limited because it does not consider the boundaries of the input sequence when forming context. Instead, it selects consecutive chunks of symbols without regard for a sentence or semantic boundaries. This can result in context fragmentation, where the model lacks the necessary contextual information to accurately predict the first few symbols in a sequence. The transformer-XL (TrXL) architecture (Dai et al., 2019) addresses these issues by dividing the input into segments and incorporating segment-level recurrence (Fig. 3c) and relative positional encodings. By caching and reusing the representation computed for the previous segment during training, TrXL can extend the context and better capture long-term dependencies. Additionally, TrXL can process elements of new segments without recomputing the past segments, leading to faster inference.

2.4. Key Advantages of Transformers in RL

This section outlines the transformer characteristics that are important for RL applications.

Attention Mechanism. The attention mechanism is crucial in transformers for sequence modeling of states (Benjamins et al., 2022). It enables the RL agent to focus selectively on relevant cues in the environment (Manchin et al., 2019) and ignore redundant features, which leads to faster training. This is particularly useful in high-dimensional state spaces where there are a large number of input elements.

Multi-Modal Architectures. For complex tasks, RL agents may require additional information from different data modalities (Zhang et al., 2018b; Kiran et al., 2022). Past approaches have used different architectures to handle multiple modalities (Ramachandram and Taylor, 2017). However, transformers can process multiple modalities of data (e.g., text, images) effectively (Jaegle et al., 2022; Xu et al., 2022c) using the same architecture.

Parallel Processing. Learning a policy in RL can be computationally expensive, especially for complex tasks requiring many samples (Obando-Ceron and Castro, 2021). RNNs require the sequential processing of inputs, which is inefficient. Transformers are well-suited for parallelization due to their self-attention mechanism, which considers all inputs simultaneously. RL algorithms can exploit transformers to learn more efficient policies in significantly less time.

Scalability. Current RL algorithms struggle to scale effectively to complex tasks that require the integration of multiple skills (Zhan et al., 2017; Kalashnikov et al., 2018). However, the performance of transformers has been shown to improve smoothly as the size of the model, dataset, and compute increases (Kaplan et al., 2020). This ability can potentially be leveraged in RL to create generalist agents capable of performing various tasks in different environments and with different embodiments (Lee et al., 2022).

These points highlight how the properties of transformers make them attractive for RL. In the following sections, we examine the use of transformers in each stage of the RL workflow, including representation learning (Sec. 3), policy learning (Sec. 6), learning the reward function (Sec. 5), and modeling the environment (Sec. 4). Additionally, we cover various training strategies (Sec. 7) and techniques for interpreting RL policies that use transformers (Sec. 8).

3. Representation Learning

Concise and meaningful representations are critical for efficient decision-making (Lesort et al., 2018) in RL. Empirical evidence has shown that training agents directly on high dimensional data, such as image pixels, is sample inefficient (Lake et al., 2016). Therefore, good data representations are crucial for learning RL policies (Lesort et al., 2018; Laskin et al., 2020) since they can enhance performance, convergence speed, and policy stability (Ghosh and Bellemare, 2020).

For instance, in self-driving cars, the raw sensory inputs ${𝐨}_t$ (e.g., camera images, LIDAR readings) are high dimensional and often contain redundant information. If these inputs ${𝐨}_t$ are mapped to a compact representation ${𝐬}_t$, the RL agent can learn more efficiently. Similarly, in game-playing scenarios (Fig. 6), it’s helpful to encode and extract relevant features from the pixel to be used as input to the RL algorithm learning the policy. Transformers can produce transferable and discriminative feature representations for diverse data modalities (Zhou et al., 2021; Brown et al., 2020; Zhang et al., 2022d; Choi et al., 2020; Ying et al., 2021).

Figure 6. Representation learning simplifies the information given to RL agents by condensing the original high-dimensional input (here a frame from an Atari game). Transformers can learn low-dimensional representations $𝐬$ from high-dimensional observations $𝐨$. In this case, this is achieved by employing a vision transformer. The input is provided in multiple patches, and the transformer encoder learns a shared low-dimensional representation contextualized using a <cls> input token.

3.1. Comparisons between Transformers, CNNs, and GNNs

Encoding high-dimensional representations using pre-trained CNN and transformer architectures is an active research area. Both approaches yield comparable performance in computer vision tasks (Woo et al., 2023; Liu et al., 2022c; Malpure et al., 2021). However, several studies (Zhou et al., 2021; Zhang et al., 2022d) have shown that transformers generate more expressive representations than CNNs for tasks where the data distribution differs at training and test time. This advantage stems from CNN’s inherent inductive bias towards local spatial features, which limits their ability to capture the global dependencies necessary for reasoning (Vo et al., 2017). Transformers can encode the image as a sequence of patches without local convolution and resolution reduction (Fig. 6). Hence, they model the global context in every layer, leading to a stronger representation for learning efficient policies (Dosovitskiy et al., 2021). Transformers exhibit comparable generalization capabilities to graph neural networks (GNNs) in graphs (Dwivedi and Bresson, 2020) and, in some instances, outperform them by capturing long-range semantics (Ying et al., 2021).

Multi-task reinforcement learning (MTRL) is a learning paradigm where an agent is trained to perform multiple tasks simultaneously. It has traditionally relied on GNNs to handle incompatible environments (i.e., differing state-action spaces) (Wang et al., 2018; Huang et al., 2020a). This is due to the ability of GNNs to operate on graphs of variable sizes. However, Kurin et al. (2021) hypothesize that the restrictive nature of message passing in sparse graphs may adversely impact performance. They propose replacing GNNs with transformers which obviates the need to learn multi-hop communication; the transformer can be considered a GNN applied to fully-connected graphs with attention as an edge-to-vertex aggregation operation (Battaglia et al., 2018). This enables a dedicated message-passing scheme for each state and pass, effectively avoiding the requirement for multi-hop message propagation. This overcomes the challenges of gradient propagation and information loss arising from such multi-hop propagation. The transformer-based model of Kurin et al. (2021), Amorpheus, learns better representations and improves performance without imposing a relational inductive bias.

3.2. Advanced Representation Learning using Transformers

Transformers, in combination with other attention mechanisms, enable the learning of expressive representations. SloTTAr (Gopalakrishnan et al., 2023) combines a transformer encoder-decoder architecture with slot attention (Locatello et al., 2020). The transformer encoder focuses on learning spatio-temporal features from action-observation sequences. Utilizing the slot attention mechanism, features are grouped at each temporal location, resulting in $K$ slot representations. The decoder subsequently decodes these slot representations to generate action logits. Notably, this parallelizable process enables faster training compared to existing benchmarks.

In multi-agent reinforcement learning, transformers have proven effective in modeling relations among agents and the environment (Zhang et al., 2022b). Li et al. (2022a) proposed replacing RNNs with a transformer encoder for robust temporal learning. Similarly, Zhang et al. (2022a) used a visual feature extractor based on the ViT architecture to obtain more robust representations for robotic visual exploration. Their network, utilizing self-attention, outperformed CNN backbones in robotic tasks.

Transformers have been widely adopted in scenarios involving the processing of multimodal information. Yang et al. (2022) introduced scene-fusion transformers that fuse observed trajectories and scene information to generate expressive representations for trajectory prediction. To reduce computational complexity, they employ sparse self-attention. Zhang et al. (2023) utilize transformers to integrate visual and text features effectively.

3.3. Enhancing Transferability and Generalization

An inherent difficulty faced by RL is generalizing to new unseen tasks (Levine et al., 2020b). This difficulty results from the intrinsic differences between various RL tasks (e.g., autonomous driving and drug discovery). While meta-learning methods such as model-agnostic meta-learning (MAML) (Finn et al., 2017) have been developed to generalize to new tasks with different distributions using limited data, these methods are hard to use in RL due to poor sample-efficiency and unstable training (Liu et al., 2019).

Transformers have shown great potential for meta-reinforcement learning (TrMRL), as demonstrated by Melo (2022). Transformers are excellent at handling long sequences and capturing dependencies over a long period, which enables them to adapt to new tasks using self-attention quickly. In TrMRL, the proposed agent uses self-attention blocks to create an episodic memory representing a consensus of recent working memories. The transformer architecture encodes the working memory and tasks as a distribution over these memories. During meta-training, the agent learns to differentiate tasks and identify similarities in the embedding space. This approach performs comparably or better than PEARL (Rakelly et al., 2019) and MAML (Finn et al., 2017). It is particularly efficient in memory refinement and task association.

Shang et al. (2022) introduced the state-action-reward transformer (StARformer) to model multiple data distributions by learning transition representations between individual time steps of the state, action, and reward. StARformer consists of the step transformer and the sequence transformer. The step transformer uses self-attention to capture a local representation that understands the relationship between state-action-reward triplets within a single time-step window. The sequence transformer combines these local representations with global representation in the form of pure state features extracted as convolutional features, introducing a Markovian-like inductive bias. This bias helps reduce model capacity while effectively capturing long-term dependencies.

4. Transition Function Learning

The transition function $p({𝐬}^{\prime },r|𝐬,𝐚)$ describes how the environment transitions from the current state $𝐬$ to the next state ${𝐬}^{\prime }$ and issues rewards $r$ in response to the actions $𝐚$ taken by the agent. Learning this function (Fig. 7) and subsequently exploiting it to train an RL agent is known as model-based RL (Sec. 2.1). Model-based RL offers a significant advantage compared to model-free RL approaches (Moerland et al., 2020); it allows the agent to plan future trajectories for each action, improving robustness and safety. Interactions with the external environment can be computationally expensive, particularly when relying on simulations that mimic the real world (Featherstone, 2014). If we learn the transition function, these interactions can be reduced.

Figure 7. A transformer can be employed to learn the dynamics of the simulator by modeling the transition function. This is accomplished by leveraging pre-collected trajectories of environment interaction. The transformer-based world model is trained using these trajectories. An RL policy $\pi$ is then trained to learn the task using this learned world model without requiring further environmental interaction.

A standard method in model-based reinforcement learning (MBRL) (Ha and Schmidhuber, 2018) involves training an end-to-end world model to represent the environment’s dynamics accurately. For instance, TransDreamer (Chen et al., 2022a) trains a single model that learns visual representations and dynamics using the evidence lower bound loss (Kingma and Welling, 2019). However, this approach can result in inaccuracies in the learned world model.

The masked world model (MWM) (Seo et al., 2022) addresses this by decoupling visual representation and dynamics learning. This framework utilizes an autoencoder with convolutional layers and ViT to learn visual representations. The autoencoder reconstructs pixels based on masked convolutional features. A latent dynamics model is learned by operating on the representations from the autoencoder. An auxiliary reward prediction objective is introduced for the autoencoder to encode task-relevant information. Importantly, this approach outperforms the strong RNN-based model, DreamerV2 (Hafner et al., 2021) in terms of both sample efficiency and final performance on various robotic tasks.

Learning the dynamics of the world ${𝐳}_{t+1}\sim p_G({𝐳}_{t+1}|{𝐳}_{\le t},{𝐚}_{\le t})$ has been formulated as a sequence modeling problem in imagination with auto-regression over an inner speech (IRIS) (Micheli et al., 2022). This approach takes advantage of the transformer’s ability to process sequences of discrete tokens. IRIS uses a discrete autoencoder to construct a language of image tokens, while a transformer models the dynamics over these tokens. By simulating millions of trajectories accurately, IRIS surpasses recent methods in the Atari 100k benchmark (Bellemare et al., 2015) in just two hours of real-time experience.

Building upon the auto-regressive nature of transformer decoders, Robine et al. (2023) introduces the transformer-based world model (TWM). Based on the TrXL architecture, TWM learns the transition function from real-world episodes while attending to latent states, actions, and rewards associated with each time step. By allowing direct access to previous states instead of viewing them through a compressed recurrent state, the TrXL architecture enables the world model to learn long-term dependencies while maintaining computational efficiency.

5. Reward Learning

The reward function is crucial in RL as it quantifies the desirability of different actions $𝐚$ for a given state $𝐬$, guiding the learning process. Typically, reward functions are predefined by human experts who carefully consider relevant factors based on their domain knowledge. However, designing an appropriate reward function is challenging in real-world scenarios, requiring a deep understanding of the problem domain. Moreover, manually designing it introduces bias and may lead to sub-optimal behavior.

Recent research has explored different approaches for learning reward functions by integrating human data in various forms, such as real-time feedback, expert demonstrations, preferences, and language instructions. Transformers have proven valuable in these contexts. The transformer architecture is particularly advantageous with non-Markovian rewards, which are characterized by delays and dependence on the sequence of states encountered during an episode (e.g., when rewards are only provided at the end). Transformers efficiently capture dependencies across input sequences, making them well-suited to handle such scenarios.

Figure 8. Nakatani et al. (2022) use a BERT-based reward function $r(𝐬,𝐚)$ to provide feedback to the RL policy $\pi$ to learn the task of machine translation. Given the input sentence $𝐬$, the policy $\pi$ translates it into a different language. This predicted action $𝐚$ is then compared with the optimal action ${𝐚}^{+}$ using BERT, which provides feedback to the agent in the form of a reward $r$.

The preference transformer (Kim et al., 2023) model captures human preferences by focusing on crucial events and modeling the temporal dependencies inherent in human decision-making processes; it effectively predicts non-Markovian rewards and assigns appropriate importance weights based on the trajectory segment. This approach reduces the effort required for designing reward functions and enables handling complex control tasks such as locomotion, navigation, and manipulation.

To train an RL policy for generating text that aligns with human-labeled ground truth, the bilingual evaluation understudy (BLEU) score (Papineni et al., 2002) is often used as a reward function. However, BLEU may not consistently correlate strongly with human evaluation. In (Nakatani et al., 2022), a BERT-based reward function is introduced, demonstrating a higher correlation with human evaluation. This approach leverages a pre-trained BERT model (Fig. 8) to assess the semantic similarity between the generated and reference sentences and update the policy accordingly.

6. Policy Learning

Policy learning is central to RL; it involves learning the policy $\pi (𝐬)$ which the agent uses to select actions $𝐚=\pi (𝐬)$ with the objective of maximizing the discounted cumulative reward $g$. Transformers have been used for modeling $\pi (𝐬)$ in various scenarios, including off-policy, on-policy, and offline RL.

6.1. Offline RL with the Decision Transformer

Offline RL trains a policy using a limited, static dataset of previously collected experiences. This is different from online or off-policy RL approaches (which continuously interact with the environment to update their policies) since the agent cannot collect experience beyond the fixed dataset, which limits its ability to learn, explore, and improve performance.

Figure 9. The decision transformer approaches offline RL as a sequence-prediction task. The network uses a sequence of states, actions, and return-to-go (RTG) to predict the next action. At inference, the RTG acts as a condition for the policy to generate a given set of actions.

The decision transformer (DT) (Chen et al., 2021a) (Fig. 9) is an offline RL method that uses the upside-down RL paradigm (see Sec. 2.1). It uses a transformer-decoder to predict actions conditioned on past states, past actions, and expected return-to-go (the sum of the future rewards). The parameters are optimized by minimizing the cross-entropy (discrete) or mean square error (continuous) loss between the predicted and actual actions.

DT uses the GPT architecture to address the credit-assignment problem; the self-attention mechanism can associate rewards with the corresponding state-action pairs across long time intervals. This also allows the DT policy to learn effectively even in the presence of distracting rewards (Hung et al., 2018). Empirical experiments demonstrate that the DT outperforms state-of-the-art model-free offline approaches on offline datasets such as Atari and Key-to-Door tasks.

The DT is a model-free approach that predicts actions based on past trajectories without forecasting new states, so it can’t plan future actions. This limitation is addressed by the trajectory transformer (TT) (Janner et al., 2021), an MBRL approach that formulates RL as a conditional sequence modeling problem. TT models past states, actions, and rewards to predict future actions, states, and rewards effectively. Using rewards as inputs prevent myopic behavior and enable the agent to plan future actions through search methods like beam-search (Negrinho et al., 2018).

This task-specific conditioning of agents offers flexibility in learning complex tasks. Prompt-based DT (Xu et al., 2022b), enables few-shot adaptation in offline RL. The input trajectory, which acts as a prompt, contains segments of few-shot demonstrations, encoding task-specific information to guide policy generation. This approach allows the agent to exploit offline trajectories collected from different tasks and adapt to new scenarios for generalizing to unseen tasks. Similarly, the text decision transformer (TDT) (Putterman et al., 2021) employs natural language signals to guide policy-based language instruction in the Atari-Frostbite environment.

However, DTs face several challenges. They struggle to learn effectively from sub-optimal trajectories. In a stochastic environment, their performance tends to degrade since the action taken may have been sub-optimal, and the achieved outcomes are merely a result of random environment transition. Insufficient distribution coverage of the environment is another challenge in offline RL approaches like DT. To overcome these challenges, solutions such as $Q$-learning decision transformer (QDT) (Yamagata et al., 2022) re-label the return-to-go using a more accurate learned $Q$-function. Environment-stochasticity-independent representations (ESPER) (Paster et al., 2022) addresses stochastic performance degradation by conditioning on average return. Additionally, bootstrapped transformer (BooT) (Wang et al., 2022e), incorporates bootstrapping to generate more offline data. By adopting these approaches, the learning capabilities of DTs can be improved, enabling more effective and robust policies in various scenarios.

6.2. Online RL with Transformers

Transformers have also been applied to online RL, where the agent interacts with the environment while learning. In realistic environments, issues such as noisy sensors, occluded images, or unknown agents introduce the problem of partial observability. This makes it difficult for agents to choose the correct action (Sec. 2.2). Here, retaining recent observations in memory is crucial to help disambiguate the true state. Traditionally, this problem has been approached using RNNs, but transformers can provide better alternatives.

Figure 10. The deep transformer $Q$-network (Esslinger et al., 2022) employs a transformer decoder to learn a policy $Q\sim \pi (𝐒)$ that maps a given state to its corresponding $Q$ value for each action. During training, the network predicts $Q$ values for all the past observations, enabling policy updates via Bellman error, which measures the discrepancy between the current and updated value estimates obtained through the TD error. The network uses the current state’s $Q$ value to predict the optimal action at inference.

The deep transformer $Q$-network (DTQN) (Esslinger et al., 2022) addresses the challenge of partially observable environments using a transformer decoder architecture. At each time step of training, it receives the agent’s previous $k$ observations and generates $k$ sets of $Q$-values. This unique training strategy encourages the network to predict $Q$-values even in contexts with incomplete information, leading to developing a more robust agent. During the evaluation, it selects the action with the highest $Q$-values from the last time-step in its history (Fig. 10).

The DTQN incorporates a learned positional encoding, which enables the network to adapt to different domains by learning domain-specific temporal dependencies. This domain-specific encoding matches the temporal dependencies of each environment and allows the DTQN to adapt to environments with varying levels of temporal sensitivity. The DTQN demonstrates superior learning speed and outperforms previous recurrent approaches in various partially observable domains, including gym-gridverse, car flag, and memory cards (Morad et al., 2023).

6.3. Transformers for Multi-Agent Reinforcement Learning

MARL (Sec. 2.1) presents unique challenges as agents learn and adapt their behaviors through interactions with other agents and the environment. One such challenge stems from the model architecture’s fixed input and output dimensions, which means that different tasks must be trained independently from scratch (Shao et al., 2019b; Wang et al., 2020b). Consequently, zero-shot transfer across tasks is limited.

Another challenge arises from the failure to disentangle observations from different agents (Hu et al., 2021). When all information from various agents or environments is treated equally, it can result in misguided decisions by individual agents. This challenge becomes particularly prominent when utilizing a centralized value function, which serves as a shared estimate of actions and state value across multiple agents to guide their behavior (Chen and Tan, 2023). As a result, appropriately assigning credit to individual agents becomes difficult.

Figure 11. The multi-agent transformer (MAT) (Wen et al., 2022) employs an encoder-decoder architecture. The encoder receives a series of observations from the agents and converts them into latent representations. These representations are then fed into the decoder. The decoder generates each agent’s optimal action step-by-step sequentially and auto-regressively. To ensure proper training, the masked attention blocks restrict agents to only access actions from preceding agents.

The universal policy decoupling transformer (UPDeT) (Hu et al., 2021) is designed to handle challenges in tasks with varying observation and action configuration requirements. It achieves this by separating the action space into multiple action groups, effectively matching related observations with corresponding action groups. UPDeT improves the decision-making process by employing a self-attention mechanism and optimizing the policy at the action-group level. This enhances the explainability of decision-making while allowing for high transfer capability to new tasks.

This characteristic is also observed in the multi-agent transformer (MAT) (Wen et al., 2022). MAT (Fig. 11) transforms the joint policy search problem into a sequential decision-making process, allowing for parallel learning of agents’ policies regardless of the number of agents involved. The encoder utilizes the self-attention mechanism to process a sequence of each agent’s observations, capturing their interactions. This generates a sequence of latent representations that are then fed into the decoder. The decoder, in turn, produces each agent’s optimal action in an auto-regressive and sequential manner. As a result, MAT possesses robust generalization capabilities, surpassing multi-agent proximal policy optimization (MAPPO) (Lohse et al., 2021), and heterogeneous-agent proximal policy optimization (HAPPO) (Kuba et al., 2022) in few-shot experiments on multi-agent MuJoCo tasks.

However, MARL faces limitations in real-world applications due to the curse of many agents (Wang et al., 2020a), which stems from the exponentially growing state-action space as the number of agents increases. This presents challenges in learning the value functions and policies of the agents, leading to inefficient relational reasoning among them and credit-assignment problems. Concatenating the state-action spaces of individual agents and treating them as a single-agent problem leads to exponential state and time complexity (Zhou et al., 2020). Additionally, independent learning of policies may struggle to converge without cooperation (Gupta et al., 2021).

TransMix (Khan et al., 2022) tackles the challenge through a centralized learning approach, enabling agents to exchange information during training. During policy execution, each agent relies on a partially observable map. The action space in the star-craft multi-agent challenge (SMAC) (Vinyals et al., 2017) encompasses various actions, including moving units, attacking enemies, gathering resources, constructing buildings, and issuing commands to control the game state. Utilizing transformers, TransMix captures global and local contextual interactions among agent $Q$-values, histories, and global state information, facilitating efficient credit assignment.

The transformer’s ability to reason about relationships between agents improves results in both model-free MARL (regardless of the number of agents) and model-based MARL (with a logarithmic dependence on the number of agents) (Guedj, 2019). Notably modeling the transformer’s self-attention with other neural network types requires an impractically large number of trainable parameters, highlighting the significance of self-attention in capturing agent interactions (Guedj, 2019). Moreover, the transformer’s performance remains stable across different agent counts, with accuracy impacted by neural network depth (Guedj, 2019), making it highly efficient for MARL.

7. Training Strategy

Training transformers poses challenges due to their reliance on residual branches, which amplify minor parameter perturbations, disrupting model output (Liu et al., 2020); specialized optimizers and weight initializers are needed for successful training. Likewise, the training of RL policies can be unstable (Nikishin et al., 2018) and require distinct strategies for achieving optimal performance. Hence, the integration of transformers into RL is particularly challenging. These challenges can manifest as sudden or extreme changes in performance during training, impeding effective learning and generalization.

The standard transformer architecture is difficult to optimize using RL objectives and needs extensive hyper-parameter tuning, which is time-consuming. Here, we review strategies for training transformers in RL. These include pre-training and transfer learning to expedite learning, improved weight initialization to mitigate gradient issues, and efficient layer utilization for capturing relevant information.

7.1. Pre-Training and Transfer Learning

Transformers can be pre-trained on large, reward-free datasets, providing opportunities to fine-tune when only small annotated datasets are available. Meng et al. (2021) propose using DT to pre-train agents on large, reward-free offline datasets of prior interactions. During pre-training, reward tokens are masked, allowing the transformers to learn to predict actions based on the previous state and action content while extracting behavior from the dataset. This pre-trained model can then be fine-tuned with a small, reward-annotated dataset to learn the skills necessary to achieve the desired behavior based on the reward function.

Transfer learning is challenging when the environment dynamics change. A training method for DT (Boustati et al., 2021) addresses this challenge by using counterfactual reasoning. It generates counterfactual trajectories in an alternative environment, which are used to train a more adaptable learning agent. This process aids in regularizing the agent’s internal representation of the environment, enhancing its adaptability to structural changes. Moreover, unsupervised pre-training of vision and sequence encoders has also improved downstream few-shot learning performance (Putterman et al., 2021). By leveraging pre-trained models, the agent can quickly adapt to new, unseen environments and achieve higher performance with limited training data.

7.2. Stabilizing Training

In the RL setting, transformer models require learning rate warmup to prevent divergence caused by backpropagation through the layer normalization modules, which can destabilize optimization. To enhance stability, Melo (2022) proposes to use T-Fixup initialization (Huang et al., 2020b). This applies Xavier initialization (Glorot and Bengio, 2010) to all parameters except input embeddings, eliminating the need for learning rate warmup and layer normalization. It is crucial in environments where learned behavior guides exploration; it addresses instability during early training stages when policies are more exploratory and prevent convergence to sub-optimal policies.

The gated transformer-XL (GTrXL) architecture (Parisotto et al., 2020) has demonstrated promising results in stabilizing RL training and improving performance. It improves upon the original TrXL architecture by applying layer normalization exclusively to the input stream within the residual model rather than the shortcut stream. This modification allows the initial input to propagate through all the layers, promoting training stability. GTrXL replaces the residual connection with a gated recurrent unit (GRU)-style gating mechanism. This gating mechanism regulates information flow through the network controlling the amount of information passed via the shortcut. This added flexibility enhances the model’s adaptability to RL scenarios and facilitates stable training.

8. Interpretability

Interpretability of the learned RL policies is desirable in safety-critical applications like healthcare and autonomous driving (Glanois et al., 2021). This helps in building trust, facilitating debugging, and promoting ethical and fair decision-making. However, achieving interpretability has been a significant challenge and a bottleneck in the progress of RL (Milani et al., 2022; Heuillet et al., 2021).

One way to interpret transformers is to visualize the attention weights using heatmaps (Zhang et al., 2022c). This helps to understand which features are used to learn the particular task. In multi-agent scenarios, these visualizations reveal the localized areas of the input space where the individual agents focus their attention, facilitating coordinated and cooperative behavior. For instance, Motokawa and Sugawara (2021) introduce a multi-agent transformer deep $Q$-network (MAT-DQN) that integrates transformers into a deep $Q$-network. Using heatmaps, MAT-DQN provides insights into the important input information that influences the agent’s decision-making process for cooperative behavior.

Analyzing attention heatmaps unveils the agent’s ability to consider other agents, relevant objects, and pertinent tasks, allowing for a clear interpretation of the policy. Such visualization is critical in sparse reward settings, where understanding which past state had the most influence on decision-making is crucial. Attention-augmented memory (AAM) (Qu et al., 2022) exemplifies this by combining the current observation with memory. This enables the agent to understand “what” the agent observes in the current environment and “where” it directs its attention in its memory.

An interesting method for enhancing interpretability involves the use of transformers in neuro-symbolic policies (Bastani et al., 2020). Neuro-symbolic policies combine programs and neural networks to improve interpretability and flexibility in RL tasks. Specifically, a neuro-symbolic transformer is a variant of the traditional transformer model that incorporates programmatic policies into the attention mechanism. Instead of utilizing a neural network, the attention layer employs a program to determine the relevant inputs to focus on. These programmatic policies can take various forms, including decision trees, rule lists, and state machines. This approach improves interpretability by providing a more precise understanding and visualization of why agents attend to specific inputs.

However, it has been demonstrated that attention weights alone are unreliable predictors of the importance of intermediate components in NLP (Serrano and Smith, 2019; Bai et al., 2021), leading to inaccurate explanations of model decisions; learned attention weights often highlight less meaningful tokens and exhibit minimal correlation with other feature importance indicators like gradient-based measures. Furthermore, relying solely on attention weights can result in fragmented explanations that overlook most other computations. Recent work has introduced the assignment of a local relevancy score (Chefer et al., 2021). These scores are propagated through layers to achieve class-based separation and enhance the interpretability of the transformers. This approach holds promise for future research to improve the interpretability of RL policies.

9. Applications

RL has traditionally been constrained to unrealistic scenarios in virtual environments. However, with modern deep neural network architectures, there has been a notable shift towards employing RL to address a broader range of practical challenges. The following section describes real-world applications where RL powered by transformers can make a substantial impact.

9.1. Robotics

In robotics, autonomous agents automate complex real-world tasks; a classic example is autonomous driving. Here, learning the RL policy for trajectory planning is essential: it involves forecasting the future positions of one or more agents in an environment while considering contextual information. This requires adequate planning and coordination among agents by modeling their spatial and temporal interactions.

Several studies have proposed to use transformers for processing sequences of high-dimensional scene observations for predicting actions. A recent study (Kargar and Kyrki, 2022) uses ViT to extract spatial representations from a birds-eye view of the ego vehicle to learn driving policies. Compared with CNNs, ViTs are more effective in capturing the global context of the scene. The attention mechanism used in ViTs allow the policy to discern the neighboring cars that are pivotal in the decision-making process of the ego vehicle. As a result, the ViT-based DQN agent outperforms its CNN-based counterparts. Liu et al. (2022b) introduce a transformer architecture to encode heterogeneous information, including the historical state of the ego vehicle and candidate route waypoints, into the scene representation. This approach enhances sample efficiency and results in more diverse and successful driving behaviors during inference. The object memory transformer (Fukushima et al., 2022) explores how long-term histories and first-person views can enhance navigation performance in object navigation tasks. An object scene memory stores long-term scene and object semantics, focusing attention on the most salient event in past observations. The results indicate that incorporating long-term object histories with temporal encoding significantly enhances prediction performance.

Transformers also excel in capturing both spatial relationships and intra-agent interactions, making them ideal for facilitating cooperative exploration and developing intelligent embodied agents. Multi-agent active neural SLAM (MAANS) (Yu et al., 2022) addresses the challenge of cooperative multi-agent exploration (Oroojlooy and Hajinezhad, 2023), where multiple agents collaborate to explore unknown spatial regions. This approach extends the single-agent active neural SLAM (Chaplot et al., 2020) method to the multi-agent setting and utilizes a multi-agent spatial planner with a self-attention-based architecture known as the Spatial-TeamFormer. This hierarchically integrates intra-agent interactions and spatial relationships, employing two layers: An individual spatial encoder that captures spatial features for each agent, and a team relational encoder for reasoning about interactions among agents. To focus on spatial information, the intra-agent self-attention performs spatial self-attention over each agent’s spatial map independently. The team relation encoder focuses on capturing team-wise interactions without leveraging spatial information. This allows MAANS to outperform planning-based competitors in a photo-realistic environment, as shown in experiments on Habitat (Savva et al., 2019).

9.2. Medicine

RL has the potential to assist clinicians; tasks involving diagnosis, report generation, and drug discovery can be considered sequential decision-making problems (Yu et al., 2023).

Disease Diagnosis. Diagnosing a medical condition involves modeling a patient’s information (e.g., treatment history, present signs, and symptoms) to accurately understand the disease. Chen et al. (2022b) propose a model for disease diagnosis called the DxFormer. This employs a decoder-encoder transformer architecture, where the decoder inquires about implicit symptoms. At the same time, the encoder is responsible for disease diagnosis, which models the input sequence of symptoms as a sequence classification task. To facilitate symptom inquiry, the decoder is formulated as an agent that interacts with a patient simulator in a serialized manner, generating possible symptom tokens that may co-occur with prior known symptoms and inquiring about them. The inquiry process proceeds until the confidence level in the predicted disease surpasses a selected threshold, thus enabling a more accurate and reliable diagnosis.

Clinical Report Generation. RL can generate medical reports from images by employing appropriate evaluation metrics such as human evaluations or consensus-based image description evaluation (CIDEr) (Vedantam et al., 2015) and BLEU metrics as rewards. Previous approaches to medical image captioning were constrained by their reliance on RNNs for text generation, which often resulted in slow performance and incoherent reports, as highlighted by Xiong et al. (2019). To address this limitation, their work introduces an RL approach based on transformers for medical image captioning. Initially, a pre-trained CNN is employed to identify the region of interest in chest X-ray images. Then, a transformer encoder is utilized to extract the visual features from the identified regions. These features serve as input to the decoder, which generates sentences describing the X-ray scans. Similarly Miura et al. (2021) used a meshed-memory transformer (M2 Trans) (Cornia et al., 2020) that generates radiology reports, proving more effective than traditional RNN and transformer models. M2 Trans incorporates a CNN to extract image regions. These regions are then encoded using a memory-augmented attention process. This involves assigning attention weights to the image based on prior knowledge stored in memory matrices that capture relationships between different regions. This model is trained using rewards, aiming to enhance generated reports’ factual completeness and consistency.

Drug Discovery. RL has the potential to accelerate drug discovery efforts. It has been utilized to bias or fine-tune generative models, enabling the generation of molecules with more desirable characteristics, such as bioactivity. Traditional generative models for molecules, such as RNNs or generative adversarial networks (GANs) (Goodfellow et al., 2014) have limitations in satisfying specific constraints, such as synthesizability or desirable physical properties. Recent research (Wang et al., 2021a; Li et al., 2022b; Yang et al., 2021; Liu et al., 2023) uses transformers as generative models for molecular generation. These approaches generate better plausible molecules with rich semantic features. A discriminator grants rewards that guide the policy update of the generator. These works demonstrate that transformer-based methods significantly improve capturing and utilizing structure-property relations, leading to higher structural diversity and a broader range of scaffold types for the generated molecules.

9.3. Language Modeling

Language modeling involves understanding the sequential context of language to perform diverse tasks like recognition, interpretation, or retrieval. Large language models like GPT leverage pre-training on vast corpora, enabling them to generate fluent, natural language by sampling from the learned distribution, thus minimizing the need for extensive domain-specific knowledge engineering. However, these models face challenges in maintaining task coherence and goal-directedness. Alabdulkarim et al. (2021) use proximal policy optimization (PPO) to fine-tune an existing transformer-based language model specifically for story generation to address this issue. This approach inputs a text prompt and generates a story based on the provided goal. This policy is updated using a reward mechanism that considers the proximity of the generated story to the desired input goal and the frequency of verb occurrence in the story compared to the goal.

Several studies use additive learning to benefit from pre-trained language models with limited data, incorporating a task-specific adapter over the frozen pre-trained language model. Jo et al. (2022) use RL to selectively sample tokens between the general pre-trained language model and the task-specific adapter. The authors argue that this enables the adapter to focus solely on the task-relevant component of the output sequence, making the model more robust to over-fitting. Cohen et al. (2022) introduce a conversational bot powered by RL where pre-trained models encode conversation history. Given that the action space for dialogue systems can be very large, the authors propose limiting the action space to a small set of generated candidate actions at each conversation turn. They use $Q$-Learning-based RL to allow a dynamic action space at each stage of the conversation.

Increasing the size of language models alone does not necessarily mitigate the risk of toxic biases in the training data. Several RL-based approaches have been proposed to better align these models with the user’s intended objectives. To align GPT-3 to the user’s preferred intentions, Ouyang et al. (2022) introduce InstructGPT. First, the authors propose to collect a set of human-written demonstrations of desired output behavior and fine-tune GPT-3 with supervised learning. Next, a reward model is trained on model outputs ranked from best to worst. Using this reward model, the model is further optimized with RL using PPO. Results demonstrate that InstructGPT, with 1.3B parameters, produces preferable outputs than much larger models such as GPT-3 with 175B parameters. Faal et al. (2023) propose an alternative method to mitigate toxicity in language models via fine-tuning with PPO. They use a reward model based on multi-task learning to mitigate unintended bias in toxicity prediction related to various social identities.

9.4. Edge and Cloud Computing

RL is a valuable tool for optimizing the performance of decision-making systems that require real-time adaptation to changing conditions, such as those used in edge and cloud computing. In edge computing, RL can optimize resource-constrained internet of things (IoT) devices’ performance (Chen et al., 2021b). In cloud computing, RL can be used to optimize resource allocation and scheduling in large-scale distributed systems (Gondhi and Gupta, 2017). Integrating transformers with RL in these two settings can be particularly useful as they can handle high-dimensional sensory states (Ho et al., 2019) and sequences of symbolic states (Bhattamishra et al., 2020).

A distributed deep RL algorithm proposed by Wang et al. (2022d) utilizes transformers to model the policy for optimizing offloading strategies in vehicular networks. These networks enable vehicle-to-vehicle communication. To represent the input sub-task priorities and dependencies, a directed acyclic graph (DAG) is used. Thereafter the attention mechanism employed by the transformer allows for the efficient extraction of state information from this DAG-based topology representation. This facilitates informed offloading decisions. The reward function used in this algorithm optimizes for latency and energy consumption providing valuable feedback. This approach enables faster convergence of the vehicular agent to equilibrium.

9.5. Combinatorial Optimization

Combinatorial optimization involves finding the values of a set of discrete parameters that minimize cost functions (Mazyavkina et al., 2021). Recently, transformer-based models have shown promise in combinatorial optimization (e.g., for the traveling salesman and routing problems) due to their ability to handle sequential data and model complex relationships between entities.

Travelling Salesman. This problem is a classical combinatorial optimization problem commonly found in crew scheduling applications. It has been formulated as an RL problem by Smith (2022). The problem is NP-hard and not NP-complete. The high polynomial complexities of brute-force algorithms necessitate the development of faster methods. This study uses a decision transformer (DT) to solve this challenge. The DT is fed random walks as input and aims to find the optimal path among all nodes. The DT has the advantage of scaling with pseudo-linear time, as it only needs to predict once per node in the route. This significantly improves over previous methods, such as dynamic programming and simulated annealing, which have polynomial and exponential complexity. However, the DT could not always accurately model the travelling salesman problem leading to inconsistent performance.

Routing. Identifying the most efficient route between two nodes in a graph is important in industries such as transportation, logistics, and networking (Mor and Speranza, 2022). Traditional heuristic-based algorithms may not always yield the optimal solution, as adapting to changing conditions is challenging (Wu et al., 2022). GNNs have been used to address these challenges (Lu et al., 2020). However, these may not be sufficient for handling data with complex inter-relationships and structures, motivating the use of transformers. In addition, routing often requires optimizing for multiple constraints, such as cost, time, or distance, which can be tackled using RL. A transformer-based policy has been proposed by Wang and Chen (2022) to tackle the routing problem using a standard transformer encoder with positional encoding, which ensures translation invariance of the input nodes. A GNN layer is used in the decoder, enabling consideration of the graph’s topological structure formed by the node relationships. The policy is then trained using the REINFORCE algorithm (Williams, 1992). This approach improves learning efficiency and optimization accuracy compared to traditional methods while providing better generalization in new scenarios.

9.6. Environmental Sciences

RL algorithms can help address climate change by optimizing the behavior of systems and technologies for reducing greenhouse gas emissions and mitigating the impacts of climate change (Strnad et al., 2019). These algorithms can learn and adapt to multiple constraints, optimizing performance without compromising productivity. However, in such settings, RL algorithms must rely on past context stored in memory, and integrating prior knowledge is crucial for their success, suggesting the use of transformers.

Nasir and Durlofsky (2023) formulate the problem of closed-loop reservoir management problem as a partially observable Markov decision process (POMDP). In subsurface flow settings, such as oil reservoirs, the goal is to extract as much oil as possible while minimizing costs and environmental impact. However, this requires making decisions about well pressure settings often complicated by geological model uncertainties. This work models the RL policy with PPO using temporal convolution and gated transformer blocks for an efficient and effective representation of the state. The framework’s training is accomplished with data generated from flow simulations across an ensemble of prior geological models. After appropriate training, the policy instantaneously maps flow data observed at wells to optimal pressure. This approach helps reduce computational costs and improve decision-making in subsurface flow settings.

Wang et al. (2022b) introduce transformer-based multi-agent actor-critic framework (T-MAAC) leveraging MARL algorithms to stabilize voltage in power distribution networks. This framework recognizes the need for coordination among multiple units in the grid to handle the rapid changes in power systems resulting from the increased integration of renewable energy and tackles this problem by using MARL algorithms. The proposed approach introduces a transformer-based actor that takes the grid state representation as input and outputs the maximum reactive power ratio that each agent in the power distribution can generate. Subsequently, the critic approximates the global $Q$-values using the self-attention mechanism to model the correlation between agents across the entire grid. The policy is reinforced through feedback in the form of reward, aiming to control voltage within a safe range while minimizing power loss in the distribution network. This approach consistently enhances the effectiveness of the active voltage control task.

9.7. Scheduling

The scheduling problem involves determining the optimal arrangement of tasks or events within a specified time frame while considering constraints such as resource availability or dependencies between tasks (Allahverdi, 2016). This problem can arise in various contexts, such as scheduling jobs in manufacturing or optimizing computer resource usage, and can be approached using various techniques (Parmentier and T’kindt, 2023). Transformers are now being used to solve scheduling problems.

The job-shop scheduling problem (JSSP) is a classical NP-hard problem that involves scheduling a set of jobs on a set of machines, where each job has to be processed on each machine exactly once, subject to various constraints. An RL approach for solving the JSSP using the disjunctive graph embedded recurrent decoding (DGERD) transformer is proposed by Chen et al. (2023). This work uses the attention mechanism and disjunctive graph embedding to model the JSSP, which allows complex relationships between jobs and machines to be captured. In the context of JSSP, the attention module learns to prioritize certain jobs or machines based on their importance or availability. By doing so, it can generate more efficient and robust schedules. The disjunctive graph embedding converts the JSSP instance into a graph representation to capture the structural properties, enabling better generalization and reducing over-fitting. This acts as an input for the DGERD transformer consisting of a parallel-computing encoder and a recurrent-computing decoder. The encoder takes the disjunctive graph embedding of the JSSP instance and generates a set of hidden representations that capture the relevant features of the input. This hidden representation is then fed to the decoder to generate an output schedule sequentially. The policy is optimized using feedback from the environment in the form of makespan (length of time that elapses from the start of work to the end) and tardiness penalties. This helps in generating schedules that are both fast and reliable.

9.8. Trading

Stock portfolio optimization involves choosing the optimal combination of assets to obtain the highest possible returns while minimizing risk (Hieu, 2020). This process can be challenging due to the various factors that can impact a portfolio’s performance (Haugh and Lo, 2001), which include market conditions, economic events, and changes in the value of individual stocks. Various techniques can be used to optimize a portfolio, including modern portfolio theory and optimization algorithms (Thakkar and Chaudhari, 2021), and RL is one such approach to automate the trading process. For this application, RL involves training a model to make trading decisions based on historical data and market conditions to maximize the portfolio’s return over time.

Although past performance may not indicate future results, data-driven approaches rely on past features to model a particular stock’s expected future performance. This is because various historical data points, such as price trends, trading volumes, and market sentiment, may hint at a stock’s future performance (Milosevic, 2016). In portfolio optimization, it is necessary to consider both short-term and long-term trends (Ta et al., 2020). The transformer architecture is well-suited for this task.

The first application of transformers in portfolio selection, as introduced by Xu et al. (2020), involved using the relation-aware transformer (RAT). This uses an encoder-decoder transformer architecture to model the RL policy. The encoder takes in the sequential price series of assets, such as stocks and cryptocurrencies, as the input state. It performs sequential feature extraction, comprising a sequential attention layer for capturing patterns in asset prices and a relation attention layer for capturing correlations among assets. The decoder has a network resembling an encoder, with an additional decision-making layer that incorporates leverage and enables accurate decisions, including short sales, for each asset. The final action is determined by combining the initial portfolio vector, the short sale vector, and the reinvestment vector. The agent then receives reward-based feedback, measured as the log return of portfolios. To evaluate the proposed method, real-world cryptocurrency, and stock datasets are used and compared against state-of-art portfolio selection methods. The results demonstrate a significant improvement over existing approaches.

9.9. Hyper-Parameter Optimization

Hyper-parameter optimization (HPO) involves finding the optimal set of hyper-parameters for training a machine learning model. Some commonly used hyper-parameters in machine learning models include the learning rate, batch size, number of hidden units in a neural network, and activation function. As such, finding the best combination of these hyper-parameters can be challenging for large models due to their correspondingly large search space (Ali et al., 2023). Manually setting hyper-parameter values is fast but requires expertise and domain knowledge (Shawki et al., 2021). Automated techniques like random search, grid search, or Bayesian optimization can automatically find ideal hyper-parameter combinations (Bergstra and Bengio, 2012; Snoek et al., 2015), but minimizing overall computational costs remains a challenge. Such auto-tuners rarely perform well for complex tasks and are prone to errors with increased model complexity (Shawki et al., 2021).

Attention and memory enhancement (AME) (Xu et al., 2022a), is a transformer-based search algorithm to enhance the selection of hyper-parameters that tackles these challenges. AME utilizes RL and addresses HPO without relying on distribution assumptions. The agent, or the searcher, is modeled using GTrXL and learns a series of state-to-action mappings based on rewards. In this context, the state refers to the combination of evaluated configurations, while an action corresponds to the new configuration chosen by the agent from the search space. Utilizing GTrXL improves the ability to capture relationships among different configurations through memory mechanisms and multi-head attention, thereby enabling attentive sampling. The agent is trained using feedback in the form of rewards, which promotes the generation of high-performance configurations and penalizes those leading to reduced performance. Consequently, it effectively locates high-performance configurations within vast search space. Results demonstrate that the AME algorithm surpasses other HPOs like Bayesian optimization, evolutionary algorithms, and random search methods in terms of adaptability to diverse tasks, efficiency, and stability.

10. Limitations

As discussed above, transformers are gradually being integrated into RL for various applications. Despite these advances, some limitations impede their widespread use. This section details these limitations and provides insights for future research.

Balancing Local & Global Context. In RL, global contextual information is required for efficient high-level planning (Barto and Mahadevan, 2003). This information is combined with additional nearby details, known as the local context, to predict low-level actions precisely. As detailed by Li et al. (2019); Wang et al. (2021b), transformers may not be as effective as other models in capturing local context. This limitation is mainly because of the self-attention mechanism, which compares queries and keys for all elements in a sequence using the dot product. This point-wise comparison does not directly consider the local context for each sequence position, which may lead to confusion due to noisy local points. Recent studies (Lin et al., 2023; Wang et al., 2022a, 2021c) inspired by CNNs have proposed modifications to the original attention mechanism to balance the local and global context more effectively. These approaches include local window-based boundary-aware attention, allowing the model to focus on a small window of nearby details and the global context when making predictions.

Weak Inductive Bias. CNNs and long short-term memory (LSTM) (Hochreiter and Schmidhuber, 1997) models have a strong inductive bias toward the dataset’s structure, which helps narrow the search space and leads to faster training (d’Ascoli et al., 2021). This makes them better suited for situations with less training data. However, transformers have a relatively weak inductive bias, making them more capable of finding general solutions (Hessel et al., 2019) but more susceptible to over-fitting, especially when less data is available. This limitation can be a significant challenge in RL, where training a policy already requires millions of trajectories. Furthermore, learning models like the decision transformer require collecting trajectories from learned policies which can be even more challenging. One approach to counter weaker inductive bias in transformers is to use foundation models (Zhou et al., 2023; Moor et al., 2023). Foundation models are pre-trained on large and diverse datasets, which allows them to learn general patterns that can be applied to a wide range of downstream tasks. The foundation model can achieve state-of-the-art results with less data by fine-tuning the pre-trained model on a smaller task-specific dataset.

Quadratic Complexity. The self-attention mechanism of transformers becomes more computationally expensive as input sequence length increases due to the quadratic increase in pairwise comparisons between tokens (Keles et al., 2023). This limitation, along with hardware and model size constraints, restricts the ability of transformers to process longer input sequences, making them unsuitable for specific tasks that require substantial amounts of contextual information, like document summarization or genome fragment classification. This limitation can also pose challenges in RL for applications requiring extended temporal modeling. However, recent works (Katharopoulos et al., 2020; Lu et al., 2021; Ren et al., 2021) have provided methods to reduce this cost to linear or sub-quadratic, providing new possibilities for using transformers in applications that require longer input sequences.

11. Conclusion

This survey explored the diverse uses of transformers in RL, including representation learning, reward modeling, transition function modeling, and policy learning. While the original transformer architecture has limitations, it can be modified for many RL applications. We showcased the advances in transformers that have broadened the scope of RL to real-world problems in robotics, drug discovery, stock trading, and cloud computing. Finally, we discussed the current limitations of transformers in RL and ongoing research in this field. Given its versatility in addressing challenges such as partial observability, credit assignment, interpretability, and unstable training — issues commonly encountered in traditional RL— we anticipate that the transformer architecture will continue to gain popularity in the RL domain.

Acknowledgement. We thank CIFAR, Google, CMLabs for funding the project, and Vincent Michalski for the valuable feedback.