Rethinking Transformers in Solving POMDPs

Reinforcement Learning (RL) in the real world confronts the challenge of incomplete information (Dulac-Arnold et al., 2019) due to partial observability, necessitating decision-making based on historical data. The design of RL algorithms under partial observability, denoted as Partially Observable RL (Kaelbling et al., 1998; Littman & Sutton, 2001; Li et al., 2015), typically employs a hierarchical structure combining $(\texttt{SEQ},\texttt{RL})$. This structure involves firstly feeding the history into a sequence model SEQ, such as Recurrent Neural Network (RNN) (Elman, 1990) or Long Short-Term Memory (LSTM) (Hochreiter & Schmidhuber, 1997), yielding a hidden state containing past information, then processing it using existing RL algorithms.

Regarding the sequence model, Transformer (Vaswani et al., 2017), renowned for its achievements in the natural language processing (NLP) domain (Radford et al., 2019; Brown et al., 2020; OpenAI, 2023), stands out as a prominent candidate. Transformers have shown a strong ability to handle contexts in a non-recurrent manner. Compared with their recurrent counterpart like RNNs, the advantages of Transformers as a sequence model shine in several aspects: 1) long-term memory capacity, as opposed to RNNs with rapid memory decay (Ni et al., 2023; Parisotto et al., 2019); 2) effective representation learning from context for specific tasks (Micheli et al., 2022; Laskin et al., 2022; Lee et al., 2022; Robine et al., 2023), benefiting meta-RL or certain environments (Bellemare et al., 2013); 3) stronger learning ability on large-scale datasets (Baker et al., 2022).

However, deploying Transformers in Partially Observable RL introduces challenges, commonly manifesting as sample inefficiency (Parisotto et al., 2019; Ni et al., 2023). This issue is similarly observed in computer vision (CV) and is attributed to the data amount that Transformers require to learn the problem-specific inductive biases (Dosovitskiy et al., 2021). While it is validated in CV, it remains unknown whether data amount is the key ingredient in decision making. Hence, a natural question arises: Can Transformers effectively solve decision-making problems in POMDPs with sufficient data?

In this work, we investigate this critical question and challenge the conventional wisdom. We argue that Transformers cannot solve POMDP even with massive data. This stance is inspired by a key observation: While most RNNs are complete for regular languages, Transformers falter to model them (Delétang et al., 2023; Hahn, 2020b) . A notable example is their struggle with tasks like PARITY, which is to determine the parity of the occurrence of “1” in a binary string. We hypothesize that, in POMDPs, this limitation becomes pronounced due to the close relationship between regular languages and POMDPs.

To elaborate further, regular languages exhibit a direct correspondence with Hidden Markov Models (HMMs) (Carrasco & Oncina, 1994), and POMDPs can be regarded as HMMs augmented with an incorporated decision-making process. We further establish that regular languages can be effectively reduced to POMDPs. From the computational complexity perspective, the parallel structure of the Transformer makes it equivalent to a constant-depth computation circuit. Some regular languages fall outside of this complexity class, making the POMDP problems derived from them harder, and Transformer would struggle to solve them. This is demonstrated both theoretically and empirically in this study.

To alleviate the limitations of Transformers caused by the parallel structure, we propose to introduce a pointwise recurrent structure. Upon reviewing current variants of sequence models with such a structure, we find that they can be broadly generalized as linear RNNs. Based on extensive experiments over a range of sequence models over POMDP tasks with diverse requirements, we highlight LRU (Orvieto et al., 2023) as a linear rnn model well-suited for Partially Observable RL. Our contributions are three-fold:

• •

  We demonstrate the theoretical limitations of Transformers as sequence model backbones for solving POMDPs, through rigorous analysis.

• •

  To better utilize the inductive bias of the sequence model, we study the advantages of the Transformer and the RNN, and advocate the linear RNN as a better-suited choice for solving POMDPs, taking advantage of both models.

• •

  Through extensive experiments across various tasks, We compare the capabilities exhibited by various sequence models across multiple dimensions. Specifically, we show that Transformers exhibit sub-optimal performance as the sequence model in certain POMDPs, while highlighting the strength of linear RNNs when assessed comprehensively.

