Real-Time Recurrent Learning using Trace Units in Reinforcement Learning

Authors: esraa  elelimy
adam  white
michael  bowling
martha  white
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 06:25:16

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 GRUs and Transformers across several partially observable environments while using significantly less computation.

Main paper: 
 introduction: 
 Agents, animals, and people perceive their surrounding environment through imperfect sensory observations. When the state of the environment is partially observable, agents construct and maintain their own state from the stream of observations. The constructed agent state summarizes past environment-agent interactions in a form that is useful to predict and control future interactions [41] . Recurrent Neural Networks (RNNs) provide a flexible architecture for constructing agent state [19, 21, 13, 6, 10] .While standard RNN architectures have been mainly supplanted by Transformers [44] , in online reinforcement learning settings where the agent learns while interacting with the environment, RNNs remain a promising direction to pursue [17, 11] . There are two main issues that limit the use of selfattention mechanisms from Transformers in online learning. First, calculating the similarity between each pair of points results in a computational complexity that is a function of k 2 , where k is the sequence length. Moreover, calculating the similarity between all pairs ignores the temporal order of the data points, which limits the usefulness of self-attention when the data is temporally correlated [47] . Second, we need access to the whole sequence of observations before taking an action or updating the learnable parameters, which is impractical in continual learning. While recent works have reduced the complexity of transformers from quadratic in the sequence length to linear [15, 16, 37] , the entire sequence length is still needed to train such architectures. Gated Transformer-XL attempts to overcome this issue by keeping a moving window of previous observations [35] . A window of past observations does not scale well to long sequences-the computation is quadratic in the sequence length-and a window is one particular fixed function of history. The simpler recursive form in RNNs, on the other hand, can learn a great variety of functions of history from the data and is well suited for updating the state online from a sequential data stream and have been shown to outperform transformers in such settings [23] .A key open question is how to efficiently train RNNs in online RL. We can divide the literature into methods that approximate Real-Time Recurrent Learning (RTRL) and those that restrict the recurrent architecture. RTRL [46] exploits the recursive nature of the gradient for RNNs, carrying forward the needed gradient information instead of unrolling the recurrent dynamics back in time like Truncated Backpropagation Through Time (T-BPTT) [45] . RTRL avoids storing past data but is so computationally expensive and is intractable for even moderately sized networks. Several methods approximate the RTRL gradient update, including NoBackTrack [33] , Unbiased Online Recurrent Optimization [43, 5] , Sparse N-Step Approximation (SnAp) [28] . All of these methods produce a biased gradient estimate. Other works have tried to approximate an unbiased gradient estimate of BPTT specifically for the case of policy gradient updates in RL. However, the approximation resulted in a high variance due to added stochasticity to the policy [1] .Methods in the second category usually restrict the RNN architecture to a diagonal RNN [9, 30] , including Columnar Networks [18] , the element-wise LSTM [17] , and Independently Recurrent Neural Networks (IndRNNs) [20] . The RTRL algorithm is computationally efficient for such architectures. However, this approach sacrifices representation power and can perform poorly [18] . Recent work suggests overcoming the poor performance of diagonal RNNs with a small modification: having a complex-valued recurrent state instead of restricting it to real values [34] . In fact, as we will show in section 3.1, there exists an equivalence between using a dense linear recurrent layer and a diagonal recurrent layer with complex values, indicating no loss of representational capacity. LRUs have been combined with RTRL [48] , though only empirically explored for supervised learning datasets.In this work, we extend the insights from LRUs into the online RL setting. Our primary contribution are our experiments showing that such a lightweight recurrent architecture can outperform standard approaches like Gated Recurrent Units (GRUs) [4] in RL, with significantly less computation. To obtain this result, we propose a small extension on LRUs, which we call Recurrent Trace Units (RTUs). RTUs incorporate nonlinearity into the recurrence and use a slightly different parameterization than LRUs, but one we find is more amenable to the use of RTRL in online RL than LRUs. We extend Proximal Policy Optimization (PPO) [40] to use RTRL, ablating the decision choices we propose. We provide an in-depth study in an animal-learning prediction benchmark, showing that RTUs scale better than GRUs with increasing compute and number of parameters and that RTUs outperform alternative diagonal recurrent architectures trained with RTRL. We then show across numerous control environments that RTUs have comparable or better performance, compared to GRUs and LRUs. 
background: 
 We formalize the problem setting as a Partially Observable Markov Decision Process (POMDP). At each time step t = 0, 1, 2, . . ., the agent perceives an observation x t , a limited view of the state s t ∈ S, and takes an action A t ∈ A(s t ). Depending on the action taken, the agent finds itself in a new state s t+1 ∈ S, observes the corresponding observation x t+1 and a reward R t+1 ∈ R. In the online control setting, the agent's goal is to maximize the discounted sum of the received rewards. It may also make predictions about its environment, such as future observations' outcomes.For prediction and control in a partially observable environment, the agent should use the stream of observations to construct its agent state. The agent state summarizes information from the history of the agent-environment interactions that are useful for prediction and control [41] . We could use the whole history up to t, namely (x 0 , A 1 , R 1 , x 1 , A 2 , R 2 , . . . x t ), as the agent state. Though the history preserves all the information, it is not feasible to use directly. We want the agent to have constant memory and computation per time step and storing the whole history causes the memory and the computation to grow with time. Instead, the agent needs to compress this history into a concise representation. We refer to the agent's internal representation of the history at time t as its agent state or its hidden state h t . The agent constructs its current agent state h t ∈ R n from its previous agent state h t-1 ∈ R n and the recent observation x t ∈ R d using a state-update function g : R n × R d → R n : h t = g(h t-1 , x t ).One way to learn this state-update function g is with a recurrent neural network (RNN). A simple form is a linear recurrent layer, where g(h t-1 , x t ) = W x x t + W h h t-1 for weight matrices W x and W h . We can also add a nonlinear activation, such as ReLU.In general, we will write h t = g(h t-1 , x t , ψ ψ ψ), where ψ ψ ψ are the learnable parameters in the network. The agent maps the agent state h t to an output ŷt and then receives a loss L t . = L(ŷ t , y t ) indicating how far the output is from a target y t . The agent updates ψ ψ ψ to minimize this loss over time.Two main gradient-based algorithms are widely used to train RNNs: Truncated Backpropagation Through Time (TBPTT) and Real-Time Recurrent Learning (RTRL). T-BPTT specifies a truncation length T , which controls the number of steps considered when calculating the gradient [45] . As a result, the computation and memory complexities of T-BPTT are functions of the truncation length. Learning with T-BPTT involves a trade-off between the network's ability to look further back in time and its compute and memory requirements. Picking a large T can be expensive, or require us to limit the network size, but picking too small of a T can cause the agent to miss important relationships and so result in poor performance. Williams and Zipser (1989) introduced the Real-time Recurrent Learning algorithm (RTRL) as a learning algorithm for continual recurrent learning. Instead of unrolling the recurrent dynamics back in time, RTRL computes the gradient using the most recent observation, and the gradient is calculated and carried from the last step [46] . Assuming the network parameters have not changed, this recursive form gives the exact gradient and does not suffer from the truncation bias inherent to T-BPTT. We provide a more detailed background on the BPTT and RTRL in Appendix A. In reality, the agent updates its parameters frequently, so the gradient information saved from previous time steps is stale, i.e., calculated w.r.t old parameters; yet, under the assumption of small learning rates, RTRL is known to converge [46] . These properties make RTRL ideal for online learning, but unfortunately, there is a catch: its computational complexity is quartic, of fourth order, in the size of h t , which can be prohibitively expensive. For this reason, we pursue a restricted diagonal form in this work, for which RTRL is efficient and linear in h t . 
