Source-linked AI summary
Stabilizing Transformers for Reinforcement Learning
Emilio Parisotto, H. Francis Song, Jack W. Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant M. Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, Matthew M. Botvinick, Nicolas Heess, Raia Hadsell
TL;DR
Partially observable RL requires memory over long horizons, but standard transformers are difficult to optimize under RL objectives. The paper proposes GTrXL, combining identity-map layer-normalization reordering with gated connections. GTrXL improves stability and learning speed, surpasses LSTMs on memory environments, and achieves state-of-the-art DMLab-30 results exceeding MERLIN.
Problem
Standard transformers have not been successfully applied as RL memory architectures and are difficult to optimize under RL objectives.
Method
GTrXL combines layer-normalization reordering with gating mechanisms that replace transformer residual connections.
Results
GTrXL has increased stability and performance, surpasses LSTMs on memory-based DMLab-30 levels, and achieves state-of-the-art results compared with MERLIN.
Takeaways & Limitations
GTrXL provides a simple-to-implement, easy-to-train alternative to standard multi-layer LSTMs for RL agents in partially observable environments.
Takeaways & Limitations
The proposed Markovian-initialization rationale assumes submodules produce values near zero at initialization, and reactive behavior may need to be learned before memory-based behavior.
Abstract
from arXiv · showhide
Owing to their ability to both effectively integrate information over long time horizons and scale to massive amounts of data, self-attention architectures have recently shown breakthrough success in natural language processing (NLP), achieving state-of-the-art results in domains such as language modeling and machine translation. Harnessing the transformer's ability to process long time horizons of information could provide a similar performance boost in partially observable reinforcement learning (RL) domains, but the large-scale transformers used in NLP have yet to be successfully applied to the RL setting. In this work we demonstrate that the standard transformer architecture is difficult to optimize, which was previously observed in the supervised learning setting but becomes especially pronounced with RL objectives. We propose architectural modifications that substantially improve the stability and learning speed of the original Transformer and XL variant. The proposed architecture, the Gated Transformer-XL (GTrXL), surpasses LSTMs on challenging memory environments and achieves state-of-the-art results on the multi-task DMLab-30 benchmark suite, exceeding the performance of an external memory architecture. We show that the GTrXL, trained using the same losses, has stability and performance that consistently matches or exceeds a competitive LSTM baseline, including on more reactive tasks where memory is less critical. GTrXL offers an easy-to-train, simple-to-implement but substantially more expressive architectural alternative to the standard multi-layer LSTM ubiquitously used for RL agents in partially observable environments.
1 INTRODUCTION
Transformers are promising for partially observable RL because they can integrate information across long episodes, yet canonical transformers are difficult to optimize under RL objectives. The paper introduces GTrXL to stabilize training and reports stronger results than LSTMs and MERLIN across challenging environments.
- Transformers avoid compressing the entire past into a fixed-size hidden state, making them attractive for long-horizon sequence processing.
- Partially observable RL episodes can last thousands of steps, with decision-critical observations spanning the entire episode.
- Canonical transformers are significantly difficult to optimize in RL, sometimes achieving performance comparable to a random policy.
- GTrXL combines reordered layer normalization with gating mechanisms to learn faster, more reliably, and with better final performance than canonical transformers.
- GTrXL achieves state-of-the-art results against MERLIN on multitask DMLab-30 and significantly surpasses LSTMs on memory-based levels.
2 TRANSFORMER ARCHITECTURE AND VARIANTS
The transformer processes embedded sequences through stacked self-attention and position-wise MLP blocks, using residual connections, layer normalization, and relative-position memory. Transformer-XL extends contextual access with a stop-gradient memory tensor.
- Each transformer layer applies multi-head attention followed by a position-wise multilayer perceptron to an input embedding sequence.
- The input embedding E^(l−1) has shape [T, D], where T is the number of time steps and D is the hidden dimension.
- Canonical blocks apply residual updates followed by layer normalization after the attention and MLP submodules.
- Relative position encodings and a T-step memory tensor extend the contextual horizon, with the memory treated as constant during weight updates.
- StopGrad prevents gradients from flowing backward through the memory during backpropagation.
3 GATED TRANSFORMER ARCHITECTURES
GTrXL stabilizes transformer optimization by moving layer normalization to submodule inputs and replacing residual connections with gates. The resulting design supports more reliable learning, with gating choices and identity initialization affecting performance and speed.
- 3.1 IDENTITY MAP REORDERING: Identity Map Reordering places layer normalization only on submodule inputs, creating an untransformed path from the first-layer input to the final output.
- 3.1 IDENTITY MAP REORDERING: The reordered architecture is hypothesized to initialize policy and value functions near a Markovian regime before memory-based behavior is learned.
- 3.2 GATING LAYERS: GTrXL replaces transformer residual connections with gating layers to improve performance and optimization stability.
- 3.2 GATING LAYERS: The paper evaluates input, output, highway, Sigmoid-Tanh, and GRU-type gating mechanisms with increasing expressivity.
- 3.2 GATING LAYERS: Gated identity initialization substantially improves optimization stability and learning speed by preconditioning the transformer near a Markovian regime.
4 EXPERIMENTS
Across DMLab-30, Numpad, and Memory Maze, GTrXL improves performance, memory-horizon scaling, learning speed, and optimization stability over transformer and LSTM baselines. Ablations identify GRU gating and gated identity initialization as especially effective, including under parameter control.
- Performance across environments: GTrXL achieves state-of-the-art DMLab-30 performance, surpassing both a deep LSTM and the external-memory architecture MERLIN.The comparison uses human-normalized return averaged across all 30 levels; MERLIN was trained with a different algorithm and for 100 billion environment steps.
- Scaling with memory horizon: GTrXL scales better than LSTM as Numpad memory requirements increase, although its performance also drops at the combinatorial 4×4 setting.On 2×2 and 4×4 Numpad, GTrXL retains a substantial improvement even when LSTM is trained twice as long.
- Performance ablation: The canonical TrXL fails to learn on DMLab-30, whereas TrXL-I improves over LSTM and GTrXL variants differ substantially in stability and sensitivity.GTrXL (Output) recovers much of GRU-gated performance on reactive tasks but is less stable; Input, Highway, and SigTanh variants perform worse or are more sensitive.
- Gating variants: GRU gating learns memory environments faster and reaches the highest final performance, while matching other gating mechanisms on reactive tasks.Figure 4 reports mean curves and individual 6–8 hyperparameter samples per model, with MERLIN shown as a reference score.
- Hyperparameter and seed sensitivity: Across 25 hyperparameter settings and seeds, GTrXL (GRU) learns Memory Maze in fewer environment steps than LSTM and substantially reduces tuning sensitivity.Only GTrXL (GRU) reaches human performance within 2 billion environment steps in the reported sensitivity experiment.
- Gated identity initialization: Gated identity initialization significantly improves optimization stability and learning speed, with initialization near a Markovian policy producing large learning-speed gains.The comparison uses otherwise identical four-layer GTrXL (GRU) models and ten hyperparameter samples per initial bias setting.
5 RELATED WORK
The paper situates its approach among recurrent and external-memory mechanisms for improving memory in partially observable RL. These alternatives can outperform LSTMs on memory tasks but have seen limited adoption partly because of implementation complexity.
- Recurrent architectures: Gating mechanisms improve recurrent-model learnability by addressing vanishing gradients.LSTMs use input, forget, and output gates, while GRUs simplify this design to two gates.
- External memory: External-memory agents typically use feedforward or recurrent policies that read from and write to a memory database.Their designs induce priors through operations resembling digital computers or environment maps.
6 CONCLUSION
The conclusion reports that standard transformers are unstable in RL, while GTrXL improves optimization stability, performance, and robustness through reordered normalization and gating. Across DMLab-30, Numpad, and Memory Maze, GTrXL consistently improves over LSTM architectures.
- Conclusion: Standard transformers are too unstable to train in RL and often fail to learn completely.
- Architecture: GTrXL combines reordered layer normalization with gating layers in place of standard residual connections.The reordered normalization enables an initially Markov regime of training.
- Ablations: GRU-type gating performs best across the evaluated metrics and has robustness to hyperparameters and random seeds comparable to an LSTM.
- Benchmark outcomes: GTrXL demonstrates substantial and consistent improvement over LSTM architectures on DMLab-30, Numpad, and Memory Maze.The authors identify this result as support for wider adoption of transformers in RL.
- Future work: The paper proposes testing how well GTrXL scales with a large and varied set of training environments.
A ENVIRONMENT DETAILS
The experiments evaluate agents in DMLab-30, Numpad, and Memory Maze using specified action spaces, observations, and image-processing components. The DMLab Arbitrary Visuomotor Mapping task is also used to assess memory-dependent performance against human-level behavior.
- Numpad: Numpad episodes last 500 steps and provide proprioceptive information, activated-pad status, the previous action, and reward.The task uses three actions, including movement and jumping.
- DMLab-30: DMLab-30 uses 72 × 96 RGB observations, a reduced action set, and an action repeat of 4.Models receive an additional 64-dimensional LSTM for language inputs on levels requiring language.
- Arbitrary Visuomotor Mapping: GTrXL reliably attains human-level performance on DMLab Arbitrary Visuomotor Mapping with a reduced action set.This task was previously used to highlight MERLIN’s ability to use memory.
- Action specifications: The DMLab native action space and simplified action sets are documented separately for the benchmark and Arbitrary Visuomotor Mapping evaluations.
- Memory Maze: Memory Maze uses 72 × 96 RGB observations and a hybrid native action distribution with 8 continuous actions and one discrete action.
- Image encoder: DMLab-30 and Memory Maze use a ResNet-based image encoder with convolutional, pooling, and residual blocks.
- Agent output: A 256-unit MLP produces policy logits, Gaussian parameters, or value estimates from the encoded representation.
B EXPERIMENTAL DETAILS
Experimental runs sample V-MPO hyperparameters while holding each sampled configuration fixed across models within an experiment. Environment-specific hyperparameters are listed in the accompanying table.
- Hyperparameter sampling: Each run samples V-MPO hyperparameters from a distribution in addition to independently sampling random seeds.The sampled hyperparameters remain fixed across all models in a given experiment.
- Environment-specific settings: DMLab-30 LSTM experiments use a more constrained hyperparameter range because preliminary experiments found it performed better.
- Hyperparameter table: Table 8 lists the V-MPO hyperparameters used per environment.
B.1 TRAINING SETUP
Experiments used a distributed actor-learner setup on TPUs, with architecture details documented for the DMLab-30, Numpad, sensitivity, and gated-identity ablations.
- Distributed training: Actors ran network inference and environment interactions on CPUs, then transmitted trajectories to a centralized learner.Training used TF-Replicator in a 16-core TPU configuration.
- DMLab-30 evaluation: DMLab-30 ablation reporting includes the number of model runs completing 10 billion environment steps.The table also states the convention that transformer embedding dimension equals head dimension multiplied by the number of heads.
- Ablation configurations: Separate architecture-detail tables cover the Numpad, sensitivity ablation, and gated identity initialization ablation experiments.These are identified as Tables 10, 11, and 12, respectively.
C.1 MULTI-HEAD ATTENTION
The attention module combines multiple masked soft-attention operations with residual and normalization components, while relative position encodings and memory extend its contextual horizon.
- Multi-Head Attention: Multi-Head Attention computes H parallel soft-attention operations at each time step and produces an output tensor Y^(l) ∈ R^(T×D).Queries, keys, and values are formed through trainable linear projections, with per-head dimension d = D/H.
- Transformer block: The resulting attention embedding receives a residual connection followed by layer normalization.These operations are applied after combining the query, key, and value tensors.
- Sensitivity ablation: Figure 10 shows 25 sampled hyperparameter settings for the sensitivity ablation, with the x-axis in log scale.Values are sampled from the corresponding ranges in Table 8.
- Attention computation: MaskedSoftmax applies causal masking to prevent the model from addressing future information.The formulation uses Einstein summation notation and omits reshaping operations for simplicity.
- Position and memory: Relative position encodings and a memory scheme provide a much larger contextual horizon than would otherwise be possible.The scheme follows Dai et al. (2019), while the basic MHA operation itself is permutation invariant with respect to sequence order.
- Position and memory: The relative-location keys use a linear projection, while Φ denotes the standard sinusoid encoding matrix and u^(l), v^(l) are trainable parameters.The passage also defines the star symbol as broadcasting.
- Identity Map Reordering: Identity Map Reordering applies layer normalization only to each submodule’s input and applies ReLU to the output stream.The resulting TrXL-I is depicted in Figure 1 (Center).
D DMLAB-30 MEMORY/REACTIVE PARTITION
DMLab-30 levels are partitioned into memory-based and reactive splits for analysis of memory-dependent and more reactive tasks.
- Level partition: Table 13 partitions DMLab-30 levels into memory-based and reactive splits.The table defines the two evaluation categories used for this analysis.