Theoretical limitations of Transformers. There is a substantial body of work investigating the theoretical limitations of Transformers from the perspective of computational complexity and formal language. For example, Delétang et al. (2023); Huang et al. (2022) experimentally verifies that RNNs can recognize regular languages, but Transformers are unable to achieve this. Additionally, Hahn (2020a) demonstrates that Transformers are not robust in handling sequence length extrapolation. Moreover, Merrill & Sabharwal (2023); Merrill et al. (2022) point out that, under limited precision, ${\mathsf{TC}^{0}}$ serves as an upper bound for the computational power of Transformers. Applying this result, Feng et al. (2023) illustrates the challenges Transformers face in solving practical problems such as arithmetic operations and linear systems of equations. Currently, works such as Ni et al. (2023); Morad et al. (2023); Deng et al. (2023) discuss the pros and cons of transformers in RL algorithms, with a focus on analyzing the advantages or providing simple evaluations. In contrast, integrating relevant theories from formal languages, we offer a new theoretical perspective on analyzing the limitations of transformers in RL.

Variants of sequence models for handling long contexts. Multiple variants of mainstream sequence models designed to handle long contexts have provided significant inspiration for this work. In the case of RNN-like models, addressing the issue of rapid memory decay has led to the emergence of linear RNNs (Gu et al., 2021; Orvieto et al., 2023), which remove activation functions in the recurrence part. These models have demonstrated excellent performances in benchmarks for long-range modeling (Tay et al., 2020). For Transformers, to tackle limited training length and inefficient inference, current studies emphasize the introduction of recurrence. Recurrence in Transformers can be categorized into two types: 1) chunkwise recurrence, which processes parts exceeding the context length using a recurrent block with minimal alterations to the original parallel structure (Dai et al., 2019; Hutchins et al., 2022); 2) pointwise recurrence, which derives recurrence representations by linearizing attention (Sun et al., 2023; Peng et al., 2023; Schlag et al., 2021; Katharopoulos et al., 2020). We argue that linear RNNs and pointwise recurrence Transformers ultimately converge to a similar solution, incorporating the strength of both approaches and are suitable for solving Partially Observable RL problems.

Applications of sequence models in RL. In recent years, there have been many applications of Transformers in RL, such as Decision Transformer (DT) (Chen et al., 2021), its variants (Yamagata et al., 2023; Wu et al., 2023) in offline RL; GTrXL (Parisotto et al., 2019), Online DT (Zheng et al., 2022) in online RL, and the Transformer State Space Model (Chen et al., 2022) as world models. There are works (Reid et al., 2022; Shi et al., 2023) showing that the inductive bias of pre-trained Transformers could help RL, where the states are fully observable. Hu et al. (2023) uses Transformer to solve a specific partially observable setting, frame drops, whereas demanding additional assumptions on the prior distribution. On the other hand, recent works comparing Transformer-based and RNN-based approaches in Partially Observable RL empirically support our idea that Transformers have weaknesses in partially observable environments (Morad et al., 2023; Deng et al., 2023). In many cases, simpler architectures like RNNs and LSTMs prove to be more effective. For instance, DreamerV3 (Hafner et al., 2023), which adopts GRU (Dey & Salem, 2017) as the backbone of the Recurrent State Space Model, has outperformed previous Transformer-based approaches like VPT (Baker et al., 2022) and IRIS (Micheli et al., 2022). Additionally, there has been a recent line of research on the application of linear RNNs in RL (Irie et al., 2021; Lu et al., 2024; Samsami et al., 2024), which has shown promising results. This paper will incorporate insights into the limitations of Transformers to analyze why this would be a natural choice.

Sequential neural network. Sequential Neural Networks are a type of deep learning model for sequence modeling. Given a input sequence $\quantity{u_{i}}_{i=1}^{n}$, the model learns a hidden state sequence $\quantity{x_{i}}_{i=1}^{n}$ and yields the output sequence $\quantity{y_{i}}_{i=1}^{n}$. There are currently two mainstream methods for computing the hidden state $x_{i}$:

• •

  Recurrent-like: $x_{t}=\sigma(Ax_{t-1}+Bu_{t-1}+c)$σしぐま ( italic_A italic_x start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT + italic_B italic_u start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT + italic_c ) where $\sigma$σしぐま is the activation function;

