Real-Time Recurrent Learning using Trace Units in Reinforcement Learning

Summary

This paper investigated using complex-valued diagonal RNNs for online RL to provide a small modification (RTUs), and the authors found policy performs significantly better in online RL across various partially observable prediction and control settings.

Strengths

1. The paper is well-written. 2. Method part has good theoretical basis and clear reasoning logic. 3. The experimental results verify the effectiveness of the algorithm in a large number of scenes.

Weaknesses

N/A

Questions

N/A

Rating

5

Confidence

2

Soundness

2

Presentation

3

Contribution

2

Limitations

The article discusses the limitations of the method in detail, which I think is reasonable.

Summary

This paper introduces Recurrent Trace Units (RTUs), a lightweight extension to linear recurrent units (LRUs) that are more suitable for online reinforcement learning. In a range of ablations and experiments the authors show that RTU trained with Real-Time Recurrent Learning (RTRL) performs on par or outperforms LRU trained with RTRL and GRUs trained with BPTT in online reinforcement learning.

Strengths

- the paper is sound - the contribution is clear - extensive ablations on toy experiments - interesting analysis

Weaknesses

- minor contribution - some tasks are not very complex (Mujoco-P and Mujoco-V do not require much memory). - expensive in the multilayer setting, making it limited to simple tasks.

Questions

- 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? - RTUs could be trained with BPTT. This would raise the question whether the increase in performance is exclusively in 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. [1] Time-Efficient Reinforcement Learning with Stochastic Stateful Policies

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors made the main limitation of this method clear. The paper would benefit from answering the above questions. **Some minor points**: - line 118: where did the non-linearity f come from (compare left and right)? - line 148: “Each element of $h_{t-1}$” should be $h_t$ - Why not keeping the $\bar{h}_t$ notation from 3.1 in 3.2? - the paper would benefit from proofreading; it has a couple of grammatical mistakes like "A window of past observations does not scalable well to long 30 sequences [...]" **Conclusion** The paper is nicely structured, with ablations and analysis, and the claimed contribution is fulfilled. Addressing the points above can improve the paper, but the contribution remains small.

Summary

This work proposes Recurrent Trace Unit (RTU) as a modified variant of the linear recurrent unit (LRU) which has gained some popularity recently as a linear complexity model. RTU adopts cosine representations of LRU to manipulate two-coupled real-valued hidden states, and introduces non-linearity (while breaking the equivalence to dense linear RNN when non-linearity is introduced in the recurrence). Like LRU, RTU allows for tractable real-time recurrent learning (RTRL) through diagonal recurrence. The focus of this work is on evaluating RTU-RTRL for reinforcement learning (RL) in partially observable environments. The main experiments are conducted on two partially observable variants of six MuJoCo environments for general evaluation; in addition, one MuJoCo environment (Reacher) and five POPGym environments are used as memory tasks. Positive results for RTU are consistently shown.

Strengths

* Evaluating variants of LRU using RTRL *for reinforcement learning* is novel and interesting. * Beyond the focus on the RL applications, several interesting details about LRU are discussed while developing RTU (Sec 3.2 and 3.4), which should be of interest to people interested in LRU in general. * Experiments are conducted on several relevant environments. Overall, I'm supportive of acceptance, provided that the authors will address the main issues described below.

Weaknesses

**One major issue is the repeated claims about outperforming "Transformers" while no such experiment is provided, certainly not in the main text (not even in the appendix as far as I can tell).** See: - Abstract Line 11, "RTUs significantly outperform GRUs and Transformers" - Line 329: "We also included a GPT2-transformer baseline" - Line 330: "We provide the results for GPT2 in Appendix G." - Conclusion: Line 356: "performed better than the more computationally intensive GRUs and Transformers" I could not find the Transformer/GPT2 results even in the indicated Appendix G, Figures 21 or 22 either. Given all these emphases on this claim, I had expected to see such results in the main text, e.g., in the same style as Figure 4 and 6. I initially considered putting 1 ("incorrect statements") to encourage the authors to fix this immediately, but in the end, I decided to put the score I would give disregarding any mentions to Transformers. Regarding Transformers, please also check the related comment in the "Questions" field below. **There are also some clarity issues:** * Clearer explanations to explain the gap between RTU and LRU are expected. In Figure 1, some explanations are needed to help the readers understand why there is such a big gap between LRU and RTU when everything is equal except the architecture. Is there anything specific to RTRL? or do you observe similar trends when T-BPTT is used for both models? Similarly, some explanations should be provided to explain why Online LTU largely underperform RTU models in Figures 4, 5, and 6. The only related comment I could find is Line 209: *"We found small choices in our implementation for LRU did not always behave as expected, partially due to how auto-differentiation is implemented in packages for complex numbers"* So is LRU's problem just an implementation issue? Please clarify. * I find Figure 2 misleading as they mix model architectures and learning algorithms within the same comparison, i.e., they compare LRU-TBPTT with RTU-RTRL. Instead, the comparison should be between LRU-TBPTT vs. RTU-TBPTT vs. GRU-TBPTT, and LRU-RTRL vs. RTU-RTRL separately. The statement *"LRU and GRU with T-TBTT is not competitive with RTUs"* (caption Figure 2) is true but does not allow us to conclude on the superiority of RTU over LRU since two different learning algorithms are used. I acknowledge that good results are sufficiently shown later in Figures 4, 5 and 6, but Figure 2 alone is not informative. **Some of the experimental designs are not convincing:** * (Related to the point above) I find the experimental setting of Sec 4.2/Figure 2 under "resources constraints" too artificial (therefore not particularly useful). Does the chosen "computational budget of 15000 FLOPs" (Line 237) correspond to something intuitive/useful? How does the memory requirement differ between RTU-RTRL and RTU-TBPTT for different values of T?

