Real-Time Recurrent Learning using Trace Units in Reinforcement Learning

We thank the reviewer for the detailed and valuable feedback. Here, we respond to the points the reviewer mentioned.

some tasks are not very complex (Mujoco-P and Mujoco-V do not require much memory).

While Mujoco-P and Mujoco-V could be considered short-term memory tasks, Reacher POMDP and some of the POPGym environments are long-term memory tasks. Table 1 in Ni et al. (2023) provides an estimate of the memory length needed to solve some tasks, and both Reacher-POMDP and Autoencode from POPGym have long memory requirements. Other environments we tested, such as CountRecall, Repeat First, and Concentration, also need a long history to be solved effectively.

[1] Tianwei Ni, Michel Ma, Benjamin Eysenbach, and Pierre-Luc Bacon. When do transformers shine in rl? decoupling memory from credit assignment. In the Thirty-seventh Conference on Neural Information Processing Systems, 2023.

expensive in the multilayer setting, making it limited to simple tasks.

It’s true that getting the exact gradient for the multi-layer case is expensive. A common solution to this problem is to stop the gradient between the recurrent layers’ hidden states [1][2], this approach has been shown to work in practice and can be used when scaling our RTUs for multi-layers. While this solution has been shown to work in practice, we think that a more principled approach is still needed.

[1] Kazuki Irie, Anand Gopalakrishnan, and Jürgen Schmidhuber. Exploring the promise and limits of real-time recurrent learning, 2023

[2]Nicolas Zucchet, Robert Meier, Simon Schug, Asier Mujika, and João Sacramento. Online learning of long range dependencies. In Advances in Neural Information Processing Systems, 2023.

in bptt the history is often needed to be truncated because of vanishing/exploding gradients. In theory, RTRL should also suffer from these problems, how are they mitigated here? Is it because of the stale gradients in practice?

That’s correct, RTRL can still suffer from vanishing/exploding gradients. In our parametrization, we can prevent vanishing/exploding gradient by restricting the magnitude of the complex number to be $\in (0,1]$. We discuss this point first in Appendix D.1, showing that different parametrization requires different restrictions to avoid vanishing/exploding gradients and that we choose the cosine representations as it has the least restrictions. In Appendix C.1, we discuss several ways to enforce the restriction on the magnitude of complex numbers and provide an ablation over these different approaches.

RTUs could be trained with BPTT. This would raise the question whether the increase in performance is exclusively in RTRL.

That’s true. RTUs can be trained with both BPTT and RTRL. We show an experiment in Appendix E.1 (Figure 11) comparing RTUs when trained with BPTT and RTRL. RTRL does play an important role in the improved performance of RTUs over other architectures; RTUs with BPTT will have the same issue of truncation length/computational complexity trade-off. However, the parametrization of RTUs was needed to get to a tractable RTRL.

It would be interesting to discuss the differnce between the RTRL approach with the approach in [1], where the authors train general stateful policies with stochastic gradient approximation instead of BPTT; especially because the experiments are related.

Thank you for pointing out this paper, we will add it to our discussion on previous work. While the paper introduces a way of calculating an unbiased gradient estimation of BPTT, they note that their estimation has high variance since they introduce stochasticity to the policy's internal state. This is in contrast to our approach where there is no added stochasticity. Additionally, the proposed approach in [1] was not able to learn both the policy and the value function at the same time without providing additional privileged information to the value function. Our approach doesn’t have these limitations.

line 118: where did the non-linearity f come from (compare left and right)?

That’s a typo. The f shouldn’t be there at this point of the derivation. Thank you for pointing it out; we will fix it.

line 148: “Each element of ht−1” should be ht

That’s correct. Thank you for pointing it out.

We thank the reviewer for the detailed and valuable feedback. Here, we respond to the points the reviewer mentioned.

One major issue is the repeated claims about outperforming "Transformers" while no such experiment is provided.

We do provide the results for GPT-2 in Appendix G. Figure 21 shows the GPT-2 performance on Mujoco-P tasks, and Figure 26 has the learning rate sweep results for GPT-2. In Figure 21, the GPT-2 line is far below all the baselines we have, so it’s easy to miss. We changed it to a brighter color now, so it’s obvious. Please check the uploaded PDF for the updated figures.

The reason for not including GPT-2 in the main text is that the number of parameters in GPT-2 is way more than any of the architectures we have, so the comparison is a bit unfair. We will modify the abstract so as to not put a lot of emphasis on this comparison since it’s not a core result of the paper.

In Figure 1, some explanations are needed to help the readers understand why there is such a big gap between LRU and RTU

Our goal from Figure 1 was twofold:

1. Understand the utility of RTU parametrization in contrast to LRU.
2. Understand the effect of different design choices, such as non-linearity and using a linear projection layer. By controlling for other factors, we can ascertain that the gap in performance is due to the difference in parametrization choices, i.e., using explicit complex numbers in LRU versus using rotation matrices in RTUs. Using rotation matrices avoids the complications around taking the real part of the hidden states; we show in Appendix D.2 that taking the real part of the hidden states could result in a biased gradient estimate.