• •

  Attention-like: $x_{t}=\operatorname{attn}\quantity(W^{Q}\tilde{U},W^{K}\tilde{U},W^{V}\tilde{U})_{t}$. Here $\tilde{U}_{i}=u_{i}+p_{i}$, $p_{i}$ is a position embedding, and $\operatorname{attn}(Q,K,V)$ is defined as $\operatorname{softmax}(QK^{\top}+M)V$, where $M$ is an attention mask.

Current sequence models use multiple layers. The overall structure is illustrated in Figure 1. For recurrent models like RNNs, GRUs, etc., the output of the previous layer is directly used as the input for the next layer., while for Transformers, the pointwise transformations typically take the form of MLPs with skip connections.

POMDP. A Partially Observable Markov Decision Process (POMDP) $\mathcal{M}$ can be defined as $(S,A,T,R,\mathrm{<ruby><rb>\Omega </rb><rt>おめが</rt></ruby>},O,\mathrm{<ruby><rb>\gamma </rb><rt>がんま</rt></ruby>})$Ωおめが𝑂𝛾(S,A,T,R,\Omega,O,\gamma)( italic_S , italic_A , italic_T , italic_R , roman_Ωおめが , italic_O , italic_γがんま ) (Åström, 1965). At time $i$, the agent is in state $s_{i}\in S$, observes $o_{i}\in\Omega\sim O(\cdot|s_{i})$Ωおめが ∼ italic_O ( ⋅ | italic_s start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ), takes action $a_{i}\in A$, receives reward $R(s_{i},a_{i})$ and would transit to $s_{i+1}\sim T(\cdot|s_{i},a_{i})$. The agent’s policy based on the observation history $h_{t}=\quantity{\quantity(o_{i},a_{i},r_{i})}_{i=1}^{t}$ is denoted as $\pi(\cdot|h_{t-1},o_{t})$πぱい ( ⋅ | italic_h start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT , italic_o start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ). We say an algorithm $\mathcal{A}$ that can solve POMDP as being able to find the optimal policy $\pi^{\star}$πぱい start_POSTSUPERSCRIPT ⋆ end_POSTSUPERSCRIPT for given POMDP $\mathcal{M}$ where:

DFA & regular language. Deterministic Finite Automata (DFA) can be defined as $A=(S,\mathrm{<ruby><rb>\Sigma </rb><rt>しぐま</rt></ruby>},T,s_0,F)$Σしぐま𝑇subscript𝑠0𝐹A=(S,\Sigma,T,s_{0},F)italic_A = ( italic_S , roman_Σしぐま , italic_T , italic_s start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT , italic_F ), where $S$ is a finite set of states, $\mathrm{<ruby><rb>\Sigma </rb><rt>しぐま</rt></ruby>}$Σしぐま\Sigmaroman_Σしぐま is a finite set of symbols, $T:S\times \mathrm{<ruby><rb>\Sigma </rb><rt>しぐま</rt></ruby>}\to S$Σしぐま𝑆T:S\times\Sigma\to Sitalic_T : italic_S × roman_Σしぐま → italic_S is the transition function, $s_{0}$ is the start state, and $F\subseteq Q$ is the set of accepting states. A string $w$ whose $i$-th symbol is $w_{i}$ is accepted by $A$ if $\exists(s_{0},s_{1},\ldots,s_{n})$, s.t. $s_{i}\in S$ for $1\leq i\leq n$, $s_{i}=T(s_{i-1},w_{i})$ and $s_{n}\in F$. $L$ is a regular language if it is recognized by DFA $A$, that is, $L=\{w:A\text{ {accepts} }w\}$.

RL algorithms typically take as inputs the current state with the assumption of the Markov property. In partially observable environments that lack the Markov property, the hidden state extracted by SEQ from the observation history is thus anticipated to contain the information of the real state to benefit subsequent RL.

Notably, Delétang et al. (2023) points out that Transformers (referred to as TFs) fail to recognize regular languages. Inspired by the significant correspondence between regular languages and POMDPs (details in Definition 4.2), we naturally conjecture that TF is not capable of retrieving the information of real state from partial observations accurately, which would lead to a decline in the performance of the pipeline $(\texttt{SEQ},\texttt{RL})$.