Questions

Most of the questions have already been asked above. These are comments and suggestions. * Tractable RTRL with a diagonal recurrent matrix dates back to [R1] [R2] at least, which should be cited. [R1] Gori et al. IJCNN 1989. BPS: A learning algorithm for capturing the dynamic nature of speech. [R2] Mozer. Complex Systems 1989. A focused backpropagation algorithm for temporal pattern recognition. * The authors' view on Transformers for POMDPs (the second paragraph of the introduction; starting at Line 20) is restricted as it lacks discussion about the "linear" variants of Transformers that are stateful and permit a "recursive" formula just like RNNs. There are several prior works evaluating [R3] [R4] [R5] and discussing [R6] such models in the context of reinforcement learning. [R3] Irie et al. NeurIPS 2021. Going Beyond Linear Transformers with Recurrent Fast Weight Programmers. [R4] Irie et al. ICML 2022. A Modern Self-Referential Weight Matrix That Learns to Modify Itself. [R5] Pramanik et al. arXiv 2023. Recurrent Linear Transformers. [R6] Lu et al. ICML 2024. Rethinking Transformers in Solving POMDPs. * In the equation just below Line 118 (sec 3.1), *"f("* is a typo at this stage (the non-linearity is not introduced yet). * Regarding the non-linearity "f", for the tractable RTRL to hold, "f" has to be a *purely* element-wise function. For example, softmax or layernorm would not work there (as they introduce interactions between recurrent units); I know nobody would use softmax or layernorm as a recurrent non-linearity in practice, but it might make sense to point out that there is a condition on "f", in the strict mathematical sense. * It is unfortunate that no further architectural advances have been integrated and evaluated. Based on [11], gating can be made compatible with RTRL, and a recent model that is closely related to LRU, "Mamba", also puts back gating to linear recurrence. So some gating could have been a natural extension too. * The proposed RTU is not specific to RL. I'm wondering if the authors considered applying it to other supervised learning applications.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Comments are already provided above.

Summary

This paper introduces a novel complex-valued, diagonally connected recurrent network model based on LRUs called Recurrent Trace Units. Due to the diagonal connectivity, the model can be trained using real-time recurrent learning in linear time. The authors further propose a version of PPO that uses stale gradients computed by RTRL during interaction. The approach is evaluated on a simple online-prediciton task and a range of POMDP RL tasks.

Strengths

The paper is well written overall. All the necessary background information is summarized in a clear fashion. The main idea is laid out concisely. The experiments used for evaluation are adequate and results are displayed in a visually appealing way. Finally, the direction of research is undoubtedly very significant.

Weaknesses

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. The work is not overly original, but rather an unavoidable next step after the success of LRUs. The claim that taking the Real-part of the latent state leads to a biased gradient is not really convincing, all I see is an expression that is harder to compute. There are many typos and misplaced words. Somehow, the Appendix was more interesting that the actual paper. Minor issues: - L 29: "... does not scalable well ..." - L 32: "... is well suited for update the state online ..." - L 48: The sentence starting on this line looks like it needs a citation, please add something like "as we will show in section..." - L 143, missing word: diagonal "matrix" - L 151: "... choices made in LRUs, that they showed ..." - L 291: The sentence starting on this line is confusing.

Questions

- 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*? - 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?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors comment on the limitation regarding multi-layer architectures. Societal impacts of the work are not considerably different than those of Artificial Intelligence in general ...

Summary