recurrent trace units: 
 In this section, we introduce Recurrent Trace Units (RTUs). We start by revisiting why complexvalued diagonals represent dense recurrent layers, and why using real-valued diagonals is insufficient. We then introduce the specific form for RTUs that leverages this relationship. We then provide the RTRL update for RTUs, highlighting that it is simple to implement and linear in the hidden dimension. We finally contrast RTUs to LRUs and motivate why this small extension beyond LRUs is worthwhile.Assume we have the recurrence relationship, with learnable parametersW h ∈ R n×n and W x ∈ R n×d , h t . = W h h t-1 + W x u(x t ), where u can be any transformation of the inputs x t before they are inputted into the recurrent layer. We can rewrite the square matrix W h using an eigenvalue decomposition W h = P Λ P -1 , where P contains the n linearly independent eigenvectors and Λ ∈ C n×n is a diagonal matrix with the corresponding eigenvalues. Then we have thath t = P(Λ P -1 h t-1 + P -1 W x u(x t )) =⇒ P -1 h t = ΛP -1 h t-1 + P -1 W x u(x t ) By defining h t . = P -1 h t ∈ C n and W x . = P -1 W x ∈ C n×d , we get a new recurrence h t = Λh t-1 + W x u(x t ).We can see h t and h t are representationally equivalent: they are linearly weighted for downstream predictions, and so the linear transformation on h t can fold into this downstream linear weighting. But it is more computationally efficient to use h t with a diagonal weight matrix Λ, meaning each hidden unit only has one recurrent relation instead of n. LRUs precisely leverage this equivalence [34] . Specifically, they learn a complex-valued h t , and use Re(W h t ) as an input to an MLP for downstream nonlinearity.Since we did not impose constraints on the matrix W h , other than being diagonalizable, the eigenvalues of W h can be complex or real numbers. Previous diagonal RNNs such as eLSTM [17] , Columnar networks [18] , and IndRNN [20] use only real-valued diagonal matrices. Having only real-valued diagonals assumes that the matrix W h is symmetric. We provide a small experiment in Appendix B.2 showing that this assumption does not hold even in the simplest setting and that complex eigenvalues do arise. We also investigate whether this result can be extended beyond linear recurrence, and largely obtain a negative theortical result (see Appendix B.1 and C ).A complex number can be represented in three ways: a + bi (the real representation), r exp(iθ) (the exponential representation), and r(cos(θ) + i sin(θ)) (the cosine representation). Mathematically, these three representations are equivalent, but do they affect learning differently? Orvieto et al. [34] empirically showed that using the exponential representation resulted in a better-behaved loss function than the real representation on a simple task; we provide some discussion in Appendix D.1 further motivating why the real representation is less stable. We chose instead to pursue the cosine representation, because it allows us to represent the complex hidden vector as two real-valued vectors. The remainder of this section outlines RTUs, with and without nonlinearity in the recurrence.Our goal is to learn a complex-valued diagonal matrix with weights λ k = r k (cos(θ k ) + i sin(θ k )) on the diagonal, for k = 1, . . . , n. Multiplying by a complex number is equivalent to multiplying by a 2x2 block matrix with a rescaling. We can use this rotational form to avoid explicitly using complex numbers, and instead use two real-values for each complex-valued hidden node. We write this real-valued matrix Λ ∈ R 2n×2n as blocks of rotation matrices2 EQUATIONEach element of h t = Λh t-1 + W x x t ∈ R 2n has two components h c1 t , h c2 t , updated recursively:h c1 t = r cos(θ θ θ) ⊙ h c1 t-1 -r sin(θ θ θ) ⊙ h c2 t-1 + W c1 x x t , h c2 t = r cos(θ θ θ) ⊙ h c2 t-1 + r sin(θ θ θ) ⊙ h c1 t-1 + W c2 x x t .We finally combine the new recurrent states into one state h t . = [f (h c1 t ); f (h c2 t )], potentially using a non-linearity f after the recurrence.We also adopt two parameterization choices made in LRUs that showed improved performance. The first is learning logarithmic representations of the parameters rather than learning them directly: instead of learning r and θ θ θ, the network learns ν ν ν log and θ θ θ log , where r . = exp(-ν ν ν), ν ν ν = exp(ν ν ν log ), and θ θ θ log .= log(θ θ θ). This re-parametrization restricts the r to be ∈ (0, 1], required for stability. We found these modifications to improve stability of RTUs (see Appendix E). The second parameterization choice we adopt from LRUs is to multiply the input (W x x t ) k by a normalization factor ofγ k = (1 -r 2 k ) 1/2. Putting this all together, the final formulation of RTUs is:EQUATIONh t = [f (h c1 t ); f (h c2 t )],where γ γ γ ∈ R n is the vector composed ofEQUATIONNote that γ γ γ can be absorbed by W, and so does not change representation capacity.There are two ways to incorporate non-linearity into RTUs: inside the recurrence or after the recurrence. In the above, in Equation ( 2), the non-linearity is after the recurrence. These RTUs maintain the equivalence to a dense linear RNN, because the recurrence itself remains linear. We refer to this definition of RTUs as Linear RTUs, because the recurrence is linear, even though we have the ability to represent nonlinear functions by allowing for any nonlinear activation after the recurrence. We also evaluated a different variation of RTUs where the non-linearity is added to the recurrence directly. These Nonlinear RTUs are written as:EQUATIONh t = [h c1 t ; h c2 t ]. Notice now f -a nonlinear activation like ReLU-is used in the update to h c1 t and h c2 t , and the final h t simply stacks these two components. Nonlinear RTUs lose the equivalence to dense RNNs, though in our experiments, we find they perform as well or better than Linear RTUs.This section shows the RTRL updates for RTUs with more in-depth derivations in Appendix E. To keep notation simpler, we write the updates as if we are directly updating r and θ; the updates for ν ν ν log and θ θ θ log are easily obtained then using the chain rule. The full derivation is in Appendix E.2.Consider the partial derivative with respect to r 1 for the first RTU with input x1 . = (W c1 x x t ) 1 :EQUATION.Since r 1 only impacts the two units in the first RTU, and derivatives w.r.t. the remaining hidden units are zero. Therefore, we just need to keep track of the vector of partial derivatives for these two components, e ∂rn ] with recursive formulas: e r,c1 t = cos(θ)⊙ h c1 t-1 +r⊙ cos(θ) ⊙ e r,c1 t-1 -sin(θ)⊙ h c2 t-1 -r ⊙ sin(θ) ⊙ e r,c2 t-1 -r √ 1-r 2 ⊙ W c1 x x t e r,c2 t = cos(θ)⊙ h c2 t-1 +r⊙ cos(θ) ⊙ e r,c2 t-1 + sin(θ) ⊙ h c1 t-1 + r ⊙ sin(θ) ⊙ e r,c1 t-1 -r √ 1-r 2 ⊙ W c2 x x tWe can similarly derive such traces for θ. The update to r involves first computing ∂Lt ∂h c 1 t, using backpropagation to compute gradients back from the output layer to the hidden layer; this step involves no gradients back-in-time. Then r is updated using the gradient ∂Lt ∂hc 1 t ⊙ e r,c1 t + ∂Lt ∂h c 2 t ⊙ e r,c2t , which is linear in the size of r ∈ R n , as the vectors e r,c1 t , e r,c2 t ∈ R n can be updated with linear computation in the above recursion. This update is the RTRL update, with no approximation.RTUs are similar to LRUs, with two small differences. First, RTUs have real-valued hidden units, because the cosine representation is used instead of the exponential representation. Second, RTUs use nonlinear activations in the recurrence, making them no longer linear. Though again a minor difference, we find that incorporating nonlinearity in the recurrence can be beneficial. RTUs can be seen as a small generalization of LRUs, moving away from strict linearity-and thus motivating the name change-but nonetheless a generalization we find performs notably better in practice.Let us now motivate the utility of moving to a cosine representation and real-valued traces. LRUs parameterize each hidden unit with λ k = r k exp(iθ k ) = exp(-exp(ν log k ) + i exp(θ log k )) and directly work with complex numbers. Consequently, the hidden layer cannot be directly used to predict realvalues. It would be biased to take Re(h t ) (see Appendix D.2), and instead an additional weight matrix W ∈ C n×n must be learned, to get Re(W h t ). To understand why this works, assume that we took the original h t from the dense NN, and handed it to an MLP. This would involve multiplying Wh t for some W. If we set W = WP, then W h t = WPP -1 h t = Wh t and we did not introduce any bias. In fact, if W is set this way, we do not need to take the real-valued part, because the output of Wh t is real-valued. Of course, learning does not force this equivalence-in fact this parameterization is more flexible than the original-and so it is necessary to take the real-part.RTUs avoid some of these complications by explicitly writing the recurrence and updates with realvalued hidden states. Implicitly, the relationship between the two real-valued hidden vectors forces them to behave like complex numbers (as rotations), but all equations and learning stay in real-valued space. RTUs consequently avoid the need to post multiply by the matrix, removing a small number of learnable parameters, allowing the use of a nonlinear activation directly on the output, and allowing the hidden state to be immediately passed to a downstream MLP. We acknowledge that others may argue that working directly with complex numbers is preferable. The preference for real-valued hidden layers may simply be our own limitations, but we suspect much of the reinforcement learning community is similarly more comfortable to work in real-valued space. 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 3 .In the end, our goal is to make these simple recurrent traces easy to use, and providing updates with real numbers may remove some barriers.We can enforce r to be ∈ (0, 1] by using a positive non-linear function. For example, learn ν log such that r . = exp(-exp(ν log )), this parameterization ensures that r ∈ (0, 1] and is suggested by Orvieto et al. [34] . Another example is learning r = σ(ν), which also ensures r ∈ (0, 1].Finally, we can also enforce stability on θ by learning θ log such that θ . = exp(θ log ) to ensure that θ is always positive. We empirically compare the different parameterizations of r and θ. In our experiments, learning r directly resulted in unstable training where the MSRE diverges. We plot the learning rate sensitivity for the other parameterization for two different sizes of RTUs in Figures 9 and 10 . While all the parametrizations produce similar performance, we notice that learning θ log and ν log has better learning rate sensitivity.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes]Justification: Details of all proofs are provided in the appendix and referenced in the main paper.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