Building upon this view, this section first shows that solving a POMDP problem is harder than solving a regular language problem. Afterward, we will introduce two theoretical results to elucidate the limitations of Transformers, supported by simple examples for illustrations. Consequently, it is inferred that $(\texttt{TF},\texttt{RL})$ cannot address POMDPs generally.

From the analysis in Section 4, it becomes evident that RNN-like models (LSTM, GRU, RNN) emerge as promising sequence model choices for Partially Observable RL. There has been considerable theoretical work demonstrating their completeness on regular languages (Merrill, 2019; Korsky & Berwick, 2019). For cases of log precision, based on definitions, we can directly map the recurrent units of RNNs to transition functions in DFAs. Therefore, RNNs do not suffer from the theoretical constraints encountered by Transformers.

However, RNN-like models face the challenge of rapid memory decay (Ni et al., 2023; Parisotto et al., 2019), leading to an inferior performance on POMDP problems that demand long-term memory when compared to Transformers (Parisotto et al., 2019; Ni et al., 2023).

Another insight from the previous section is that the attention mechanism of Transformers is primarily to blame for their limitations (see Figure 1b). As articulated in  Merrill & Sabharwal (2023), there exists a trade-off between the highly parallel structure of Transformers and their computational capacity.

To alleviate these limitations of Transformers, a natural idea is to endow Transformers with the ability of pointwise recurrence (see Figure 1a). This line of development has been the focus of numerous efforts, resulting in several Transformer variants that incorporate this mechanism, as detailed in Table 1. The shared feature of these methods in their recurrence representation for a single head can be found in the simple linear operations they employ, such as $x_{t}=\lambda x_{t-1}+u_{t}$λらむだ italic_x start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT + italic_u start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT. While the non-linear components can be amortized across the layers through the FFN between layers. If the number of heads is set equal to the dimension of the hidden state $h$, then

where $\mathbf{\Lambda}=\operatorname{diag}\quantity(\lambda_{1},\ldots,\lambda_{h})$Λらむだ = roman_diag ( start_ARG italic_λらむだ start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT , … , italic_λらむだ start_POSTSUBSCRIPT italic_h end_POSTSUBSCRIPT end_ARG ). In RetNet, operations involving $x_{t}$, $\lambda_{i}$λらむだ start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT, and $u_{i}$ are performed over $\mathbb{C}$, while the remaining operations are carried out over ${\mathbb{R}}$. As for RWKV, $\lambda_{i}$λらむだ start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT is a learnable parameter, while in the rest variants are hyperparameters.

From the perspective of RNN, if we linearize and diagonalize the RNN’s recurrent unit $\mathbf{x}_{t}=\mathbf{A}\mathbf{x}_{t-1}+\mathbf{B}\mathbf{u}_{t}$, we obtain the following form:

where $\mathbf{A}=\mathbf{P}\mathbf{\Lambda}\mathbf{P}^{-1},\mathbf{\tilde{x}}_{t}=\mathbf{P}^{-1}\mathbf{x}_{t},\mathbf{\tilde{u}}_{t}=\mathbf{P}^{-1}\mathbf{B}\mathbf{u}_{t}$Λらむだ bold_P start_POSTSUPERSCRIPT - 1 end_POSTSUPERSCRIPT , start_ID over~ start_ARG bold_x end_ARG end_ID start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT = bold_P start_POSTSUPERSCRIPT - 1 end_POSTSUPERSCRIPT bold_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , start_ID over~ start_ARG bold_u end_ARG end_ID start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT = bold_P start_POSTSUPERSCRIPT - 1 end_POSTSUPERSCRIPT bold_B bold_u start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT. Since almost all matrices can be diagonalized over $\mathbb{C}$, the operations mentioned above are defined in $\mathbb{C}$ (Horn & Johnson, 2012).

Comparing (1) and (2), the pointwise-recurrence Transformer can be viewed as a linear RNN with certain constraints, and the linear RNN serves as a balance point between Transformers and RNNs. To summarize, we expect linear RNN to be more suitable as a sequence model in Partially Observable RL for the following reasons.

