Recurrent Action Transformer with Memory

TL;DR

A Recurrent Action Transformer with Memory (RATE) is proposed, a novel model architecture that incorporates a recurrent memory mechanism designed to regulate information retention that can significantly improve performance in memory-intensive environments, while maintaining or improving results in classic environments.

Abstract

Recently, the use of transformers in offline reinforcement learning has become a rapidly developing area. This is due to their ability to treat the agent's trajectory in the environment as a sequence, thereby reducing the policy learning problem to sequence modeling. In environments where the agent's decisions depend on past events (POMDPs), it is essential to capture both the event itself and the decision point in the context of the model. However, the quadratic complexity of the attention mechanism limits the potential for context expansion. One solution to this problem is to extend transformers with memory mechanisms. This paper proposes a Recurrent Action Transformer with Memory (RATE), a novel model architecture that incorporates a recurrent memory mechanism designed to regulate information retention. To evaluate our model, we conducted extensive experiments on memory-intensive environments (ViZDoom-Two-Colors, T-Maze, Memory Maze, Minigrid-Memory), classic Atari games, and MuJoCo control environments. The results show that using memory can significantly improve performance in memory-intensive environments, while maintaining or improving results in classic environments. We believe that our results will stimulate research on memory mechanisms for transformers applicable to offline reinforcement learning.

Figures (16)

• Figure 1: Recurrent Action Transformer with Memory (RATE). $R$ -- returns-to-go, $o$ -- observations, $a$ -- actions, $M_n$ -- segment's $S_n$ memory embeddings.
• Figure 2: Memory-intensive environments with different observation spaces and reward functions used to test the performance of the memory mechanism in the RATE model.
• Figure 3: Results for the ViZDoom-Two-Colors: with (a) and without (b) pillar in the first $45$ steps of the episode; calculated at environment steps $0$ -- $89$ (c) and $90$ -- $179$ (d); depending on the return-to-go (e, f, g). Pillar disappears after first $45$ steps, $K_{eff} = 90$.
• Figure 4: Results for the T-Maze environment with $K_{eff}=30\times 3 = 90$. The notations are represented as MODEL-N, where N is the number of segments into which the trajectories are divided.
• Figure 5: Results for the Minigrid-Memory environment, $K_{eff}=10\times 3=30$.