online prediction learning: 
 In this section, we explore different architectural variants of RTUs and LRUs in a online prediction task and then move on to study the tradeoffs between computational resources and performance when using RTUs with RTRL compared to GRUs and LRUs with T-BPTT.In this first experiment, we investigate the impact of several architectural choices on the performance of RTUs and LRUs varying where nonlinearity is applied. We use a simple multi-step prediction task called Trace conditioning [39] inspired by experiments in animal learning. The agent's objective is to predict a signal-called the Unconditional Stimulus (US)-conditioned on an earlier signal-the Conditional Stimulus (CS). The prediction is formulated as a return, G t . = ∞ k=0 γ k US t+k+1 , where the agent's goal is to estimate the value function for this return. More details on this environment and experimental settings are in Appendix F. Figure 1 summarizes the results.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.In this section, we investigate the tradeoffs between computational resources and performance when using RTUs with RTRL compared to GRUs and LRUs with T-BPTT.In the following experiments, all agents consist of a recurrent layer followed by a linear layer generating the prediction. We measure performance of the agents online by calculating the Mean Square Return Errors (MSRE), which is the mean squared error between the agent's prediction at time t and G t . In all the experiments, we used the Adam optimizer. We first ran each agent with different step sizes for 5 runs and 2 million steps. We then averaged the MSRE over the 2 million steps and selected each agent's best step size value. Finally, we ran the agents with the best step size value for 10 runs, which we report here. We also report all agents' step size sensitivity curves in Appendix F.3. Learning under computational constraints: We first investigate: how well do different agents exploit the available computational resources? We specified a fixed computational budget of 15000 FLOPs. Since RTUs are learning using RTRL and have a linear computational complexity, the computational budget only determines the number of hidden units in the architecture. For GRU and LRU, both the truncation length and the hidden dimension contribute to the budget. We tested several configurations of truncation lengths and parameters such that the overall computations fit the computational budget. Figure 2 .a shows the results of this experiment. As we move along the horizontal axis, the number of parameters for GRU and LRU decreases as T increases to fit the computational constraints. However, the RTU agents do not depend on T , so their performance and computation is constant.Scaling with computation: The computational complexity of T-BPTT depends on the truncation length and the number of parameters in the neural network. Thus, the agent can use the additional resources in two ways: (1) Increasing the truncation length, and (2) Increasing the number of parameters. On the other hand, RTUs use all the computations to have more parameters. Now, we move to our second question: how well do different methods scale with increasing the computational budget? We answer this question in two stages: Firstly, we study T-BPTT with increasing T and a fixed number of parameters. For RTU, the computation increases by adding more parameters such that all corresponding points from GRU and RTU use the same amount of computation. Secondly, we fixed the truncation length for GRU to 45, which is more than the maximum distance between the CS and the US, and increased the computation by increasing the number of parameters for GRU. Again, for RTU, we increased the computation by increasing the number of parameters.. While GRU's performance improved as the truncation length increased, RTU outperformed GRU across all different computational budgets. Figure 2 .c shows the results of the second experiment. The RTU agent's performance consistently improves as we increase the computation available. However, the performance improvement for the GRU agent is inconsistent. The inconsistency of GRU performance could be associated with the trade-off between the truncation length and the number of parameters.Scaling With Parameters: Finally, we study the performance of RTU and GRU when given the same number of parameters and allow the GRU agent to use more computation. We fixed the truncation length for GRU to 45 as before and used the same number of parameters for both agents. Figure 2 .d shows the results of this experiment. For RTU, we see the same consistent performance improvement as we increase the number of parameters. For GRU, the performance improvement is also consistent, though it degrades slightly towards the end. The RTU agent outperforms the GRU agent even though the GRU uses more computation.We provide additional experiments comparing RTUs to two other approaches that use RTRL: online LRUs and a real-valued diagonal RNN in Appendix F.1.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: all hyper-parameters and experiment details are shared in the appendix and referenced in the main paper. 
real-time recurrent policy gradient: 
 This section first highlights some differences in using linear RTRL methods, i.e., RTRL with linear complexity, in incremental and batch settings. We then investigate different ways of integrating linear RTRL methods with policy gradient approaches, and we use PPO as a case study for this investigation. Finally, we compare the performance of RTRL methods with T-BPTT methods and other baselines.The benefits of linear RTRL methods over T-BPTT are more evident in the incremental rather than the batch setting. In the incremental learning setting, where the agent updates its parameters after each interaction step, linear RTRL methods have a constant computational complexity per update that depends only on the number of parameters. In contrast, T-BPTT methods have a complexity proportional to the truncation length T since T-BPTT methods require storing a sequence of past T activations to perform one gradient update. Figure 3 shows the time it takes to make one update with linear RTRL and T-BPTT given the same number of parameters. For T-BPTT, the time to make one update scales with the truncation length T, while for linear RTRL, it is constant.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. The computational analysis for the batch setting is different than the incremental setting. In the batch setting, linear RTRL still have a constant cost per update and provide an untruncated yet stale gradient for all the samples. When using T-BPTT in the batch setting, there are two possibilities for the gradient updates. The first option, the typical strategy, is to divide the batch into non-overlapping sequences, each of length T, and perform T-BPTT on each sequence. In this case, the cost of one gradient update per sequence is a function of T, resulting in an effective update cost per sample independent of T. As a result, T-BPTT is computationally efficient in this case, albeit at the expense of a worse gradient estimate; in each sequence, only the last sample has a gradient estimate with T steps [25] . Figure 3 shows the time it takes to make one batch update with linear RTRL and T-BPTT given the same number of parameters. In this case, both methods use similar time per update. The second option is to divide the batch into overlapping sequences, where each gradient uses a sequence of T steps [25] . This approach increases the cost of updates per sample to be proportional to T, as in the incremental setting, with the benefit of better gradient estimates. However, all standard implementations of RL methods with T-BPTT use the computationally efficient option [38, 14, 24] .Integrating Linear RTRL Methods with PPO When performing batch updates, as with PPO, the RTRL gradients used to update the recurrent parameters will be stale, as they were calculated during the interaction with the environment w.r.t old policy and value parameters. One solution to mitigate the gradient staleness is to go through the whole trajectory after each epoch update and re-compute the gradient traces. However, this can be computationally expensive. In Appendix G, Algorithm 1, we provide the pseudocode for integrating RTRL methods with PPO with optional steps for re-running the network to update the RTRL gradient traces, the value targets, and the advantage estimates. We also performed an ablation study to investigate the effect of the gradient staleness in RTRL when combined with PPO, Appendix G. The results from the ablation study show that using a stale gradient results in better performance with RTUs and suggests that the staleness might help PPO maintain the trust region.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? 
