Real-Time Recurrent Learning using Trace Units in Reinforcement Learning

TL;DR

Recurrent Trace Units (RTUs) are introduced, a small modification on LRUs that are nonetheless find to have significant performance benefits over LRUs when trained with RTRL, and significantly outperform other recurrent architectures across several partially observable environments while using significantly less computation.

Abstract

Recurrent Neural Networks (RNNs) are used to learn representations in partially observable environments. For agents that learn online and continually interact with the environment, it is desirable to train RNNs with real-time recurrent learning (RTRL); unfortunately, RTRL is prohibitively expensive for standard RNNs. A promising direction is to use linear recurrent architectures (LRUs), where dense recurrent weights are replaced with a complex-valued diagonal, making RTRL efficient. In this work, we build on these insights to provide a lightweight but effective approach for training RNNs in online RL. We introduce Recurrent Trace Units (RTUs), a small modification on LRUs that we nonetheless find to have significant performance benefits over LRUs when trained with RTRL. We find RTUs significantly outperform other recurrent architectures across several partially observable environments while using significantly less computation.

Figures (27)

• Figure 1: Ablation over different architectural choices for RTUs and LRUs. The RTU variants are blue, and the LRU variants are orange. In each subplot, we restrict both architectures in a particular way, reporting prediction error (MSRE) as a function of hidden state size. Across variations, RTUs are often better and, at worst, tie LRU. Here, both architectures were using RTRL.
• Figure 2: Learning under resources constraints in Trace Conditioning. Each of the four subplots shows how each algorithm's performance varies as a function of resources. (a) LRU and GRU with T-TBTT is not competitive with RTUs even as $T$ is increased while restricting the number of hidden units in LRU and GRU so that all algorithms use about the same computation per step. (b) If we allow GRU and LRU's computation to increase (fixed network size) while increasing $T$, the performance gap remains. (c) Fixing $T$ to a large value to solve the task, we can increase the number of parameters, holding the computation equal for all methods. (d) If we do not require compute to be equal across methods as we scale parameters, then the LRU can eventually match the error of RTU, but GRU cannot. The black dashed line represents the near perfect prediction performance.
• Figure 3: Contrasting runtime in incremental and batch settings. In the incremental setting, evaluated in the animal-learning prediction task, T-BPTT updates scale with truncation length, whereas linear RTRL is constant. With batch updates, evaluated in Ant-P with PPO, linear RTRL remains linear and T-BPTT is slightly more efficient.
• Figure 4: Learning curves on the Mujoco POMDP benchmark. Environments with -P mean that velocity components are occluded from the observations, while -V means that the positions and angles are occluded. All architectures have the same number of recurrent parameters ($~24$k parameter). For each architecture, we show the performance of its best-tuned variant.
• Figure 5: Reacher, $30$ runs with standard errors.

Theorems & Definitions (7)

• Proposition B.1
• proof
• Definition C.1
• Lemma C.2
• proof
• Theorem C.3
• proof