Rethinking Transformers in Solving POMDPs

TL;DR

This work investigates the suitability of Transformer architectures for solving POMDPs, highlighting theoretical limitations that arise from their parallel, non-recurrent structure. By constructing POMDPs from regular languages and analyzing fitting and generalization limits under circuit complexity arguments, the authors demonstrate that Transformers struggle with POMDP-induced inductive biases. They propose the Deep Linear Recurrent Unit (LRU) as a practical, principled alternative that blends recurrence with attention, and provide extensive experiments across regular-language–derived POMDPs, PyBullet occlusion tasks, and long-term memory environments showing LRUs' strong performance relative to Transformers and LSTMs. The findings suggest that incorporating recurrence, even in a linear form, yields substantial benefits for partially observable decision-making, guiding future sequence-model design in RL settings.

Abstract

Sequential decision-making algorithms such as reinforcement learning (RL) in real-world scenarios inevitably face environments with partial observability. This paper scrutinizes the effectiveness of a popular architecture, namely Transformers, in Partially Observable Markov Decision Processes (POMDPs) and reveals its theoretical limitations. We establish that regular languages, which Transformers struggle to model, are reducible to POMDPs. This poses a significant challenge for Transformers in learning POMDP-specific inductive biases, due to their lack of inherent recurrence found in other models like RNNs. This paper casts doubt on the prevalent belief in Transformers as sequence models for RL and proposes to introduce a point-wise recurrent structure. The Deep Linear Recurrent Unit (LRU) emerges as a well-suited alternative for Partially Observable RL, with empirical results highlighting the sub-optimal performance of the Transformer and considerable strength of LRU.

Figures (18)

• Figure 1: The Left figure indicates the general structure of recurrent-like sequential neural networks and the right figure represents the attention-like ones. Both of them can be deepened by pointwise transformations and skip connections.
• Figure 2: Above: Illustration for DFA of PARITY. There are two states $q_0$ and $q_1$, where $q_0$ is both the initial state and the accepting state. The transitions are plotted in gray arrows. Below: Illustration for $\mathcal{M}^{\texttt{PARITY}\xspace}\xspace$. The states are $(q_i,w)$ where $i\in\{0,1\}$ and $w\in\{0,1,\#\}$, and the agent could observe $w$. The initial state are randomly sampled from $(q_0,w)$. The stochastic transitions are plotted in gray arrows. At final state $(q_i,\#)$, blue arrows stand for choosing accept, and red arrows stand for choosing reject.
• Figure 3: Hidden state for regular language tasks. We visualize the hidden states of each sequence model during evaluation at length 25 using t-SNE tsne and annotate them according to their real states. Our classification corresponds to the state the observation history maps to in the reduced POMDP, namely $(q, w)$, while 'T' stands for the terminal state. The states with similar colors in the diagram generally produce the same type of observation.
• Figure 4: Learning curves for PyBullet occlusion tasks. Mean of 5 seeds. The shaded area indicates $95$% confidence intervals.
• Figure 5: Mean Squared Error (MSE) ratios for state space modeling tasks. 'V' refers to 'only velocities observable', and 'P' refers to 'only positions observable'. To enable a comparative analysis of the performance among the three models, we present the MSE ratio, as defined in Appendix \ref{['appendix:supplementary']}, Equation \ref{['eq:mse_ratio']}.

Theorems & Definitions (19)

• Proposition 4.1
• Definition 4.2: POMDP derived from regular language $L$
• Remark 4.3
• Theorem 4.4
• Lemma 4.5: Lemma 5 in attentionlim
• Theorem 4.6
• Corollary 4.7
• Definition 1.1: Syntactic congruence semigroups
• Definition 1.2: Syntactic monoid semigroups
• Theorem 1.3: Theorem IX.1.5 in semigroups