experiments in memory-based control: 
 In this section, we evaluate the memory capabilities of RTUs when solving challenging RL control problems. We divide the problems in this section based on the source of partial observability. 1) Missing sensory data, where we mask out parts of the agent's observation. The agent must accumulate and integrate the sensory observations over time to account for the missing information. 2) Remembering important cues, where the agent must remember an essential cue about the environment that happened many steps in advance.Integrating Sensory Observations: We use the standard Mujoco POMDP benchmark widely used in prior work for evaluating memory-based RL agents [31, 12, 27, 32] . The benchmark consists of several challenging tasks where the agent controls a multi-joint dynamical body while only observing the joints' positional (Mujoco-P) or velocity information (Mujoco-V). To increase experiment throughput, we use the Jax implementation of Mujoco from the Brax library [7] and implemented wrappers to mask either the velocity (Mujoco-P) or positional information (Mujoco-V).We evaluated our Linear and Non-linear RTUs against GRU, LRU, and Online LRU. All agents use PPO [40] as the control algorithm, and the difference between the agents is the recurrent component. For all agents, we fixed the number of parameters for the recurrent part to be ∼ 24k. We tuned the learning rate for all agents in all environments and selected the best learning rate for each agent per environment. We also included a GPT2-transformer baseline. We followed the implementation details in previous work [32] , and used a GPT2 variant with 200k parameters. We provide the results for GPT2 in Appendix H. 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 ( 24k parameter). For each architecture, we show the performance of its best-tuned variant.When given the same number of parameters, RTU agents outperform other baselines in all environments in Figure 4 . Furthermore, we show in Appendix H that even when increasing the truncation length of both GRU and LRU agents to use significantly longer history, they outperform RTUs in only one task. Of particular note is again that RTUs outperform online LRUs, highlighting again that our simple modifications have a large impact on performance in this online RL setting.Remembering Important Cues: Next, we test the agents' ability to remember essential environmental cues. We use several tasks from the POPGym benchmark [29] in addition to the Reacher POMDP task, a modified version of Mujoco Reacher where the agent observes the target position only at the beginning of the episode. The POPGym tasks we consider along with the Reacher POMDP are all long-term memory tasks [32] as the agent must remember and carry the information for a long time.Figure 5 summarizes the results for the reacher POMDP task and the POPGym results can be found in Figure 6 . In both cases, we can see that RTUs outperform the other approaches. Non-linear RTUs achieve a better performance than linear RTUs in reacher POMDP, and both achieve a better performance in all tasks than online LRUs. In Reacher POMDP, GRU was able to achieve a similar performance to that of linear RTUs. Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: all hyper-parameters and experiment details are shared in the appendix and referenced in the main paper.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
