[Paper Review] Stabilizing Transformers for Reinforcement Learning

Emilio Parisotto, Hao Song+11 more|arXiv (Cornell University)|Oct 13, 2019

Reinforcement Learning in Robotics39 references131 citations

TL;DR

The paper introduces Gated Transformer-XL (GTrXL), a gated, reordered-layer-transformer architecture that stabilizes training and improves memory-based reinforcement learning performance, outperforming LSTMs and external memory on DMLab-30 and scalable memory tasks.

ABSTRACT

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.

Motivation & Objective

• Motivate the use of transformers for long-horizon, partially observable RL problems.
• Identify training instability of canonical transformers in RL settings.
• Propose architectural modifications (identity map reordering and gating) to stabilize learning.
• Demonstrate that GTrXL surpasses LSTMs and external memory on memory-based benchmarks.
• Show robustness to seeds and hyperparameters while maintaining competitive performance.

Proposed method

• Adapt the Transformer-XL architecture with relative position encodings for RL memory.
• Apply Identity Map Reordering by placing LayerNorm only on the input stream of submodules (TrXL-I).
• Introduce gating mechanisms to replace residual connections in MHA and MLP submodules (GTrXL).
• Explore GRU-type gating as the strongest variant (GTrXL GRU) and several ablations (Input, Output, Highway, SigTanh).
• Initialize gating biases to encourage near-identity mappings to bootstrap Markovian policy learning.
• Train with V-MPO (an on-policy MPO variant) to evaluate learning stability and performance across domains.

Experimental results

Research questions

• RQ1 Can transformers be stabilized sufficiently to serve as memory architectures in RL agents?
• RQ2 What architectural changes (layer normalization order and gating) can improve RL training stability for transformers?
• RQ3 How does GTrXL perform relative to LSTMs and external memory architectures on memory-demanding RL benchmarks?
• RQ4 How robust is GTrXL to hyperparameters, seeds, and varying memory horizons?
• RQ5 Does GTrXL scale with memory horizon and continue to outperform traditional architectures on complex tasks?

Key findings

• GTrXL with GRU gating substantially outperforms a competitive 3-layer LSTM baseline on DMLab-30 in memory-based environments.
• GTrXL (GRU) achieves state-of-the-art results on the multitask DMLab-30 benchmark, surpassing external memory MERLIN in final performance.
• GTrXL demonstrates better memory horizon scaling than LSTMs in Numpad tasks, maintaining superior performance as memory demands grow.
• GTrXL variants with gating outperform other gating options in stability and learning speed, with GRU gating providing the strongest results across tasks.
• GTrXL remains competitive or superior on reactive tasks where memory is less critical, indicating broad applicability as an RL memory replacement for LSTMs.