For Figure 1, we chose to use RTRL as using T-BPTT would complicate the comparison. Bad performance could be attributed to the truncation length rather than the architecture itself, while with RTRL, there is no confounding factor other than the parametrization.

Similarly, some explanations should be provided to explain why Online LTU largely underperform RTU models

When comparing RTU with Online LRU, the core difference is RTU’s parametrization of complex numbers as rotation matrices. As we noted above, using rotation matrices avoids the complications around taking the real part of the output. We also noticed that using complex numbers combined with auto-diff libraries resulted in unexpected results when taking the gradients. This appears to be an issue related to the use of complex numbers in general, not specific to LRU. We pointed to a similar discussion in the footnote on page 5. We also emailed Online LRU authors to verify that our implementation is correct, and we used their suggested implementation of Online LRU in our experiments.

I find Figure 2 misleading as they mix model architectures and learning algorithms within the same comparison.

The goal of Figure 2 was to highlight the trade-off between computational resources and performance when using T-BPTT versus RTRL. So, we could compare RTU-RTRL and RTU-TPBTT. However, we believe the reader would be more interested in seeing the performance of LRU-TBPTT and GRU-TBPTT to see how much benefit we get on top of these more widely used approaches when moving to RTU-RTRL. The plot becomes unreadable with too many choices. But your point is a good one, and we do provide a comparison of RTU-BPTT with LRU BPTT in Appendix E, Figure 11. As expected, when using T-BPTT with RTU, the truncation length affects the performance of RTU as well.

Does the chosen "computational budget of 15000 FLOPs" (Line 237) correspond to something intuitive/useful?

We wanted to start with a computational budget under which all architectures can learn and produce relatively good performance so the comparison is meaningful. We looked at the sizes of the GRU architectures used in [1], where the animal learning experiments were introduced, and chose similar hidden dimensions to start with. That’s how we came to the choice of 15k flops specifically.

[1] Banafsheh Rafiee, Zaheer Abbas, Sina Ghiassian, Raksha Kumaraswamy, Richard S Sutton, Elliot A Ludvig, and Adam White. From eye-blinks to state construction: Diagnostic benchmarks for online representation learning. Adaptive Behavior, 2020.

diagonal recurrent matrix dates back to [R1] [R2]

Thanks for pointing out these references. We will cite them in the introduction.

The authors' view on Transformers is restricted

Thanks for pointing out those relevant papers. We will modify the discussion on transformers to include them. While the approaches in [R3] [R4][R5] are all interesting and relevant, they face the same challenges as T-BPTT-based approaches in terms of the trade-off between sequence length, computational complexity, and performance.

Regarding the non-linearity "f", for the tractable RTRL to hold, "f" has to be a purely element-wise function.

We agree that the use of layer norm as an activation function is not conventional. The discussion we provided in Appendix A.2 was to mathematically state what conditions we need on f for the equivalence to hold. However, as we point out in the main body, we think of these theoretical results as negative results; as you pointed out, those functions are not element-wise and not conventional activation functions.

Based on [11], gating can be made compatible with RTRL. So some gating could have been a natural extension too.

Thank you for the suggestion. We agree that gating is a natural extension of this work.

The proposed RTU is not specific to RL.

We agree that RTUs are not specific to RL. However, we think RTRL approaches are more needed in the context of online RL than supervised learning. As in supervised learning, offline datasets are usually available, and having access to a long history is not as hard as in online RL. There are areas of online supervised learning where RTRL approaches could be useful, which could be explored in future work.

We thank the reviewer for the detailed and valuable feedback. Here, we respond to the points the reviewer mentioned.

A thorough ablation study of the proposed architecture is missing, i.e. taking discrete steps from LRUs towards RTUs and comparing all of them in one plot, including BPTT.

Thank you for raising this point. Currently, in Figure 1, we provide an ablation on several architectural choices and test them one at a time while fixing everything else. This means that the only difference between RTU and LRU lines in Figure 1 is the RTU parametrization, which is the main difference between LRU and RTU.

In the left plot of Figure 1 (titled: linear recurrence), we are testing the minimal architectural choice that could be made, a linear recurrence. Hence, this plot is only testing the utility of the RTU parametrization over LRU’s with the absence of any other factor. We then move to test different architectural choices that could be made for either of them, such as adding non-linearity or a projection layer.

For Figure 1, we chose to use RTRL as the learning algorithm rather than T-BPTT. We made this choice to prevent any other confounding factor (such as truncation length in T-BPTT) from affecting the performance of either architecture. We still tested RTUs with T-BPTT in the appendix, Figure 11, but as expected with T-BPTT, the performance of RTUs would reduce due to truncated gradients.

Somehow, the Appendix was more interesting that the actual paper.