conclusion and limitations: 
 In this work, we investigated using complex-valued diagonal RNNs for online RL. We built on LRUs, to provide a small modification (RTUs) that we found performed significantly better in online RL across various partially observable prediction and control settings. We also found RTUs performed better than the more computationally intensive GRUs. Overall, RTUs are a promising, lightweight approach to online learning in partially observable RL environments.A primary limitation of RTUs is the extension to multilayer recurrence. This limitation is inherent to all RTRL approaches; with multilayers, we need to save the gradient traces of the hidden state w.r.t the weights from all the preceding layers [17] . Previous work [17, 48] showed that using stop gradient operations between the layers and not tracing the time dependencies across layers is a viable choice. However, we need a more principled approach for tracing the gradient across layers.One advantage of the linearity restriction in LRUs is that it allows the use of parallel scans for training [26] . However, recent works have shown the possibility of employing parallel scans to non-linear RNNs [8, 22] . A future direction is to investigate the use of parallel scans for training RTUs.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: we report error bars/shaded standard error regions in all our experiments.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: Compute details were provided in the appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.• The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: The research conducted here doesn't include large models or scraped datasets. Guidelines:• The answer NA means that the paper poses no such risks.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: This section provides a brief background on BackPropagation Through Time (BPTT) and Real-Time Recurrent Learning (RTRL) algorithms.BPTT calculates the gradient, ∇ ψ ψ ψ L, by unfolding the recurrent dynamics through time and incorporating the impact of the parameters on the loss from all observed time steps. Formally, we can write ∇ ψ ψ ψ L as:EQUATIONApplying the chain rule, we re-write Eq.5 as:EQUATIONWhen calculating ∂hi ∂ψ ψ ψ , we need to consider the effect of ψ ψ ψ from all the time steps. To illustrate this effect, consider unrolling the last 2 steps of the RNN dynamics:EQUATIONEquation 7 shows that the network parameters ψ ψ ψ affect the construction of the recurrent state h t through two pathways: a direct pathway, i.e., using ψ ψ ψ to evaluate f (h t-1 , x t , ψ ψ ψ), and an implicit pathway, i.e., ψ ψ ψ affected constructing all previous recurrent states, h t-1 , . . . , h 1 , and all those recurrent states affected h t construction. Thus, to calculate ∂ht ∂ψ ψ ψ , we need to consider those two pathways:EQUATION)Once again, we need to consider the two pathways when evaluating ∂ht-1 ∂ψ ψ ψ in 8. For simplicity, letJ t . = ∂ht ∂ψ ψ ψ , B t = ∂f (ht-1,xt,ψ ψ ψ) ∂ψ ψ ψ , C t = ∂f (ht-1,xt,ψ ψ ψ) ∂ht-1, and re-write 8:EQUATIONWriting ∂ht ∂ψ ψ ψ using the results from 9:EQUATIONAccording to Eq. 10, the agent needs to store all the previous inputs to calculate ∂ht ∂ψ ψ ψ which is impractical; the computation and memory complexity will be increasing with t.Williams and Peng (1990) introduced Truncated-BackPropagation Through Time (T-BPTT) which solves the issue of increasing memory and computational complexities of BPTT. In T-BPTT, we specify a truncation length T , which controls the number of steps considered when calculating the gradient in 10. We now write the truncated version of 9 which takes into consideration the gradient from the last T steps only:EQUATIONUsing results from 11, we then write the approximated gradient of the loss w.r.t the learnable parameters:EQUATIONA.2 Real-Time Recurrent LearningWilliams and Zipser (1989) introduced the Real-time Recurrent Learning algorithm (RTRL) as a learning algorithm for continual recurrent learning. RTRL employs the recurrent formulation of the gradient in 8; instead of unrolling ∂ht-1 ∂ψ ψ ψ further back in time, RTRL saves its calculated value from the previous time step and use it later when needed. It is worth emphasizing that after the agent updates its parameters, the gradient information saved from previous time steps would be stale, i.e., calculated w.r.t old parameters, however, under the assumption of small learning rates, RTRL is known to converge. The gradient formulation of RTRL can be written as:EQUATIONThis section explains why we need complex-valued diagonals to represent dense recurrent layers.We first show when it is equivalent to use complex-valued diagonal and a dense recurrent layer. We highlight that using a real-valued diagonal is like restricting the weights to be symmetric-because the (complex) diagonal corresponds to the eigenvalues of the weight matrix-which can severely limit representability. We provide a small experiment to show that complex eigenvalues naturally arise when training both a dense linear and nonlinear RNN, further motivating the utility of moving towards complex-valued diagonals.Let us first consider when we can perfectly represent a dense, linear recurrent layer with a complexvalued diagonal recurrent layer. Assume we have the recurrence relationship, with learnable parametersEQUATION)where f is a potentially nonlinear function that inputs a vector and outputs the same-sized vector and u can be any transformation of the inputs x t before they are inputted into the recurrent layer. The following equivalence result is straightforward but worthwhile formalizing. Proposition B.1. Assume f •P = P•f for any full rank, potentially complex-valued P ∈ C n×n with unit-length column vectors. Then given any W h and W x for Equation ( 14), there is a corresponding complex-valued diagonal weight matrix Λ ∈ C n×n andEQUATION)whereh t ∈ C n is a linear transformation of h t ∈ R n .Proof. We can rewrite the square matrix W h using an eigenvalue decomposition W h = P Λ P -1 , where P contains the n linearly independent eigenvectors and Λ is a diagonal matrix with the corresponding eigenvalues. Then, we can re-write ( 14) as:EQUATIONwhere P came outside of f under our assumption that it commutes with such matrices of eigenvectors. By defining h t . = P -1 h t and W x . = P -1 W x , we get Eq. (15) .We can see h t and h t are representationally equivalent: they are linearly weighted for downstream predictions, and so the linear transformation on h t can fold into this downstream linear weighting. But it is more computationally efficient to use h t with a diagonal weight matrix Λ, meaning each hidden unit only has one recurrent relation instead of n.Since we did not impose constraints on the matrix W h , other than being diagonalizable, the eigenvalues of W h can be complex or real numbers. Previous diagonal RNNs such as eLSTM [17] ,Columnar networks [18] , and IndRNN [20] use only real-valued diagonal matrices. Having only real-valued diagonals implicitly assumes that the matrix W h is a symmetric matrix. [34] suggested using complex-valued diagonal matrices for better performance.The above equivalence has only been used without any activation, namely in linear-recurrent units (LRUs) [34] . A natural question is if only the identity f (x) = x (and linear functions) satisfy this property of commuting with eigenvector matrices. Intuitively, this seems like the only option, as imagining a nonlinearity that commutes is hard. Surprisingly, for the more restricted case of symmetric W h , we can show a slightly more general class of activations can be used, proving an if-and-only-if relationship (see Appendix C). However, even for this restricted setting, this generalized class is limited and such activations unlikely to be preferable to a linear recurrence. We see this as a negative result, that suggests this equivalence only holds for the linear setting.We empirically investigate whether complex eigenvalues appear when training dense RNNs in a simple task. The goal is to show that the weight matrix, W h , is not a symmetric matrix, even in the simplest tasks. Hence, having only real-valued diagonals is too restrictive.We used the Three State POMDP [42] , depicted in Figure 7 , for this experiment. In this task, the agent needs to remember one cue from the previous time-step ago, to make a prediction about the next time-step. The MDP has three states, s 1 , s 2 , and s 3 , and no actions. If the agent is in either s 1 or s 2 , it transitions to any of the three states with equal probability. However, if the agent is in s 3 , it transitions to the state preceded by s 3 . A sequence of observations would look like 1, 3, 1, 2, 2, 3, 2, • • • , and we ask the agent to predict the next observation.Figure 7 : Illustration of the Three State MDP. We used dashed lines for the transitions starting in s 3 to make them more visible.We trained a vanilla RNN with 3 hidden states with T-BPTT with truncation length 2, which is a sufficient history length in this problem to predict the next observation. Since we have 3 hidden states, the matrix W h is ∈ R 3×3 and could have at most 2 complex eigenvalues.We measured the performance in terms of the percentage of correct predictions made in S 3 . We recorded the number of complex eigenvalues of W h after each parameter update, shown in Figure 8 . This agent reaches 100% accuracy in this problem relatively quickly. We can also see that the agent oscillates between having two complex eigenvalues and zero eigenvalues. The average number of complex eigenvalues across 30 run is above 1.5, which means that on more than 3 4 of the steps, the RNN has two complex eigenvalues. The primary point is that we see complex eigenvalues appear frequently. As discussed in the main body, we likely only have an equivalence between using a full dense weight matrix and a complex-valued diagonal matrix for linear recurrent layers. However, we can obtain a slightly more general equivalence in the restricted setting where the weight matrix for the recurrence is symmetric. This restricted setting is not of general interest, but we include the result here because it could be of interest to a few.For the case where we have a symmetric weight matrix, we need the activation to commute with orthonormal matrices. Consider again the formEQUATIONwhere f is a potentially nonlinear function that inputs a vector and outputs the same-sized vector. To obtain the equivalence, assume that for any orthonormal matrix A, f • A = A • f . We can rewrite the square and symmetric matrix W h using an eigenvalue decomposition W h = A Λ A ⊤ , where A contains the n linearly independent eigenvectors and is an orthonormal matrix and Λ is a diagonal matrix with the corresponding eigenvalues. Then, we can re-write ( 14) as:EQUATIONwhere A came outside of f because commutes with orthonormal matrices. By defining h t . = A ⊤ h t and W x . = A ⊤ W x , we get:EQUATIONEach hidden unit now has one recurrent relation instead of n, because our weight matrix Λ is diagonal.A natural question is if only the identity f (x) = x-namely linear recurrence-satisfies this property of commuting with orthonormal matrices. We show below that it holds for a slightly more general class of recurrent layers, proving an if-and-only-if relationship. We see the below result as a negative result, highlighting that this equivalence largely only holds for the linear setting and does not generalize to other activations of interest. It provides even further evidence that likely the only setting of interest for the general non-symmetric case is also with a linear recurrence.Nonetheless, let us obtain the if-and-only-if for completeness. A simple extension that continues to satisfies this property is f (x) = xc(||x|| 2 ) for any c : R → R. We can see that for any orthonormal A, we havef (Ax) = Axc(||Ax|| 2 ) = Axc(||x|| 2 ) = Af (x).This means that we can have activations that rescale the input x depending on the norm of that input.More generally, the activation can involve matrices and rotations. We can define θ(x) = [x U(x)] ⊤ where U(x) ∈ R n×n-1 is a matrix where the columns are orthogonal vectors to each other and to x. Then for any vector-valued g : R → R n , we have that f (x) = θ(x) ⊤ g(||x|| 2 ) also satisfies this property:f (Ax) = θ(Ax) ⊤ g(||Ax|| 2 ) = [Ax; U (Ax)]g(||x|| 2 ) = A[x; U(x)]g(||x|| 2 ) = Af (x).The last line follows because AU(x) = U(Ax). Namely, for any orthogonal vector u with u ⊤ x = 0, we have thatũ . = Au satisfies ũ⊤ Ax = u ⊤ A ⊤ Ax = u ⊤ x = 0 because A ⊤ A = I. Now we show this formally. Denote O n ⊂ R n×n : A ∈ O n ⇐⇒ A T A = AA T = I. Denote e i =        0 . . . 1 . . . 0       , where the i th element is 1.Definition C.1. θ θ θ(x) is a matrix such that:EQUATIONLemma C.2. θ θ θ(x) exits and∀A ∈ O n , θ θ θ(Ax) = θ θ θ(x)A T T θ θ θ(x) = I θ θ θ(Ax) T = θ θ θ(x)A T (19) Theorem C.3. Let f : R n → R n and A ∈ O n . Then, f • A = A • f ⇐⇒ ∃g : R → R n , f (x) = θ(x) T g(∥x∥ 2 )Proof. Define g(α) = f (α 2 e 1 ).EQUATIONEQUATIONIn this section we provide a few additional insights on alternative ways to handle complex numbers within an RNN, and why they are not preferable.We look at the gradient when using each complex representation to understand how different representations affect learning stability. Since each hidden unit has only one recurrent relation in diagonal RNNs, it is sufficient to consider one unit, in isolation. To keep the below intuition simple, we also omit the input of x, and considerEQUATION)where h 0 is the initial hidden state.Real Representation a + bi: Substituting λ in (22) with the real representation, we get:h t = (a + bi) t h 0 = h 0 t k=0 t k a t-k b k i kThen it follows that the gradient w.r.t the learnable parameters a and b is:∂h t ∂a = h 0 t k=0 t k (t -k)a t-k-1 b k i k ∂h t ∂b = h 0 t k=0 t k ka t-k b k-1 i kTo prevent the gradient from vanishing/exploding, we need to restrict both |a| and |b| to be ∈ (0, 1].Exponential Representation r exp(iθ): Substituting λ in (22) with the exponential representation, we get: h t = r t exp(itθ)h 0 and the gradient w.r.t the learnable parameters r and θ is:∂h t ∂r = tr t-1 exp(itθ)h 0 , ∂h t ∂θ = r t exp(itθ)ith 0To prevent the gradient from vanishing/exploding, we need to restrict r ∈ (0, 1].Cosine Representation r(cos(θ) + i sin(θ)): Substituting λ in (22) with the cosine representation, we get: h t = r t (cos(tθ) + i sin(tθ))h 0 and the gradient w.r.t the learnable parameters r and θ is:∂h t ∂r = tr t-1 (cos(tθ) + i sin(tθ))h 0 ∂h t ∂θ = r t (it cos(tθ) -t sin tθ)h 0 .To prevent the gradient from vanishing/exploding, we need to restrict r ∈ (0, 1]. It is simpler to maintain stability with the exponential and cosine representations, since we only need to constrain r ∈ (0, 1]. , whereas the real representation requires us to restrict both the complex number's magnitude and phase.We can attempt to get the benefits of having a real-valued hidden state by simply converting the complex-valued state to a real-valued one within the LRU. However, taking only the real part results in a biased gradient, as we show in this section.Consider again the one recurrent unit example, but now also consider the output obtained by taking only the real part of the recurrent state: y t = wRe{h t }, where w is a learnable parameter. The gradient w.r.t r and θ is:EQUATIONMultiplying by a complex number z is equivalent to a rotation by the matrix Re{z} -Img{z} Img{z} Re{z} and a scale by Re{z} 2 + Img{z} 2 . We re-write the recurrent unit using this property as:EQUATIONNotice that we don't need to take the real part of the recurrent state in this formulation. We now write the gradient using this new formulation:∂y t ∂r = w(r cos(θ)(h c1 t-1 + h c2 t-1 ) + r cos(θ)( ∂h c1 t-1 ∂r + ∂h c2 t-1 ∂r )+ r sin(θ)(h c1 t-1 -h c2 t-1 ) + r sin(θ)(EQUATION))EQUATIONComparing Eq. 25 and Eq. 23, we can see that using only the real part of the recurrent state leads to a loss of information in the gradient.This appendix details the parametrization used for RTUs, the derivation of the RTRL update rules, and the extension of RTUs to multi-layers.Since r represents the magnitude of a complex number, then r ∈ R + , it is preferred to have r ∈ (0, 1] to avoid vanishing/exploding gradients as discussed in the previous section. Let w r be a learnable parameter which could directly represent r or represent a function of r, we can enforce the constraints on r in several ways:1. Direct learning: Learn r directly, and clip it after each parameter's update to be in (0, 1].2. Enforcing r ∈ R + : Learn ν such that r . = exp(-ν). This parameterization enforces r to be positive. However, additional clipping is needed to enforce r ∈ (0, 1].We now outline the details for the RTRL update rules for RTUs. The set of learnable parameters for RTUs isψ ψ ψ . = {ν ν ν log , θ θ θ log , W c1 x , W c2 x }.At each time step t, the learner receives a loss L t (ŷ t , y t ; ψ ψ ψ) where y t is the network output at time t, then the gradient of the loss w.r.t the parameters is:EQUATIONwhere∂h c 1 t ∂ψ ψ ψ = ∂h c 1 t ∂ν ν ν log , ∂h c 1 t ∂θ θ θ log , ∂h c 1 t ∂W W W c 1 x , ∂h c 1 t ∂W W W c 1 x and ∂h c 2 t ∂ψ ψ ψ = ∂h c 2 t ∂ν ν ν log , ∂h c 2 t ∂θ θ θ log , ∂h c 2 t ∂W W W c 2 x , ∂h c 2 t ∂W W W c 2 x .We can derive the following gradients:EQUATIONwhereEQUATIONTo efficiently compute the gradient w.r.t W c1 x and W c2 x , we look at the influence of each when considering a single element from each recurrent state, h c1 t and h c2 t :h c1 t,i = g(ν i , θ i )h c1 t-1,i -ϕ(ν i , θ i )h c2 t-1,i + γ i d j=0 w c1 x,(i,j) x t,j h c2 t,i = g(ν i , θ i )h c2 t-1,i + ϕ(ν i , θ i )h c1 t-1,i + γ i d j=0 w c2x,(i,j) x t,j .We then get:∂h c1 t,i ∂W c1 x,(i,j) = g(ν i , θ i ) ∂h c1 t-1,i ∂W c1x,(i,j)-ϕ(ν i , θ i ) ∂h c2 t-1,i ∂W c1x,(i,j)+ γ i x t,j ∂h c2 t,i ∂W c1 x,(i,j) = g(ν i , θ i ) ∂h c2 t-1,i ∂W c1x,(i,j)+ ϕ(ν i , θ i ) ∂h c1 t-1,i ∂W c1x,(i,j)∂h c1 t,i ∂W c2 x,(i,j) = g(ν i , θ i ) ∂h c1 ∂W c2x,(i,j)-ϕ(ν i , θ i ∂h c2 t-1,i ∂W c2x,(i,j)∂h c2 t,i ∂W c2 x,(i,j) = g(ν i , θ i ) ∂h c2 t-1,i ∂W c2x,(i,j)+ ϕ(ν i , θ i ) ∂h c1 t-1,i ∂W c2x,(i,j)EQUATIONWe see that each h c1 t,i gets affected by weights from only one row of W c1 x , thus,∂h c 1 t W c 1 xcan be written as a matrix of the same dimension as W c1x . The same is true for∂h c 2 t W c 2 x , ∂h c 1 t W c 2 x, and∂h c 2 t W c 1 x.We extend the linear RTUs to non-linear RTUs by adding a non-linear activation function f to the recurrent states. We can write the non-linear RTUs as follows:EQUATIONwhere f : R n → R n . Following the same procedure as in the linear case, we can derive RTRL update rules for the non-linear RTUs.EQUATION)EQUATIONh c1 t,i = f (g(ν i , θ i )h c1 t-1,i -ϕ(ν i , θ i )h c2 t-1,i + γ i d j=0 w c1x,(i,j) x t,j )EQUATIONWe then get:∂h c1 t,i ∂W c1 x,(i,j) = f ′ (•)(g(ν i , θ i ) ∂h c1 t-1,i ∂W c1 x,(i,j) -ϕ(ν i , θ i ) ∂h c2 t-1,i ∂W c1x,(i,j)+ γ i x t,j ) ∂h c2 t,i ∂W c1x,(i,j)= f ′ (•)(g(ν i , θ i ) ∂h c2 t-1,i ∂W c1x,(i,j)+ ϕ(ν i , θ i ) ∂h c1 t-1,i ∂W c1x,(i,j))∂h c1 t,i ∂W c2x,(i,j)= f ′ (•)(g(ν i , θ i ) ∂h c1 t-1,i ∂W c2x,(i,j)-ϕ(ν i , θ i ) ∂h c2 t-1,i ∂W c2x,(i,j))∂h c2 t,i ∂W c2x,(i,j)= f ′ (•)(g(ν i , θ i ) ∂h c2 t-1,i ∂W c2x,(i,j)+ ϕ(ν i , θ i ) ∂h c1 t-1,i ∂W c2x,(i,j)EQUATIONWe now move to calculate the computation and memory complexity of RTUs when learning using the RTRL rules introduced in the previous section.For an input x t ∈ R d and hidden statesh t = [f (h c1 t ); f (h c2 t )] ∈ R 2n ,we have g(ν ν ν, θ θ θ), ϕ ϕ ϕ(ν ν ν, θ θ θ), γ γ γ ∈ R n and W c1 x , W c2 x ∈ R d×n .An agent using the RTU with RTRL needs to store the gradient information,∂h c 1 t-1∂ψ ψ ψ and∂h c 2 t-1∂ψ ψ ψ , from one step to the next. We denote the set of saved gradient information as:EQUATIONThe saved gradient information has the dimensions:EQUATIONThen, it follows that memory complexity for RTU with RTRL is O(n + nd). i.e., linear in the number of parameters.For the computational complexity, a forward pass according to 2 has a computational complexity of O(n + nd). Additionally, after doing the forward pass, the learner needs to update the saved gradient information according to equations 27 through 31 which has a computational complexity of O(n + nd). To summarize, using Real-Time Recurrent Learning with RTUs has linear computational and memory complexities.We now extend RTUs to a multilayer setting. We show that in the multilayer case, we lose the computational advantages. However, prior work suggested that treating each recurrent layer independently is a sensible choice and allows us to gain a computational advantage [18] . Consider an RTU with n layers, we refer to the hidden dimension of a layer i where 0 < i ≤ n as d i . The network has the following set of parameters:EQUATIONTo update the network parameters, we need to calculate the gradient of the loss w.r.t the parameters from all the layers:EQUATIONwhereEQUATIONTake as an example the gradient of h c1,n t w.r.t ν ν ν log,n , ν ν ν log,n-1 , . . . , ν ν ν log,1 . Unrolling the last two layers of the network, we get:EQUATIONThe gradient of h c1,n t w.r.t ν ν ν log,n can be calculated in linear complexity as indicated in the previous section.Calculating the gradient w.r.t the parameters from the earlier layer:EQUATION)∂h c1,n t ∂ν ν ν log,n-2 = g(ν ν ν n , θ θ θ n ) ∂h c1,n t-1 ∂ν ν ν log,n-2 -ϕ ϕ ϕ(ν ν ν n , θ θ θ n ) ∂h c2,n t-1 ∂ν ν ν log,n-2 + γ γ γ n ⊙ W c1,n x ∂h n-1 t ∂ν ν ν log,n-2 = g(ν ν ν n , θ θ θ n ) ∂h c1,n t-1 ∂ν ν ν log,n-2 -ϕ ϕ ϕ(ν ν ν n , θ θ θ n ) ∂h c2,n t-1 ∂ν ν ν log,n-2 + γ γ γ n ⊙ W c1,n x f ′ (•)   g(ν ν ν n-1 , θ θ θ n-1 ) ∂h c 1 ,n-1 t-1 ∂ν ν ν log,n-2 -ϕ ϕ ϕ(ν ν ν n-1 , θ θ θ n-1 ) ∂h c 2 ,n-1 t-1 ∂ν ν ν log,n-2 + γ γ γ n-1 ⊙ W c1,n-1 x ∂h n-2 t ∂ν ν ν log,n-2 g(ν ν ν n-1 , θ θ θ n-1 ) ∂h c 2 ,n-1 t-1 ∂ν ν ν log,n-2 + ϕ ϕ ϕ(ν ν ν n-1 , θ θ θ n-1 ) ∂h c 1 ,n-1 t-1 ∂ν ν ν log,n-2 + γ γ γ n-1 ⊙ W c2,n-1 x ∂h n-2 t ∂ν ν ν log,n-2   ∂h n-2 t ∂ν ν ν log,n-2 ∈ R 2dn-2Can be calulated with linear complexity.EQUATION) Let's define J i,i-1 = { ∂h c 1 ,i t-1 ∂ν ν ν log,i-1 , ∂h c 2 ,i t-1∂ν ν ν log,i-1 }. Then, to calculate the gradient of the hidden units from layer n w.r.t the parameters from layer n -1, we need save J n,n-1 , and to calculate the gradient of the hidden units from layer n w.r.t the parameters of layer n -2, we need to save J n,n-2 and J n-1,n-2 . If we keep going, to calculate the gradient of the hidden units of layer n w.r.t the parameters of the first layer, we need to save J n,1 , J n-1,1 , J n-2,1 , . . . , J 2,1 .For a function f : R n → R m , we have the Jacobian ∂f (x) ∈ R m×n and we can calculate this Jacobian in two ways: forward-mode or reverse-mode differentiation. In the forward-mode differentiation, the chain rule is applied to each operation while traversing the computational graph in the forward pass [2] . While computing the derivatives during the forward pass is appealing, we need to do n forward passes to get the full Jacobian, where each forward pass would allow us to compute the derivative w.r.t only one of the inputs. i.e., with forward mode differentiation, we evaluate the Jacobian one column at a time. As a result, forward-mode differentiation is inefficient for neural networks; neural networks map from learnable parameters, which can be in millions, to a loss function, hence, have very wide jacobians, n ≫ m.Reverse-mode differentiation, on the other hand, offers a more efficient approach. It allows us to evaluate the Jacobian one row at a time, which is particularly advantageous for neural networks. This efficiency comes at the cost of two passes through the network: a forward pass for function evaluation and a backward pass for derivative evaluation [2] .RTRL is an instance of forward-mode differentiation; during the forward pass, the gradient information is evaluated along with the recurrent function computation. As a result, there is no need to perform a backward pass for the recurrent component. To efficiently use a recurrent layer with RTRL within a larger neural network, we combine RTRL for the recurrent layer with the reverse mode for the rest of the network. We use a stop gradient operation on the recurrent layer hidden state and do a normal reverse-mode differentiation. Due to the stop gradient operation, the gradient from the reverse mode assumes no time dependencies between the recurrent states. We then use the gradient traces calculated during the forward pass of the recurrent layer to correct the gradient from the reverse mode and account for the time dependencies between the recurrent states [3] . 4In animal learning, Trace conditioning is a type of experiment where animals predict the occurrence of a stimulus (e.g., food), based on the occurrence of another stimulus like a tone. There is no prior connection between the two stimuli. However, after enough repetitions of pairing them togetherplaying the tone and then serving the food-the animal learns to anticipate food arrival when it hears the tone [36] . We use an open-source trace conditioning benchmark introduced in prior work [39] . Two signals appear sequentially: the Conditional Stimulus (CS) and the Unconditional Stimulus (US). The CS is the trigger signal, similar to the tone, and the US is the signal of interest and appears several time steps after the CS, similar to the food. The agent also observes several distractor signals which are uncorrelated with the CS and US; the agent must learn ignore them and focus only on predicting the US. The agent's objective is to predict the onset of the US, which we model as a prediction of the discounted sum of the future US, G t :EQUATIONwhere γ is a discount factor determining the prediction horizon. This problem is challenging because the CS appears, then disappears, and sometime later the US appears; the agent must construct an internal state that represents the time period between the two signals. Comparison to Other RTRL-based Architectures: We now compare RTUs to other RTRL-based approaches with similar architectures: an online version of LRU [48] and a vanilla block diagonal RNN. The block diagonal is a recurrent formulation similar to RTU but ignores the relation between the learnable parameters. i.e., replaces c k in 3.2 with c k = a k b k c k d k .The results in figure 11 indicate that these seemingly small differences between the diagonal RNNs can result in significantly different behavior. RTUs outperform online LRUs, with the differences discussed in-depth in Section 3.4. RTUs also outperform the block diagonal RNN. We emphasized using real-valued diagonals implicitly assumes symmetric matrices, but that is for a single real-value. This block diagonal has more representational capacity than the RTU. This result suggests it is beneficial for learning to enforce these constraints on the learnable parameters, that they correspond to the rotational representation of complex numbers.On the role of RTRL in RTUs: To highlight the role of RTRL in RTUs, we evaluated the performance of both linear and non-linear RTUs with T-BPTT. In this experiment, we all agents use the same number parameters; the only difference is whether they use RTRL or BPTT.Figure 12 summarises the results of this experiment. We can see that the performance of T-BPTT approaches the performance of RTRL as the truncation length increases to cover the whole context of the task. All agents have one recurrent layer, either an RTU or a GRU, and one linear layer. At each time step t, the agent passes the observation o t to the recurrent layer, which outputs the recurrent state, the agent state. The recurrent state is then passed to the linear layer generating the prediction.For each agent, we swept over the learning rate α used to update the network parameters, α ∈ {10 -1 , 10 -2 , 10 -3 , 10 -4 , 10 -5 , 10 -6 }, and averaged the performance for each learning rate over 5 independent runs. We then selected the best-performing learning rate for each agent and ran 30 independent runs using it. For all the experiments, we ran the agents for 2 million steps, and the performance was the mean squared prediction error averaged over the 2 million steps.We show the learning rate sensitivity for all agents in the animal learning benchmark. When performing batch updates, as with PPO, the RTRL gradients used to update the recurrent parameters will be stale, as they were calculated during the interaction with the environment w.r.t old policy and value parameters. One solution to mitigate the gradient staleness is to go through the whole trajectory after each epoch update and re-compute the gradient traces. However, this can be computationally expensive. In Algorithm 1, we provide the pseudocode for integrating RTRL methods with PPO with optional steps for re-running the network to update the RTRL gradient traces, the value targets, and the advantage estimates. In the next experiment, we investigate the effect of the gradient staleness in RTRL when combined with PPO and how this staleness interacts with the targets and advantage estimates. To understand this interaction, we evaluate all combinations of stale gradient and stale targets with increasing the number of epoch updates. We perform this analysis on the Ant-P environment from the Mujoco POMDP benchmark [31, 12, 27, 32] . Surprisingly, Figure 18 shows that using a stale gradient results in better performance with RTUs than re-computing the gradient traces. This performance improvement is also consistent when we increase the number of epochs from 4 to 8. It also shows that re-computing the value targets and advantage estimates has a minimal effect on the performance. We repeated the same experiments for NonLinear RTUs and Online LRU with consistent results in figures 19 and 20.One hypothesis for the superior performance of stale gradients is that the staleness is helping PPO maintain the trust region. We investigate this hypothesis by measuring the KL divergence between the policy used to collect the trajectory and the most recent policy. We use the following estimate for the KL divergence between the two policies as (r -1) -log(r), where r = (π θnew )/(π θ old ). The rightmost subplot of Figure 18 shows that at the beginning of learning, agents with stale gradients move away from the old policy more slowly than agents with fresh gradients; perhaps stale gradients help with maintaining the trust region. However, this hypothesis still needs more investigation in future work. In our implementations, we use a shared representation learning network followed by two MLPs for the actor and the critic heads, as illustrated in Figure 21 . The shared representations consist of a feedforward layer with 64 hidden units and memory components: an RTU, LRU, or a GRU. The actor and the critic's heads consist of two feedforward layers with tanh activation function.In Figures 22 and 23 , we set the truncation length for GRU and LRU to be 64, which is larger than needed to solve the Mujoco POMDP tasks. These results show that even when the truncation length is larger than needed, RTUs still outperform T-BPTT baselines. We also show that the transformerbased models, GPT2, perform worse than RNN-based models. This is consistent with previous work suggesting that transformers might not be suitable for RL tasks [32] .• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.• If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.Justification: A link for a public repo is provided.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.Code available at https://github.com/esraaelelimy/rtus 38th Conference on Neural Information Processing Systems (NeurIPS2024).We assume the matrix Λ has only complex eigenvalues, as the network can easily turn a complex eigenvalue into a real one by setting the imaginary component to 0.Autodiff can give unexpected results when dealing with complex numbers. For example, see the discussion https://github.com/google/jax/discussions/6817.This can be implemented by defining a custom vjp for the recurrent layer, which modifies the backward pass for the recurrent layer to include the gradient traces. 