Regular language. Many studies suggest that the recurrence with non-linear activation functions plays a crucial role in the completeness of RNNs in regular languages (Chung & Siegelmann, 2021), while linear RNNs may lose this completeness. However, some researches indicate that linear RNNs can effectively approximate RNNs (Huang et al., 2022; Lim et al., 2023) and perform well on formal language tasks similar in form to NLP (Huang et al., 2022; Irie et al., 2023). Compared to Transformers, their inductive biases are closer to HMMs. In subsequent experiments (cf. Section 6), we validate that $(\texttt{LRNN},\texttt{RL})$ can implicitly learn the states in POMDPs.

State space model. Linear RNNs have been proven to efficiently fit partially observable linear dynamic systems (Wang et al., 2022). While the transformer’s fitting capability has theoretical proofs only under certain specific conditions (Balim et al., 2023; Li et al., 2023), with no similar conclusion for more general situations. The linear dynamic system can be considered as a first-order approximation of a state space model, indicating the potential of linear RNNs in addressing a broader range of POMDPs.

Long term memory. The primary reason for the long-term dependency issues in RNNs is the challenge of gradient explosion or vanishing when input length increases during training (Pascanu et al., 2013). Transformers, due to their parallel structure, are less susceptible to this issue. To mitigate this problem, gate mechanisms are introduced to RNNs (Dey & Salem, 2017; Hochreiter & Schmidhuber, 1997). However, Kanai et al. (2017) indicates that the non-linear recurrence is the primary cause of gradient explosions. For linear RNNs, effectively managing the range of parameters $\lambda_{i}$λらむだ start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT between $[0,1]$ during initialization successfully addresses both gradient explosion and vanishing issues. This has been validated in certain supervised learning tasks with long-term dependencies (Orvieto et al., 2023; Gu et al., 2021).

In this section, we compare the effectiveness of three different sequence models — Transformer, RNN, and linear RNN — in addressing partially observable decision-making problems within the realm of Partially Observable RL (SEQ, RL). We choose GPT (Radford et al., 2019), LSTM (Hochreiter & Schmidhuber, 1997), and LRU (Orvieto et al., 2023) as the representative architectures for these three types of models. To substantiate our hypotheses, we conduct experiments in three distinct POMDP scenarios, detailed in Sections 6.1 to 6.3. These experiments are designed to assess the models from various perspectives: 1) POMDPs derived from certain regular languages, including EVEN PAIRS, PARITY, and SYM(5); 2) tasks from Pybullet Partially Observable environments (Ni et al., 2022) that require the ability of state space modeling; 3) tasks that require pure long-term memory capabilities, such as Passive T-Maze and Passive Visual Match (Hung et al., 2018). Comprehensive implementation details, task descriptions, and supplementary results are presented in Appendix D. We also conduct a comparison with some published Transformer in RL there.

In this work, we challenge the suitability of Transformers as sequence models in Partially Observable RL. Through theoretical analysis and empirical evidence, we reveal Transformer’s limitations in solving POMDPs, particularly their struggle with modeling regular languages, a key aspect of POMDPs. As a remedy to these issues, We propose LRU as a more effective alternative, combining the strengths of recurrence and attention. Supported by extensive experiments, our findings challenge the prevailing use of Transformers in sequential decision-making tasks, and open new avenues for exploring recurrent structures in complex, partially observable environments.

It is also important to acknowledge the limitations of our work. After introducing recurrence, LRU serves as a choice to combine the advantages of Transformer and RNN, while still lacking theoretical guarantees for modeling regular languages. Although LRU demonstrates satisfactory performance in experiments, there remains a need for further exploration in this direction. Additionally, the theoretical analysis in this paper focuses more on the exploitation aspect of RL, while lacking discussion on exploration. Complex POMDP tasks not only require suitable sequence models but also need to be paired with appropriate RL algorithms.

Our work revisits the application of Transformers in RL, aiming to advance the development of decision intelligence. If misused in downstream tasks, it has the potential to lead to adverse effects such as privacy breaches and societal harm. Nevertheless, this is not directly related to our research, as our primary focus is on theoretical investigations.