We are glad that you found the appendix interesting. We tried to extensively answer any questions that might arise about our approach, but we still didn’t want to overwhelm the reader with a lot of mathematical details in the main body and rather focus on the core ideas, so we chose to delegate most of the mathematical work to the appendix.

Can you comment on the claim of a biased gradient? Surely, information is lost by discarding the real part, but the gradient should not be biased?

The discarded information from taking the real part also results in discarded gradient components. This makes the gradient estimation different from the true gradient and biased towards the information coming from the real components. We provide a simple derivation of that in Appendix D.2. It is worth noting, though, that the biased gradient issue happens when the recurrent part of LRU is not followed by a linear projection. Hence, for LRU to have an unbiased gradient estimate, it needs to always have a linear projection layer afterward, which is restrictive.

Additionally, while the original LRU paper proposes always having a linear projection after LRU, our empirical evaluations in section 4.1 show that this linearity restriction actually harms LRU's performance.

One major advantage of LRUs is the possibility to employ parallel scans reducing the complexity to sublinear in the number of units. Is this advantage thwarted by adding a non-linearity?

Currently, Linear RTUs in Eq.2, where we apply non-linearity afterward, not on the recurrence, can still benefit from parallel scans, but the Non-Linear RTUs cannot since the recurrence part is non-linear. However, there are some recent works on parallelizing non-linear recurrence, which could be explored with Non-Linear RTUs as well [1][2].

[1] Lim, Y. H., Zhu, Q., Selfridge, J., & Kasim, M. F. Parallelizing non-linear sequential models over the sequence length. In The Twelfth International Conference on Learning Representations.

[2] Gonzalez, X., Warrington, A., Smith, J. T., & Linderman, S. W. (2024). Towards Scalable and Stable Parallelization of Nonlinear RNNs. arXiv preprint arXiv:2407.19115.

We thank the reviewer for the detailed and valuable feedback. Here, we respond to the points the reviewer mentioned.

Explaining the importance and benefits of addressing problems with online reinforcement learning in partially observable environments would strengthen the motivation.

We motivated the importance of learning under partial observability in the first paragraph of the introduction when talking about how in the real world, agents perceive their environments through imperfect (i.e., partial) sensory observations. Hence, for agents to be deployed and learn online in the real world, they need the ability to predict and control under partial observability. We agree that the motivation for learning under partial observability is core to our work, and we will expand our current motivation in the introduction to emphasize the importance of learning under partial observability.

Some terminology not completely clear, such as difference between “Reinforcement learning” and “Online reinforcement learning”

Thanks for pointing that out. When we say "online reinforcement learning" we mean that the agent learns while interacting with the environment. This contrasts with offline reinforcement learning or reinforcement learning with access to a simulator. We will add a couple of sentences to the intro to make this clear.

is comparing with GRU and LRU enough for online reinforcement learning tasks? Or, would testing different architectures or explaining why the chosen architectures are sufficient would be better?

We chose GRU as our main baseline since it’s widely used for reinforcement learning problems with partial observability and has been shown to outperform most of the other memory components in several domains, which makes it a strong baseline to compare against [1][2][3][4]. Additionally, the authors of the original POPGym paper that we use in section 6 have extensively tested various recurrent and transformer architectures on the benchmark, and GRU was always the best-performing architecture [1]. Since we are reusing their benchmark and similar other tasks, it is natural to compare against GRU as they have already done extensive testing with the most known architectures. While LRU has not been explored in reinforcement learning domains, we added it to our experiments because it is similar to our proposed architecture.

[1] Steven Morad, Ryan Kortvelesy, Matteo Bettini, Stephan Liwicki, and Amanda Prorok. POPGym:Benchmarking partially observable reinforcement learning. In International Conference on Learning Representations, 2023

[2] Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhutdinov. Recurrent model-free rl can be a strong baseline for many pomdps. In International Conference on Machine Learning, 2022.

[3] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104, 2023.

[4] Pleines, M., Pallasch, M., Zimmer, F., & Preuss, M. (2023). Memory gym: Partially observable challenges to memory-based agents. In The eleventh international conference on learning representations.

One downside is that the paper has a lot of mathematical formulations that may be difficult to follow for someone not well versed in this topic. However, these formulations are necessary and each step of the formulations is explained in detail, making it easier to follow.

We are glad that you found our mathematical formulations easy to follow. We agree that mathematical formulations can be dense, but we think they were all needed to provide the reader with a thorough understanding of our approach and answer any question that might arise about it. We chose to delegate most of the mathematical analysis to the appendix to allow the reader to understand the main contributions easily and then dive into more details about the specifics in the appendix.

• Line 188-190: “RTUs can be 189 seen as a small generalization of LRUs, moving away from strict nonlinearity” shouldn’t this be linearity, since LRU use linear operations?

That’s correct. The sentence should be “RTUs can be seen as a small generalization of LRUs, moving away from strict linearity”