Decision: Overall the reviewers agreed that the paper makes a nice contribution to deep architectures for reinforcement learning. Given the success of GRUs in the deep RL literature it is impressive to see this level of improvement. 

The reviewers noted a number of issues that can use more clarification in the final paper and the authors should work hard to address them.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: •	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: 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?
limitations: 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?
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: 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 ...
paper_summary: paper_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.
strength_weakness: strength_weakness: **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: 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.
limitations: limitations: Comments are already provided above.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: N/A
questions: questions: N/A
limitations: limitations: The article discusses the limitations of the method in detail, which I think is reasonable.

Comments: title: Thank you for the response
comment: Thank you for the response and for updating the score.
Thank you for the response and feedback.

> When only using the real part I expect my gradients to be biased towards the real part because this is the only quantity I am considering after all.

That is correct. We were unclear in our initial response, and we meant that though a reader *might* think we can just take the gradient of the real part, this would not be the same as the actual gradient, and this is what we meant by bias. We realize this word is really not the right word, and we should simply say it would be incorrect to do that to get the original behavior, as taking the real part discards some information.
I appreciate the detailed response to my questions. The authors did clarify my concerns, which is why I will raise my score to weak accept.
> tuned the context length and the learning rate for GPT-2.

Thank you. I had overlooked this too as I could not find any description regarding this in any part of the text.