In online reinforcement learning, Recurrent Neural Networks (RNN) are still better than transformers, but there are still problems that need to be addressed. In this paper, the authors propose a modified form of the recurrence equation used for Linear Recurrent Units (LRU) – a type of RNN – called a Recurrent Trace Unit (RTU). The main differences RTU introduces compared to LRU is that it uses a diagonal matrix in place of a full matrix, thus reducing computations and it uses complex valued numbers in the diagonal. Note that multiplications with complex values can be represented as multiplying by a 2x2 matrix with real valued numbers. Additionally, RTUs are a more generalized form of LRUs since they allow nonlinearity in the recurrence equation, unlike LRUs which are strictly linear. This leads to improved accuracy. The authors first explain their recurrent formulation in detail and justify their design choices. In the second half of the paper the authors empirically perform multiple comprehensive tests that show that RTUs outperform other architectures (GRUs, LRUs, and Transformers [note: Transformer results included in supplementary] ) in online reinforcement learning tasks, providing better performance with less computation. These tests include ablation, scaling of learning under various constraints, comparison of two different training methods – Real-Time Recurrent Learning (RTRL) (i.e. incremental updates) and Truncated Backpropagation Through Time (T-BPTT) (i.e. batched updates), and analysis of how well RTUs ‘remember’ information compared to other architectures.

Strengths

• The paper addresses an important problem in online reinforcement learning and proposes a computationally more efficient solution, RTUs, which outperform other architectures in various online reinforcement settings • Has comprehensive tests covering and analyzing a variety of cases, which empirically show the improvements from the proposed architecture, RTU. • Provides detailed explanations and derivations of the RTUs formulas, making it easier for readers to understand and replicate it • Results are explained well using figures • Supplementary section goes into detail on many areas which are not fully explained in the main text and answers many questions that may arise.

Weaknesses

• Most of the comprehensive tests compare the proposed RTU only with GRU and LRU architectures (there are a few transformer (GPT) tests in the supplementary). It would be better if the author uses more architectures or explains why these other two architectures are sufficient for a comparison • Explaining the importance and benefits of addressing problems with online reinforcement learning in partially observable environments would strengthen the motivation. • Some terminology not completely clear, such as difference between “Reinforcement learning” and “Online reinforcement learning”

Questions

However, 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?

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

Paper is very comprehensive; it motivates and both mathematically (in the supplementary) and empirically shows the improvement of their proposed architecture very well. 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. The author performs many types of tests comparing the proposed architecture with GRU and LRU, including a few transformer tests in the supplementary. However, is comparing with GRU and LRU enough for online reinforcement learning tasks? Testing different architectures or briefly explaining why the chosen architectures are sufficient would be better. In the motivation, the author motivates why RNNs are still state of the art in reinforcement learning in partially observable environments and why transformers are not suitable for this task. However, the reason why reinforcement learning in partially observable environments is important is not explained. Having a few sentences on the importance addressing of this topic would strengthen the motivation. Lastly, in the introduction and abstract, the author mentions both reinforcement learning and online reinforcement learning, but does not explain the differences. The authors should clarify whether these are the same concepts or are distinct. Grammar: • 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?

Thank you for the response

I thank the authors for their response. > We do provide the results for GPT-2 in Appendix G. Figure 21 Thank you for the pointer to Figure 21 and 26, and the updated figures. The authors' comment regarding the "bad" colors in the old figure is accurate; the new colors help. But I also do confirm that it does not make sense to emphasize on Transformer-related claims that are entirely based on results that are hidden in Appendix G, Figure 21 and 26 (and for other reasons I explain below). > 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. It is unreasonable to complain about the model's parameter count, as it is the deliberate choice made by the authors. On the contrary, the current comparison seems unfair to the Transformers. GPT-2 is a language model instantiation of Transformers. To properly evaluate Transformers in specific RL tasks, their hyper-parameters have to be adjusted accordingly. In particular for POMDPs, a discussion on the context length is a minimum requirement in a proper evaluation, since if their context length is shorter than the memory span required for the task, they will obviously fail (we won't even need to run experiments to figure that out). In this sense, I find all the current “GPT-2” related results reported in this paper insignificant, so are the claims regarding the "superiority" over "Transformers". I understood that this choice of Transformer architecture/GPT-2 was based on a prior work, but I do not see any good reason to follow such a choice which does not make sense. I will maintain my currently score on the condition that the authors promise to remove all these claims from the main text, as I initially stated. Alternatively, the authors should conduct a new set of proper experiments by tuning the Transformer hyper-parameters/architecture. What follows are minor comments: > Our goal from Figure 1 was twofold: Understand the utility of RTU parametrization in contrast to LRU. This was actually my original critique: Figure 1 does not fulfill this goal. Figure 1 currently *shows* the utility of RTU but fails at helping us "understand" why it is better than LRU. It will make sense to add a sentence toward the end of Sec 4.1 to explain this gap and guide readers to take a look at Sec 3.4 and Appendix D.2. > 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, I do not agree with this argument. I asked whether the gap between RTU and LRU would remain the same when T-BPTT was used. If we train both models with T-BPTT, the condition is the same for the two models, so the author’s argument does not hold. That said, I had also overlooked some details in Sec 3.4 and Appendix D.2, so I no longer find this experiments crucial.