But then this is one more reason not to call those models GPT-2. I am not aware of any GPT-2 architectures with a context length size of 64. I do not understand what makes this model a GPT-2.

Overall, given that there are prior works successfully applying Transformers to RL settings (e.g., Parisotto et al. [25]); I would not find these Transformer experiments very informative, unless there were more in-depth studies. 

I acknowledge the author's promise to modify the corresponding claims. Please ensure these modifications are implemented.
We thank the reviewer for the detailed response and feedback. 

We agree that the GPT-2 results are not core to the paper, and we will modify the main text not to emphasize those results. We also wanted to point out that we tuned the context length and the learning rate for GPT-2. In Figure 26 (or Figure 2 in the uploaded pdf), we show two variants of GPT-2, with context lengths of 1024 and 64. Both of these variants should have enough context length for the Mujoco-p benchmark. Nevertheless, we agree with the reviewer, and we will change the main text accordingly. 

> 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.

Thank you for the suggestion. We will add a couple of sentences to clarify this part.
Thank you for the response and for pointing to some literature on parallelizing non-linear RNNs.

> 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. 

I feel like showing that the model performance increases significantly when trained using RTRL over when trained with BPTT is the main contribution of this paper. Paired with the timing comparison of RTRL and TBPTT over different truncation horizons, this would make a great point and I lament that it did not make it into the main part of the manuscript.

> This makes the gradient estimation different from the true gradient and biased towards the information coming from the real components.

I am still not convinzed. When only using the real part I expect my gradients to be _biased_ towards the real part because this is the only quantity I am considering after all. Whereas, when using both the real and the imaginary part, I also expect the gradients to contain some portion corresponding to the imaginary part. Comparing the two equations in Appendix D2 seems to be misguided. 

Nonetheless, I still feel comfortable with keeping my current rating.
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.