Real-Time Recurrent Learning using Trace Units in Reinforcement Learning
Authors: Esraa Elelimy
Adam White
Michael Bowling
Martha White
Venue: NeurIPS 2024 poster (2024)
Keywords: Reinforcement Learning
RNNs
Primary area: reinforcement_learning

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.

Real-Time Recurrent Learning using Trace Units in Reinforcement Learning
Esraa Elelimy, Adam White", Michael Bowling", Martha White"
University of Alberta, Alberta Machine Intelligence Institute (Amii)
*Canada CIFAR AI Chair
elel amw8, mbbwiing, whitembualberta.ccm
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 train RNNs with real-time recurrent learning (RTRL); unfortunately, RTRL is prohibitively expensive for standard RNNs. A promising direction is to linear recurrent architectures (LRUs), where dense recurrent weights are replaced with a complex-va efficient. In this work, build these insights to provide lightweight but effective approach for training RNNs in online RL. We introduce Recurrent Trace Units (RTUs), small modification on LRUs that nonetheless find have significant performance benefits over LRUs when trained with RTRL. We find RTUs significantly outperform other recurrent architectures across several partially observable environments while using significantly less computation."
diagonal, making RTRL
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 self- attention 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 k2, 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 the whole sequence of observations before taking an action updating the learnable parameters, which is impractical in continual learning. While recent works have reduced the complexity of transformers from quadratic the sequence length to linear [15, 16, 37], the entire sequence length is still needed to train such architectures. Gated Transformer-XX XL attempts to overcome this issue by keeping moving window of previous observations [35]. window of past observations does not scale well to long sequences theeees---eeeeeeeeeee-eeeeeeee computation is quadratic in the sequence length-ann and a window is one particular fixed function of history. The simpler recursive form in RNNs, on the other hand, can learn great variety of functions of history from the data and is well
Code available at https:://githen httpttttuu.cmoooeeetettetee htt.:/aaw..ccomeeesreeeeeaat//aa
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

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 NoBack [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 RI.. However, the approximation resulted in high variance due to added stochasticity to the policy [1].
Methods in the second category usually restrict the RNN architecture 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-valuee vaauued 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 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 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.
2 Background
We formalize the problem setting as a Partially Observable Markov Decision Process (POMDP). At each time step
1,....... the agent perceives an observation Xt, limited view of the state E S, and takes an action A, E A(s:). Depending on the action taken, the agent finds itself in a new state $1+1 E S, observes the corresponding observation X;+1 and a reward Re41 E 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 eegermenvvronmmet interactions that are useful for prediction and control [41]. We could use the whole history up to f, namely (Xo: A1, R1, x1, A2, Rg....x,, 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 concise representation. We refer to the agent's internal representation of the history at time i as its agent state or its hidden state h, The agent constructs its current agent state h, E RR" from its previous agent state h E IR" and the recent observation x, E Rdd using a state-update function g: IR" x Fd R": h g(h;-1,x1).
2

---

One way to learn this state-update function is with a recurrent neural network (RNN). A simple form is linear recurrent layer, where 1, x,) = ,,,, + W for weight matrices W, and W h- We can also add nonlinear activation, such as ReLUU. In general, we will write
h = g(: 1b),
where & are the learnable parameters in the network. The agent maps the agent state h, an output de and then receives a loss Lt = [(t yt) indicating how far the output is from a target yt- 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 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 's ability to look further back in time and its compute and memory requirements. Picking large T be expensive, or require us limit the network size, but picking too small of a 'I' 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... 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 a catch: its computational complexity is quartic, of fourth order, in the size of h, which can be prohibitively expensive. For this reason, we pursue restricted diagonal form in this work, for which RTRL is efficient and linear h,
3 Recurrent Trace Units
In this section, we introduce Recurrent Trace Units (RTUs). We start by revisiting why complex- valued 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 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.
3.1 Revisiting Complex-valued Diagonal Recurrence
Assume we have the recurrence relationship, with learnable parameters W, E Ra and W. E Raxd, h, = W,,,- + W uxu(.., where they are inputted into the recurrent layer. We can rewrite the square matrix W,, using an eigenvalue decomposition W. P A P-1, where P contains the n linearly independent eigenvectors and A E Cax" is a diagonal matrix with the corresponding cigenvalues. Then we have that
be any transformation of the inputs x; before
h, = P(AP-1 h + P
W₂ u(x:))
= P h = AP-1 h + P W. uxx:)
By defining h P-1 h, E C" and W. - P W. e Caxd, we get a new recurrence h = Ah; + Wxu(x:).
We see h, and h, are representationally equivalent: they are linearly weighted for downstream predictions, and so the linear transformation h, fold into this downstream linear weighting. But it is more computationally efficient to use h, with a diagonal weight matrix A, meaning each hidden unit only has one recurrent relation instead of n. LRUs precisely leverage this equivalence [34]. Specifically, they learn a complex-valuee h,, and use Re( W h2) as an input to an MI.P for downstream nonlincarity.
3

---

Since we did not impose constraints on the matrix Wt. other than being diagonalizable, the eigenval- ues of W can 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. is symmetric. We provide 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 extended beyond linear recurrence, and largely obtain a negative theortical result (see Appendix B.1 and C).
3.2 The RTU Parameterization
complex number can represented three ways: a + hi (the real representation), r rex((if) (the exponential representation), and r(cos(0) + (0)) (the cosine representation). Mathematically, these three representations are equivalent, but do they affect learning differently? Orvieto al. [34] empirically showed that using the exponential representation resulted in 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 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 complex- -aauued valued diagonal matrix with weights A = ry "cooss00 (0x) + sin)))))
on the diagonal, for k = -,...., Multiplying by complex number is equivalent to multiplying by 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 E IR 222x2 as blocks of rotation matrices?
C1
[cos(0s) - -iin(0+) (0)) cOs(Os)
A
where
C₁ r
(I)
Cn
Each element of hr Ah;-1
W, x Xx E 2 has two components h h updated recursively;
h = r cos
h
C h + W X,: s0)
(i
h = r cos
h C2 1-1
+
) @ h + W Xx-
We finally combine the new recurrent states into one state h, [[(h])); ((h;))) potentially using 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: exp(y and glor l0g. This re- re-parametrization restricts the to be E (0. 11. 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 (W,).), by normalization factor of ?
instead of learning r and 0, the network learns poo and glog where r = exp(-v).
r2) Putting this all together, the final formulation of RTUs is:
h = g(y log 0 log
@ h 00plog, glog) + h + W
h gov glog h + 6i gloR) O h +7 W
(2)
h [f(h:t); f(h"))]
where 7 E IR" is the vector composed of 7 = - (1-exp(- (esp(-eap) (x () and
g(vn, 0x) exp( exp(x )) cos (os(exp( (f1
(3)
0(v2, 82) exp( exp(1 // sin exp g
Note that y 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 after the recurrence. In the above, in Equation (2), the nin-linyarity is after the recurrence. These RTUs
We assume the matrix A has only complex eigenvalues, as the network can easily turn complex eigenvalue into real one by setting the imaginary component to 0.

---

maintain the equivalence to dense linear RNN, because the recurrence itself remains linear. We refer to this definition RTUs Linear RTUs, because the recurrence 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 added the recurrence directly. These Nonlinear RTUs are written as:
h f(g(u log o
o h
$(v OloR
+7 O Wx CD
h = f(g(v jog 0
@ h + + og Blog)
h +Y O W
(4)
h, [ba] h
Notice now f-a nonlinear activation like ReLU-is used in the update to h," and h and the final h 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.
3.3 The RTRL Update for 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 0; the updates for ylog and goos are easily obtained then using the chain rule. The full derivation is in Appendix E.2.
Consider the partial derivative with respect to T1 for the first RTU with input σ1 (W2x)) xt)
1 = 11 cos(() (0)) - (in( + v(1 r 1-
Then
ac. 0L. Th dr, Th dr, 3i dr
+ oc, ah
Since Tı 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, cf = [ ( 0h2)
047
Srx and c
with recursive formulas:
Dr:
dr.
or,
e, = cos C h tro cos 2 e
sin( @ h r sin( C e
W X
VI-
c, c003 c h tro cos( o 5 + sin @ h r sin c
W
VI-
We can similarly derive such traces for 0. The update to r involves first computing ah, backpropagation to compute gradients back from the output layer to the hidden layer; this step involves gradients back-in-ting Then is updated using the gradient DC, Sh, e + ac, Dh,
using
@ o
which is linear in the size of E I,, the vectors e , E E can be updated with linear computation the above recursion. This update the RTRL update, with no approximation.
3.4 Contrasting LRUS
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 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 A exp(i0x) = exp(- exp(-exp(r) log) ) +iexp work with complex numbers. Consequently, the hidden layer cannot be directly used to predict real- values. It would be biased to take Re(h)) (see Appendix D.2), and instead an additional weight matrix W C must be learned, to get Re( W h.). To understand why this works, assume that we took the original h, from the dense NN, and handed it to an MLP. This would involve multiplying Wh, for some W. If we set W = WP, then W h = WPP h, = Wh, 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 fact this parameterization more flexible than the original- and it is necessary to take the real-part.
exp(1
and directly
Wh, real-valued. Of course, learning does not force this equivalence
5

---

RTUs avoid some of these complications by explicitly writing the recurrence and updates with real- valued 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 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 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 implementation for LRU did not always behave as expected, partially due to how uuttodifferreniittion is implemented in packages for complex numbers 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.
4 Online Prediction Learning
In this section, we explore different architectural variants of RTUs and LRUs in online prediction task and then move on to study the tradeoff's between computational resources and performance when using RTUs with RTRL compared to GRUs and LRUs with T-BPTT.
4.1 Ablation Study on Architectural Choices for RTUs and LRUs
In this first experiment, we investigate the impact of several architectural choices 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 's objective is to predict a signal signal-called the Unconditional Stimulus (US)-conéae conditioned on an earlier signal- the Conditional Stimulus (CS). The prediction is formulated
return, G, = =k= *US+++++ where the agent's goal is to estimate the value function for this return. More details on this environment and experimental settings in Appendix F. Figure summarizes the results.
MBRE 15 run)) Linea Recurence
Uneer Recurrence
NonLinear Recurrence
Unear Returrence Noniiinar Projection
Uieer Recumence
URU
0.10
6.16-
0.36
0.2
0.1-
MTU
0.08
32B
Hidden Units
Hidden
Hidden
Hidden Units
Hidden Units
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, restrict both architectures in 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.
4.2 Learning under resources constraints
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 recurrent layer followed by 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 's prediction at time t and G.. 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 is 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.
2 'Autodiff' can give unexpected results when dealing with complex numbers. For example, see the discussion https:/ htops://giii //tpt/bggttttttt/// www/googleee/ooooooooooooooooooo jax/ wtt.d/scussooo//ee6. /6817.

---

Compucatiio Opuon
(33
Condiin enC (154 (PS)
0.30
Option
0.50
Soal Fig Paramatera
(R)
LRU LT
IRU
MoAe 1.5
033
GIU, T-46
Urear KTll
Nov
13
Trncaloa erugh
R.OM
Approsttoon Compure in Oee (x11
number of Poramderr (4:27))
Figure 2: Learning under resources constraints in Trace Conditioning. Each of the four subplots shows how each algorithm's 's performance varies function of resources. (a) LRU and GRU with T-TBTT is not competitive with RTUs even as is increased while restricting the number of hidden units in LRU and GRU that all algorithms use about the same computation per step. (b) If we allow GRU and LRU's computation to increase (fixed network size) while increasing T, the performance gap remains. (c) Fixing T to large value to solve the task, we can increase the number of parameters, holding the computation equal for all methods. (d) If we do not require compute to be equal across methods as we scale parameters, then the LRU can eventually match the error of RTU, but GRU cannot. The black dashed line represents the near perfect prediction performance.
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 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 TT' 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, 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 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,
Figure of 2.b shows the first experiment's results. 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.
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.
Scaling With Parameters:
We provide additional experiments comparing RTUs to two other approaches that use RTRL: online LRUs and a real-valued diagonal RNN in Appendix F.I.
7

---

5 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.
5.1 Linear RTRL Methods in Incremental and Batch Settings
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 the number of parameters. In contrast, T-BPTT methods have complexity proportional to the truncation length since T-BPTT methods require storing a sequence of past activations to perform one gradient update. Figure 3 shows the time it takes make 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.
Thee by
Incremental Serting
Sutting (MG)
Figure Contrasting runtime in incremental and batch settings. In the incremental set- ting, evaluated in the amimal-learning predic- tion task, T-BPTT updates scale with trunca- tion 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.
0.05
GRU
12.0
0.04
12.0
0.02
11.2
0.00
255
1024
Lengrh
Truncation Length
The computational analysis for the batch setting is different than the incremental setting. In the batch setting, linear RTRL still have 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, 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 worse gradient estimate; in each sequence, only the last sample has a gradient estimate with 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.
this case, both methods use similar time per update. The second option is to divide the batch into overlapping sequences, where each gradient uses 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 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.
6 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
8

---

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 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) V).
We evaluated our Linear and Non-liness RTUs against GRU, LRU, and Online LRU. All agents 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 GPT2-transforme 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.
Unduscentted 150 runs)
Vozo GAU
1250
LonLincor RTU
Uinor RT1
1500
NumLins NT
Hoppar-P
Liverr RTU
3000
KonLineer RTLL
Chestan-
Unwar STU
Morti Liww RTT
1200
GRU
2000
2000
750
500
1000
On LRU
230
bire Stapps
1e6
Ery Steps
Lc6
Ery Stage
Erre Staps
Je6
Indeerounted (30 runs]
Hoppentl
1.000
4000
G00
750
1800
2000
800
400
Dnn Steps
Jab
Ery Steps
1.0
Zre Stops
200
Ere Stops
Jeb
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 basclines in all environ- ments 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 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 In both cases, we can see that RTUs outperform the other ap- proaches. Non-linear RTUs achieve better performance than linear RTUs in reacher POMDP, and both achieve better performance in all tasks than online LRUs. In Reacher POMDP, GRU was able to achieve a similar performance to that of linear RTUs
-308
Figure 5: Reeachr, 30 runs with standard errors.

---

-0854 hortieer dlis
La
3.0
400
-35
One I0p3
305
See
Eny Stape
Dw Stage
One Dapt
Figure 6: Results across several tasks from the POPGym benchmark.
7 Conclusion and Limitations
In this work, we investigated using complex-vai valuee diagonal RNNs for online RL. We built on LRUs, to provide 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 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 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 allows the 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 to investigate the use of parallel scans for training RTUs
8 Acknowledgments
We would like thank Nicolas Zucchet for advice about the online LRU implementation, and Subhojeet Pramanik for many discussions on transformers and RNNs. We would like to thank NSERC, CIFAR and Amii for research funding and the Digital Research Alliance of Canada for the computational resources.
References
[1] Firas Al-Hafez, Guoping Zhao, Jan Peters, and Davide Tateo. Time-efficient reinforcement learning with stochastic stateful policies. In The Twelfth International Conference on Learning Representations.
[2] Atilim Gunes Baydin, Barak Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: survey. Journal of machine learning research, 2018.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Minnh and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
[4] Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, Yoshua Bengio. the properties of neural machine translation: Encoder-decoder approaches. arXiv preprint #rXiv:1409.122 1259, 2014.
[5] Tim Cooijmans and James Martens. On the variance of unbiased online recurrent optimization. arXiv preprint arXiv:1902.0 024000 2019.
[6] Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-leamer architectures. In International conference on machine learning, 2018
[7] C. Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax differentiable physics engine for large scale rigid body simulation, 2021.
[8] Xavier Gonzalez, Andrew Warrington, Jimmy TH Smith, and Scott W Linderman. Towards scalable and stable parallelization of nonlinear mas. arXiv e-prints, pages arXiv-2400 2024.
[9] Gori, Bengio, and De Mori. Bps: learning algorithm for capturing the dynamic nature of speech. In International Joint Conference Neural Networks, 1989.
10

---

[10] Audrumas Gruslys, Will Dabney, Mohammad Gheshlaghi Azar, Bilal Piot, Mare Bellemare. and Remi Munos. The reactor: A fast and sample-efficient actor-critic agent for reinforcement learning. In Interna- tional Conference on Learning Representations, 2018.
[11] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04D0 2023.
[12] Dongqi Han, Kenji Doya, and Jun Tani. Variational recurrent models for solving partially observable control tasks. International Conference on Learning Representations, 2019.
[13] Matthew Hausknecht and Peter Stone. Deep recurrent q-learning for partially observable mdps. In 2015 auai fall symposium series, 2015.
[14] Shengyi Huang, Rousslan Fernand Julien Dossa, Chang Ye, Jeff Braga, Dipam Chakraborty, Kinal Mehta, and JoAGo GM AraAšjo. Cleanrl: High-quality single-fil file implementations of deep reinforcement learning algorithms. Journal of Machine Learning Research, 2022.
[15] Kazuki Irie, Imanol Schlag, Róbert Csordás, and Jürgen Schmidhuber. Going beyond linear transformers with recurrent fast weight programmers. Advances in neural information processing systems, 2021.
[16] Kazuki Trie, Imanol Schlag, Róbert Csordás, and Jürgen Schmidhuher. A modern self-referential weight matrix that learns to modify itself. In International Conference on Machine Learning, 2022.
[17] Kazuki Irie, Anand Gopalakrishnan, and Jürgen Schmidhuber. Exploring the promise and limits real-time recurrent learning, 2023.
[18] Khurram Javed, Haseeb Shah, Rich Sutton, and Martha White. Online real-time recurrent learning using sparse connections and selective learning. Journal of Machine Learning Research, 2023.
[19] Steven Kapturowski, Georg Ostrovski, Will Dabney, John Quan, and Remi Munos. Recurrent experience replay in distributed reinforcement learning. In International Conference on Learning Representations, 2019.
[20] Shuai Li, Wanging Li, Chris Cook, Ce Zhu, and Yanbo Gao. Independently Recurrent Neural Network (In- dRNN); Building A Longer and Deeper RNN. In Conference Computer Vision and Pattern Recognition, 2018.
[21] Xiujun Li, Lihong Li, Jianfeng Gao, Xiaodong He, Jianshu Chen, Li Deng, and Ji He. Recurrent reinforcement learning: hybrid approach. arXiv preprint arXiv:1509.03044, 2015.
[22] Yi Heng Lim, Qi Zhu, Joshua Selfridge, and Muhammad Firmansyah Kasim. Parallelizing non-linear sequential models over the sequence length. In The Twelfth International Conference on Learning Representations, 2023.
[23] Chenhao Lu, Ruizhe Shi, Yuyao Liu, Kaizhe Hu, Simon Shaolei Du, and Huazhe Xu. Rethinking transformers in solving pomdps. In Forty-first International Conference on Machine Learning.
[24] Chris Lu, Jakub Kuba, Alistair Letcher, Luke Metz, Christian Schroeder de Witt, and Jakoh Foerster. Discovered policy optimisation. Advances in Neural Information Processing Systems, 2022.
[25] Owen Marschall, Kyunghyun Cho, and Cristina Savin. unified framework online learning algorithms for training recurrent neural networks. Journal of Machine Learning Research, 2020.
[26] Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. In Interna- tional Conference on Learning Representations, 2018.
[27] Lingheng Meng. Rob Gorbet, and Dana Kulić. Memory- deep reinforcement learning for pomdps. In International Conference on Intelligent Robots and Systems (IROS), 2021.
[28] Jacob Menick, Erich Elsen, Utku Evei, Simon Osindero, Karen Simonyan, and Alex Graves. Practical real time recurrent learning with a sparse approximation. In International Conference on Learning Representations, 2021. 2021..
[29] Steven Morad, Ryan Kortvelesy, Matteo Bettini. Stephan Liwicki. and Amanda Prorok. POPGym: Benchmarking partially obscrvable reinforcement learning. In International Conference on Learning Representations, 2023.
[30] Michael Mozer. A focused backpropagation algorithm for temporal pattern recognition. Complex Systems. 1989.
11

---

[31] Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhutdinov. Recurrent model-free rl can be strong baseline for many pomdps. In International Conference Machine Learning, 2022.
[32] Tianwei Ni, Michel Ma, Benjamin Eysenbach, and Pierre-Lue Bacon. When do transformers shine in rI? decoupling memory from credit assignment. In Thirty-sevemth Conference Neural Information Processing Systems, 2023.
[33] Yann Ollivier, Corentin Tallec, and Guillaume Charpiat. Training recurrent networks online without backtracking, 2015.
[34] Antonio Orvieto, Samuel Smith, Albert Gu. Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks long sequences, 2023.
[35] Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg. Raphael Lopez Kaufman, Aidan Clark, Seb Noury, Stabilizing transformers for reinforce- ment learning. In International conference on machine learning, 2020.
[36] Ivan Pavlov. Conditioned reflexes: an investigation of the physiological activity the cerebral cortex, 1927.
[37] Subhojeet Pramanik, Esraa Elelimy, Marlos C Machado, and Adam White. Recurrent linear transformers. arXiv preprint arXiv:2310.1577 2023.
[38] Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-haselines:: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 2021.
[39] Banafsheh Rafiee, Zaheer Abbas, Sina Ghiassian, Raksha Kumaraswamy, Richard Sutton, Elliot Ludvig, and Adam White. From eye-blinks to state construction: Diagnostic benchmarks for online representation learning. Adaptive Behavior, 2020.
[40] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017.
[41] Sutton. Toward new approach to model-based reinforcement learning, 2020.
[42] Sutton. Markov and agent state, 2020.
[43] Corentin Tallec and Yann Ollivier. Unbiased online recurrent optimization, 2017.
[44] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances neural information processing systems, 2017.
[45] Williams and Peng. An efficient gradient-based algorithm for on-line training of recurrent network trajectories. Neural computation, 1990.
[46] Williams and Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural computation, 1989.
[47] Ailing Zeng, Muxi Chen, Lei Zhang, and Qiang Xu. Are transformers effective for time series forecasting?, 2022.
[48] Nicolas Zucchet, Robert Meier, Simon Schug, Asier Mujika, and João Sacramento. Online learning of long range dependencies. In Advances in Neural Information Processing Systems, 2023.
12

---

A Background on BackPropagation Through Time and Real-Time Recurrent Learning
This section provides a brief background on BackPropagation Through Time (BPTT) and Real-Time Recurrent Learning (RTRL) algorithms.
A.1 BackPropaggttion Through Time
BPTT calculates the gradient, V&C. by unfolding the recurrent dynamics through time and incorpo- rating the impact of the parameters the loss from all observed time steps. Formally, write V as:
VoC
VoL.
V
(5)
1=
Applying the chain rule, re-write Eq.5
V
ViLi
(6)
OL Oh.
Uh Önp
0h, we need to consider the effect of &b from all the time steps. To illustrate this When calculating 86 effect, consider unrolling the last steps of the RNN dynamics:
h. f(h,1, Xx: +b)
Re-write h;- as f(, x;
f(f(hr-2, Xt- 1b), Xr: vb)
Re-write h; as f(h, X 1/)
=f(E([(h-=3, 2 vb), xr Xx, $)
Equation 7 shows that the network parameters & affect the construction of the recurrent state h through two pathways: a direct pathway, i.c., using & to evaluate f(h,-1: Xx, $6, pathway, ie., 1b affected constructing all previous recurrent states, h h1, and all those recurrent states affected h, construction. Thus, to calculate Sh. we need to consider those two pathways:
and an implicit
Oh. Oup
(f +) Örb
of (h, Oh.-
(8)
Oh
Onp
Once again, need to consider the two pathways when evaluating
in 8. For simplicity, let 8h 8h d
J. = Sh, B,
-1. C.
-.(,,-1xx..) and re-write 8:
dh.
J. B, + C,J,-1
B, (B C.-Jr-2)
Unrolling J,-1
B, +
B, C,B,, C,C,,, (B, C2-2J1-3)
Unrolling J,-2
B, + C,B,-I CCC, B, + C.C, C,-
B, C,B,, C,C, B, t... C,CC,-CC--- C,B, C,C,-, CjJo
Keep unrolling
Σ
C
B
[ c
Jo.
M=R+1
(9)
13

---

Writing The using the results from 9:
Oh,
1,, 16) Th:-1
+
eup
@h.-1
Erb
(10)
Σ k=1
II
or 11
Or(h: 1, XR, 2b)
II
df of Xi; 2b)
Tho Ōrb
Th;-1
Erb
bh.-1
According Eq. 10, the agent needs store all the previous inputs calculate
Uhs
which is
impractical: the computation and memory complexity will be increasing with
A.1.1 Truncated-Ban BeraateddBackPPeee Paaneatte-BackkPrppaattoo Through Time
Williams and Peng (1990) introduced Truncated-BackPPPe solves the issue of increasing memory and computational complexities of BPTT. In T-BPTT, specify truncation length 7', 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:
Through Time (T-BPTT) which
J. = Σ II
C.
B
(11)
k=t-T
\1=++
Using results from 11, we then write the approximated gradient of the loss w.r.t the learnable parameters:
0C, Oh
v Σ
-= Th Tup
Σ 0L,
Σ
II ki=
Of 0f 1b)
x;;
Of 0f Xb,
Of (,,
Oho
Th k=j-T
Oh.11
Orb
Oh:-1
Oub
(12)
A.2 Real-Time Recurrent Learning
Williams and Zipser (1989) introduced the Real-time Recurrent Learning algorithm (RTRL) learning algorithm for continual recurrent learning. RTRL employs the recurrent formulation of the gradient in 8; instead of unrolling 8h,-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 converge. The gradient formulation of RTRL can be written as:
UL Oh,
Dh Onb
OC,
(13)
Oh,
1=0
0f 0f: X4, v))
af ,,, @h;-1 Oup
Oup
Oh.-1
B More Details on Representability with Complex-valued Diagonal Recurrence
This section explains why need complex-valuee diagonals to represent dense recurrent layers. We first show when is equivalent to use complex-valued diagonal and a dense recurrent layer. We highlight that using real-valued diagonal is like restricting the weights to be ceametrrc--eccusee
14

---

the (complex) diagonal corresponds to the eigenvalues of the weight maarix--hich can severely limit representability. We provide 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-valus -valued diagonals.
B.1 Representability with Complex-valued Diagonals
Let us first consider when we can perfectly represent a dense, linear recurrent layer with complex- valued diagonal recurrent layer. Assume we have the recurrence relationship, with learnable parame- ters W E IRaxn and WT E Raxdd
h f(Wab:-1+ + Wx)(x)))
(14)
where f is a potentially nonlinear function that inputs a vector and outputs the same-sized vector and can be any transformation of the inputs x before they are inputted into the recurrent layer. The following equivalence result is straightforward but worthwhile formalizing.
Proposition B.1. Assume foP Pof f for any full rank, potentially complex-valued valued P E C with unit-length column vectors. Then given any W and W., for Equation (14), there is corresponding complex-valued diagonal weight matrix A E C and W E Caxd
h = f(AR W u(x()).
(15)
where h E C is a linear transformation of h, E R".
Proof. We can rewrite the square matrix Wt, using an eigenvalue decomposition Wa = PAP-1, where P contains the linearly independent eigenvectors and A is diagonal matrix with the corresponding eigenvalues. Then, we can re-write (14) as:
h = f(P AP h + PP PP- Wz u(x,))
= Pf(A P-1 h + P P_ u(x:))
(16)
P- f(AP-) h11 + p-1 W. u(x:))
where P came outside of f under our assumption that it commutes with such matrices of eigenvectors. By defining h, = P h, and W. = P P-1 we get Eq. (15).
We can see h and h, are representationally equivalent: they are linearly weighted for downstream predictions, and so the linear transformation on h, can fold into this downstream linear weighting. But it is more computationally efficient to use h, with a diagonal weight matrix A, meaning each hidden unit only has one recurrent relation instead of
Since we did not impose constraints on the matrix Wh. other than being diagonalizable, the eigen- values of W, A can be complex or real numbers. Previous diagonal RNNs such as eL [17], Columnar networks [18], and IndRNN [20] use only real-valued diagonal matrices. Having only real-valued diagonals implicitly assumes that the matrix W, is symmetric matrix. [34] suggested using complex-vaii -valued diagonal matrices for better performance.
The above equivalence has only been used without any activation, namely in linear-recurren units (LRUs) [34]. A natural question 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 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-oonly-i relationship (see Appendix C). However, even for this restricted setting, this generalized class is limited and such activations unlikely to be preferable to linear recurrence. We see this negative result, that suggests this equivalence only holds for the linear setting.
B.2 Complex Figenvalues in Vanilla RNNs
We empirically investigate whether complex eigenvalues appear when training dense RNNs in simple task. The goal is to show that the weight matrix, Wn, is not symmetric matrix, even in the simplest tasks. Hence, having only real-valued diagonals is too restrictive.
15

---

We used the Three State POMDP [42], depicted in Figure 7, for this experiment. In this task, the agent needs to remember one eue from the previous time-step ago, make a prediction about the next time-step. The MDP has three states, $1, $2, and 83, and no actions. If the agent is in either or 82, it transitions any of the three states with equal probability. However, the agent in 83, transitions to the state preceded by $3. A sequence of observations would look like 1,3, 1, 2, 2, 3, "* and we ask the agent to predict the next observation.
$
$2
$3
Figure 7: Illustration of the Three State MDP. We used dashed lines the transitions starting 83 to make them more visible.
We trained vanilla RNN with 3 hidden states with T-BPTT with truncation length 2, which is sufficient history length in this problem to predict the next observation. Since we have 3 hidden states, the matrix WT R3x3 and could have at most complex eigenvalues.
We measured the performance in terms of the percentage of correct predictions made in S3. We recorded the number of complex eigenvalues of W, after each parameter update, shown in Figure 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 cigenvalues across 30 run is above 1.5, which means that on more than 1 of the steps, the RNN has two complex eigenvalues. The primary point is that we see complex eigenvalues appear frequently.
Prediction Accuracy
1.0
# Complex 2.0
Eigenvalues
1.5
1.0
0.4
0.5
0.2
0.0
0.0
200
400
200
400
Steps (K)
Steps (x 1K)
Figure 8: Left: The percentage of correct predictions when training an RNN the Three State MDP. Right: Number of complex eigenvalues when training RNN in the Three States MDP. The solid line is the mean over 30 runs, the shaded region area is the standard error, and the lines are individual runs.
C More on the Equivalence of Non-Linear RTUs and Dense RNNs
As discussed in the main body, we likely only have equivalence between using a full dense weight matrix and a complex-valued diagonal matrix for linear recurrent layers. However, we can obtain slightly more general equivalence in the restricted setting where the weight matrix for the recurrence symmetric. This restricted setting not general interest, but we include the result here because it could be of interest to few.
16

---

For the case where we have symmetric weight matrix, we need the activation to commute with orthonormal matrices. Consider again the form
h f(Wah--1+ WxX)
(14)
where f is potentially nonlinear function that inputs a vector and outputs the same-sized vector. To obtain the equivalence, assume that for any orthonormal matrix A, fo A A o f. We can rewrite the square and symmetric matrix W, using an eigenvalue decomposition WA = A A A where A contains the n linearly independent eigenvectors and an orthonormal matrix and A is a diagonal matrix with the corresponding eigenvalues. Then, we can re-write (14) as:
h& f(A A T h AA TAW,, x1)
Af((A A h + A J.W, x²) A h, AAT h + A W
(17)
where A came outside of f because commutes with orthonormal matrices. By defining h; = A h and W. = AWW,, 2. get:
h. = J(Ah + W₂ x.).
(15)
Each hidden unit now has recuirent relation instead of because weight matrix A is diagonal.
natural question if only the identity f(x) = x --nnally 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 negative result, highlighting that this equivalence largely only holds for the linear setting and does not generalize to other activations of interest. provides even further evidence that likely the only setting of interest for the general non-symmetric case is also with linear recurrence.
Nonetheless, let obtain the in-l-n-only-ii for completeness. simple extension that continues to satisfies this property is f(x) = xc(||x|2) for any c IR + R. We can see that for any orthonormal A, we have
f(Ax) = Axe Ax| Axe( x x1)) = Af(x).
This means that we can have activations that rescale the input depending on the norm of that input.
More generally, the activation involve matrices and rotations. We can define ((x) [x Ux))] where U(x) E Rexx-1 is a matrix where the columns are orthogonal vectors to each other and to x. Then for any vector- g R IR"", we have that f(x) = 0xx) (([xx|12) also satisfies this property:
f(Ax) = ((A)) Ax|AA||2 = [Ax (AA)]]((())))
A|x: U(x)]9((2x(1)) Ax(x).
The last line follows because AU(x) = U(Ax). Namely, for any orthogonal vector u with u x O.
we have that û Au satisfies û Ax u T A Ax = u x = because A A = I.
Now we show this formally. Denote O Rax= A E O" > A A AAT I. Denote
where the 8th element is 1.
6,-
Definition C.1. (xx) is matrix such that:
(xx) E n
(18)
0 0(x)x 1x112201
Lemma C.2. xxx) exits and VA O", ((Ax) 0(x)AT

---

Proof. Let (x) be E O". Then, by definition of O", 0(x) 0(x)0((x)* = I. Let [(x][1]
x, and [2 :] be any orthogonal vectors, also orthogonal to x. Where ((x) be the 2th row of (x)). Then, for x E I2 00())x |x 2211
A(Ax)Ax ||Ax||2 e = |xx|2 - 0(x)x
0 0(Ax) ((AA))((xx)) 0(x)x
Ax 0(Ax)* 0(Ax))0(x)x 0(xxxx
A A²O(AX) 0(x)x
(19)
= A ( xx))x
A 0(() = 0(Ax) = =(x)
Theorem C.3. Let f: :::18 IR" + R" and A E O". Then, f o A Aof = 3g: f(2)) 0(x)" (([=]))
Proof. Define g(a:) = f(a2e1).
f(x) O(x)x)
0(x) f(0(x)xx)
(20)
=0(x)" 0(x)"{([|x[2> (2)1))
(f A) f(Ax)
= -(Ax)) )' g((1Ax|12)
=((A)) "g(1x(12)
(21)
A0(xx) ((([x)12)
=A o (x)
D Issues with Two Alternative Parameterizations
In this section we provide a few additional insights on alternative ways to handle complex numbers within an RNN, and why they not preferable.
D.1 Stability
We look at the gradient when using each complex representation to understand how different repre- sentations 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 consider
hr Ah4-1 = Xhho
(22)
where top is the initial hidden state.
Real Representation a + bi: Substituting À in (22) with the real representation, we get:
h4 (a bi hp Σ
k= 1
18

---

Then it follows that the gradient w.r. the learnable parameters a and b is:
The
ho
k) (t k)aaa>-- 1
da
k=
The
ho
ka 2- b
0b
k=
To prevent the gradient from vanishing/explodingg we need to restrict both |a| and [b] to be E (0, 1]. Exponential Representation cx)( Substiuuting À in (22) with the exponential representation, we get:
h e'
and the gradient w.r.t the learnable parameters and 0 is:
The Ir -arrrrlexppl exp((l)))
The
r'exp(i10)ilhh
Ör
ae
To prevent the gradient from vanishing/exploding, need restrict (0, 1]. Cosine Representation (0) + iiin())): Substituting À in (22) with the cosine representation, we get:
h; r' (cos((0) + iiin
and the gradient w.r. the learnable parameters and is:
dh; tr' (cos(tt) + isin(t)))ho
Òr
0h: r' (il cos((1) - ssn 10)
00
To prevent the gradient from vanishing/exploding. we need to restrict r E (0. 1]. It is simpler to maintain stability with the exponential and cosine representations, since only need to constrain r E (0, 1]. whereas the real representation requires us to restrict both the complex number's magnitude and phase.
D.2 Biased gradient when only using the real part of the hidden state
We can attempt to get the benefits of having a real-valued hidden state by simply converting the complex-valued state real-valued one within the LRU. However, taking only the real part results 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: Yt = wRe(hh)) where w is a learnable parameter. The gradient w.r.t and 0 is:
Oy. cos(())4-11 cos(E) Thg-1) Or
(23)
Oy: w -rssin + (O) The-1 00
00
Re{z} [mg[z}
-Img(z}
Multiplying by a complex number is equivalent to rotation by the matrix
Re(2)
and scale by / Re[z]² + [m] We re- writt the recurrent unit using this property h = r cos (0)h - - (in
h = r cos( (2 + + (i
(24)
yt = wh + h(a))
19

---

Notice 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:
ay dr
w(r 1 +
c0)
Ör
0
+ rsin(0)(h - 11 + rinn
Ör
(25)
Dy: = w(-rssn(0)(()))) h, cos
ah
dbiー1)
0)(
+
00
00
00
a
0
+ r (0) - sin
00
00
Comparing Eq. 25 and Eq. 23, we can see that using only the real part of the recurrent state leads to loss of information in the gradient.
E
Recurrent Trace Units
This appendix details the parametrization used for RTUs, the derivation of the RTRL update rules, and the extension of RTUs to multi-layers.
E.1 Empirical Analysis for different and & parameterizations:
Since r represents the magnitude of a complex number, then r E Rt, is preferred to have r E (0, avoid vanishing/exploding gradients discussed in the previous section. Let wi be a learnable parameter which could directly represent or represent function of we can enforce the constraints on in several ways:
1. Direct learning: Learn directly, and clip it after each parameter's update to be in (0, 1].
2. Enforcing r E +R+ Learn such that r = exp(-v). This parameterization enforces r to be positive. However, additional clipping is needed to enforce (0, 1].
3. Enforcing stability
be (0, by using a positive non-linear function. For example, learn ylor such that r exp( - cxp() this parameterization ensures that r (0, 1] and suggested by Orvieto et al. [34]. Another example is learning (V), which also ensures (0, 1].
We enforce
Finally, we can also enforce stability on 0 by learning olog such that = exp(Alo5) to ensure that always positive.
MSRE
Linear RTU with hidden units
exp(=exp(p 3), 0 exp(
0.35
yo exp(-exp(u )2, exp(-v)_
0.30
0.25
0.20
0.15
0.10
Step Size
Figure Learning rate sensitivity for different parameterizations of r and for RTUs with hidden units.
We empirically compare the different parameterizations of and our experiments, learning directly resulted in unstable training where the MSRE diverges. We plot the learning rate sensitivity
20

---

Uiber XTU with 512 hidcen units
MSRE
exp(-exp(u*F)) exp(p
1.5
z= exp(=exp(o"9)) r- exp(-v),
1.0
0.5
0.0
Step Size
Figure 10: Learning rate sensitivity for different parameterizations of and 0 for RTUs with 512 hidden units.
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 glog and plug has better learning rate sensitivity.
F.2 Real-Time Recurrent Learning for Single Layer Linear RTUs
We now outline the details for the RTRL update rules for RTUs. The set of learnable parameters for RTUs is ab = {ylog, Glos, W W }. At each time step f, the learner receives a loss (i ye; vb:b)) where V/ the network output at time 1, then the gradient of the loss w.r.t the parameters is:
Oc.
0L, Th, Oh
aL, Oh, Oh
+
(26)
Th Th
Oep
Th
Th
Oxp
where
8h 0
8h Tig
Sh. 3H
Sh TW
8h Sw
and Oh 0V
0h
Oh
Sh
Sh
50
DW
OW.
We derive the following gradients:
Oh
Ög(v.00 O)
Th f-1
OΦ(v.07 0)
0h
ey
h + g)
② h1
$(v.0)
+
WWix,
Oplog
Oulog
Oylog
Oplog
Oylog
Iolog
Oh
Ög(v.0)
2
ÖØ(v.0 0)
en
h + ((
h (
Oygog
Oylog
Oylog log
Oplog
dolog
dvlog
(27)
Oh
0g(v, 0)
Th
0Ø(v.0
@h
h + g(v,0 0
h Φ(v.0)
10
00 log
log
O0log
80log
(28)
Oh
0g(v, 0
Ih C2
0Φ(v,0) 0)
@h
h + g(v, 0
h $(v.0)
00
30 log
00 og
O0
80
where
0g(v, 0)
g(v, (x)(
Oy
0g(v, 0) $lv, 0) exp()
O0log
(29)
00(v, 0)
Φ(v. 0) (ok(
Oylog
00(v, 0) g(v. 0) cxls)
O0log
21

---

To efficiently compute the gradient w.r.t W and W we look at the influence of each when considering a single element from each recurrent state, h and h
M - o(s,, 0:),,, ous 0,)00 + 74
w 2. "tj
(30)
13 - p(u,, 0:)6/1, + oes 0,)11 + 74
w (t.5)) F1.j
5
We then get:
Ohi
0h
g(v,, 0,)
0(,,,0))
+ 71015
dwa (.
ow
ow
(,,1))
2,(1,))
ah
0h1-1
0h
g(us 0.)
601 0.)
aw (3.5
owel
oWel
7,(3.)
2,(1.)
(31)
Oh
0hi-1
g(v,, 01)
Sw 2.655)
o (1
OW 2,(1.5)
0hi-1a
g(vs 0:)
0(V,, 0:)
+ 76
aW x.((.)
W 2,((.)
(t,f))
0h
We see that each h gets affected by weights from only one row of W thus,
w can be written
matrix of the same dimension Wa. The same is true for 8h 0h?2
22 and 8h.2
W W
W
E.3 Single Layer Non-Linear RTUs Formulation:
We extend the linear RTUs to non-linear RTUs by adding non-linear activation function f to the recurrent states. We can write the non-linear RTUs as follows:
h a f(g(v, 0)
Φ 0
h
+
h fgg(v,0)
h + (
W xx)
(32)
h [h,;; h
where f _ R Following the same procedure as the linear case, we can derive RTRL update rules for the non-linear RTUs.
ah
Ög(u, 0)
aepi 0)
2h
Oy ()))
c h + g(v
@ h $(v.0) 0)
Oylog
Oying
Oyplox
(33)
Oh
0g(v,00 0) do
Th
0p(v, 0)
Th
(())
h 1-1 g(v,0)
+
h + $lv 0)
Oylog
Oylog
Oy
OyloR
Th
0g(v.0 0)
0h.
05(v, 0)
Oh
r' r'(-)
h g(v.0)
④ h,
Φ(v,6 0
00log
aolog
a0log
a0loz
o0log
(34)
0hF 62
(g(v.
Oh
00(v.0
Oh
(
2 h + g(v
h
+ $lv 0
aging
O0log
00
30 ing
h f(g(,,
0(,,,) -(v + 78 Σ
w 9
Tt.
j=
(35)
f 0.),, 0:)
((() It.j)
1=
22

---

We then get:
dh
(g((, 01)
664, 01)
aw 2.(6
ow (.()
aw 21(1,))
Oh
0h5-15
OW ('))) 0:)
+ 0:)
ow
Ow
(,5))
x.(i.)
(36)
('())) 0:) 0h1-1.
o(vi, 04) 0h1-1.5
2 2(()))
aw (b,)))
aWa 2,(,,)
Oh
0h
0h
='(())(( 0:)
+ +lu, 0:)
+ +za.))
aw 2,(6))
ow? (1.5)
oW> (1
E,4 Complexity Analysis of RTUs
We 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₁ E Fd and hidden states h, = [fbb");; f(b,"")) E 2 .w have RR and Wel 20 W E Rddxn
2h axp
0h Cnp
An agent using the RTU with RTRL needs to store the gradient information, one step to the next. We denote the set of saved gradient information as:
and
from
ah
Ov
Oviog
Th
ah
V
0
00log
00
(37)
Oh
Oh
Th
W
aw
oW
aw
OW
The saved gradient information has the following dimensions:
V E 2
0 E R
(38)
V w E ((dx)))
Then, 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, forward pass according to 2 has 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.
E.5 Multi-Layers Recurrent Trace Units
We now extend RTUs to multilayer setting. We show that in the multilayer case, we lose the compu- tational advantages. However, prior work suggested that treating each recurrent layer independently 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 layer where S d,. The network has the following set parameters:
- {vb1, 22., ths 3n vb lologi glogi Wenk W
(39)
23

---

To update the network parameters, we need to calculate the gradient of the loss w.r.t the parameters from all the layers:
ac,
dh OC,
0rb Oh
(40)
Oh
Oh!
Ohe ah
abg Orp
Örp
erb
eup 0rb
where
h = [f
(41)
f(h, m)
Take as an example the gradient of h layers of the network, we get:
cl,x
w.r.t plog.n log
log.1. Unrolling the last two
h g(v", 00) о h
Ø(v"
W m-1
[h cl A(h 22
g(v", + h
"1v" 0")
22.R
W
(42)
g(v" 0²) 2 h Fn.
Φ(v" 0t)
② h
W
f(g(v=-1 0 ② ha
6(v
n--
y
Wa h 1hm-3 3)
f(g(vm-1,
ha.a- h
6(v=~1,
g
y n- W h
The gradient of h,l.n w.r. can be calculated in linear complexity indicated in the previous section.
Calculating the gradient w.r.t the parameters from the earlier layer:
Oh
Th
Sh c2.R
Oh n-1
g(v":
W Övlog,n-1
Tylog.n-T
Ovlog.n-
Tvlog.n-
Oh:-1
Ovlog.n-1
E
R22d
2d Can be calulated with linear complexity.
(43)
0h
E R di xd, Saved from previous timestep.
Jylog.n-I
@h C2,7
1
ddnxdd-1 Saved from previous timestep.
Th C2.7
Oh
Öv iylog.n-2
g(v" 0n)
6(v" 0")
dvllggn-2
+ y
Wa.n
dylog.n-2
Th
g(v"
Φ(v )0
+ y" 4 Worn
Ey 1-
Ey 2
g g n-1 0 0n-1)
Φ n-1 0n11)
n- 3 W
Oh; dplos.n
((()
(gv 0a-1)
p2 0n
+ We
Oh." 2
Bulok.n
Bulos.n
Oh,- n-2
E R 2dn Can be calulated with linear complexity.
0h -1
E R dn-1 xda-2 Saved from previous timestep.
n22
Oh c1.n-1
E R dn xda-2 Saved from previous timestep.
Oglog.n-g
RR x Saved from previous timestep.
nylog,n-2
ah.za
Rd xd._2 Saved from previous timestep.
Oylognn-2
(44)
24

---

Let's define J Bylngg 1 Then, to calculate the gradient of the hidden units from layer n w.r.t the parameters from layer n2 1, we need save Jn I, and to calculate the gradient the hidden units from layer n. w.r. the parameters of layer n 2, we need to save Jn,n-2 and Jn-1, 1-- 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 In 1/- 2n,2,,,..... J2,1.
E.6 Implementing RTRL within the reverse-mode automatic differentiation
For function : R - R we have the Jacobian 0f(z) R and calculate this Jacobian in two ways: forward-mode or reverse-mode differentiation. In the forward-mode differ- entiation, 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 forward passes to get the full Jacobian, where each forward pass would allow to compute the derivative w.r.t only one of the inputs. i.c., with forward mode differentiation, we evaluate the Jacobian one column at a time. As a result, forward-mod differentiation is inefficient for neural networks; neural networks map from learnable parameters, which can be in millions, to 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 a time, which is partieularly advantageous for neural networks. This efficiency comes at the cost of two passes through the network: a forward pass for function evaluation and backward pass for derivative evaluation [2].
RTRL is an instance of forward-mode diffcrentiation: during the forward pass, the gradient infor- mation is evaluated along with the recurrent function computation. As result, there is 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 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]."
F Additional Details on Trace Conditioning Experiments
In animal learning, Trace conditioning a type of experiment where animals predict the occurrence of a stimulus (c.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 together- playing the tone and then serving the the animal learns to anticipate food arrival when it hears the tone [36]. We use an open-source rrace conditioning benchmark introduced in prior work [39]. Two signals appear sequentially: the Conditional Stimulus (CS) and the Unconditional Stimulus (US). The CS 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, Ge:
8
G
US
(45)
where y 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.
'This can be implemented by defining a custom vjp for the recurrent layer, which modifies the backward pass for the recurrent layer include the gradient traces.
25

---

0.5
0,4
0.3
Diagonal
MSRE
Online LRU
0.2
0.1
Linear RTUs
0.0
4K
8K
64K
Number of Parameters
Figure 11: Comparison to a block diagonal RNN.
F1 Additional Trace Conditioning Experiments
Comparison Other RTRL-based Architectures: We now compare RTUs other RTRL-based approaches with similar architectures: online version of LRU [48] and vanilla block diagonal RNN. The block diagonal a recurrent formulation similar RTU but ignores relation between by
the learnable parameters. i.e., replaces C₂ in 3.2 with C2
Jos
dx
The results figure indicate that these scemingly small differences between the diagonal RNNs result in significantly different behavior. RTUs outperform online LRUs, with the differences discussed in-depth Section 3.4. RTUs also outperform the block diagonal RNN. We emphasized using real-valued diagonals implicitly assumes symmetric matrices, but that is for single real-value. This block diagonal has more representational capacity than the RTU. This result suggests it is beneficial for learning enforce these constraints on the learnable parameters, that they correspond to the rotational representation of complex numbers.
To highlight the role of RTRL in RTUs, we evaluated the performance of both linear and non-lineas RTUs with T-BPTT. this experiment, we all agents use the same number of parameters; the only difference is whether they use RTRL BPTT.
On the role of RTRL in RTUs:
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.
Fvaluating
0.25
0.20
0.15
RTUS
T8PTT
0.10
0.05
Linear RTUs w. RTRL
30
60
120
Truncation Length
Figure 12: Evaluating RTUs with BPTT.
26

---

F.2 Trace Conditioning Experiments Details
All agents have one recurrent layer, either an RTU or a GRU, and one linear layer. At each time step the agent passes the observation to the recurrent layer, which outputs the recurrent state, the agent state. The recurrent state is then passed the linear layer generating the prediction. For each agent, we swept over the learning rate used to update the network parameters, E 1,0-1,10-2 10 10-3 10-1, 10-5, 10-%), and averaged the performance for each learning rate over 5 independent runs. We then selected the best- perforrennn learning rate for each agent and ran 30 independent runs using it. For all the experiments, we ran the agents for million steps, and the performance was the mean squared prediction enror averaged over the 2 million steps.
F.3
Learning Rate Sensitivity
We show the learning rate sensitivity for all agents in the animal learning benchmark.
CHU Leorn Compuranenel
CHU SOB Compusation
CIU SOOF Computar on
017u2
(NT)
1.0
3473 2005
01712004
Learning hate
Leininng Pete
Lecen Rete
Loanning Park
Figure 13: GRUs used in the animal learning benchmark. The (H: 7) in the label refers to the (hidden dimension: truncation length) for the GRU.
roas
Cotiolrcttia
Optio1
Octon
UIU Scaling w. Naamssts DIDH
G.44
coe-
0:35
8.30-
9.1%v
48 47 d5
29
07
-1 :00
27
<1 40
27
15
Figure 14: LRUS used the animal learning benchmark, The (H: 7) in the label refers to the (hidden dimension: truncation length) for the GRU.
HERE (5 runs)
ITUs Learning Computational Constraintts
Linear Scaling Computation Option
NonLinear RTUs Computation
Option
0.15
neer ETUIS0
0.15
0.15
KTUs:110 57U5:2200
0.10
NorLinear STU-500
0.10
KTU0
0.10
-8.
-17
-66 -15
44
-13
42
41 -10
4
-18
-17
-16
-15
-34
-3
42
+
0,00
-18 -17 -66
+15 -14 -13
-10
log.0 eamiing Rate)
l0go 2lemming Rane)
10g:0 wamiing Rate]
Figure 15: RTUs used in the animal learning benchmark. The number in the label refers to the number of hidden units in the RTU.
G 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 to go through the whole trajectory after each epoch update and re-compute the gradient traces. However, this can
27

---

Linear RTUs in in Scaling w. Computation NonLinear RTUs in in Scaling w. Computation Option
MSRE (5 runs)
Option
0.20
0.15
0.15
0.10
0.10
160
0.05
800
h:450
h: 230
0.05
800
450
230
3200
h: 61000
3200
1650
0.00
0.00
-18
-16
-15
-13
-12
-10
-18
-16
-15 -14
-13
-12
-11
-10
-9
log2 (Learning Rate)
(Learnngg Rate)
Figure 16: RTUs used in the animal learning benchmark.
MSRE runs)
Linear RTUs in Scaling w. Parameters
NonLinear RTUs in Scaling w. Parameters
0.20
0.20
h:80
0.15
h:160
0.15
h: 80
h1600
h: 300
0.10
0.10
160
h:600
h:300
0.05
h:1150
0.05
600
h:2300
h: 4600
1150
h: 2300
0.00
0.00
-18
-17
-16
-15
-14
-13
-12
-11
-10
-9
-15
-14
-13
-12
-11
-10
-9
-8
logz (Learning Rate)
log2 (Learning Rate)
Figure 17: RTUs used in the animal learning benchmark.
be computationally expensive. In Algorithm I, 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 ad advantage estimates.
Undioounec 5922 Fes
Liverr AFV, epeensi
Stale proderssiiin targe
6000
Linerr RILL spocts
Make
4000
Wie Falie
2033
Fobec
Falle
Staps
Stops
Figure 18: The impact of stale gradients and stale targets when combining RTRL and PPO on Ant.
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 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 minimal effect 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 KI. divergence between the policy used to collect the trajectory and the most recent policy. We use the following estimate for
28

---

the KL divergence between the two policies as (r- log(r), where (Te ) (Teaaa) 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.
Approx.
Approx.
Non Linear RTUs, apochs:
Non Linear RTUs, apochs:
0.08
Gradient: Target:
Gradient:
True Targel: False
0.03
Gradiant: Targat: The
Gradient: Raisz Target: Faise
0.06
0.04
0.01
0.02-
Gradient: Thue Target False
Gradient: Target: Thue
larget:
Gradients Faise Target: Faise
0
4
0
Env Steps
1e6
Env Steps
1e6
Figure 19: (a) Approximate KL divergence for NonLinear RTU with epochs.(b) Approximate KL divergence for NonLinear RTU with 8 epochs.
Approx, IS runs)
Approx. KL IS runs)
KL
LRU, spochs: 4
LRU, epochs:
Gradient: True Target: True
Stale gradient: TrueState target: True
Gradient: Tree Target: False
Stale gradient: TrueStale target: False
Gradient: False Target: True
0.09
Stale gradient: FalseStale target: Thue
False Target: Faise
gradient: FaiseStale target: False
0.030
0.06
0.015
0.00
Erw Steps
1e6
Env Steps
1e6
Figure 20: (a) Approximate KL divergence for LRU with 4 epochs..b) Approximate KL divergence for LRU with epochs.
H More Details on the Memory-Based Control Experiments
In our implementations, we use a shared representation learning network followed by two MI.ss for the actor and the critic heads, illustrated in Figure 21. The shared representations consist of a feedforward layer with 64 hidden units and memory components: an RTU, LRU, or GRU, The actor and the critic's heads consist of two feedforward layers with tanh activation function.
Additional Mujoco Results:
In Figures 22 and 23, we set the truncation length for GRU and LRU to be 64, which is larger than needed 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 transformer- based models, GPT2, perform worse than RNN-hased models. This consistent with previous work suggesting that transformers might not be suitable for RL tasks [32].
29

---

Algorithm Pseudocode for integrating RTRL methods with PPO
Inputs: differentiable policy parametrization saas, Wp).
Inputs: differentiable state-value function parametrization ö(s, W.).
loop
Generate a trajectory using the current policy Oo. Ao. Ri..... OM AM. RM-
Store the transition and the gradient traces for the recurrent components along the way for = O..... .M
Compute the advantage estimates and the target value for each timestep for epoch = ...... k do
Divide the trajectory into minibatches and shuffle them. for minibatch = l.... m do
Calculate PPO loss
Perform gradient step with AutoDiff and correct it with the RTRL saved gradient as discussed E.6.
end for [Optional] Re-run network to update hidden states and the gradient traces for the trajectory. [Optional] Update value targets and advantages estimates for the trajectory. end for end loop
Actor Head
MLP
<(a/ss))
Memory
MLP
Model
MLP
i(ss)
Critic Head
Figure 21: Agents architectures used in our control experiments.
Chastahv
4000
1200
1340
1000
-3000
Figure 22: Additional results on Mujoco-P, where we allow GRU and LRU to use larger truncation length than needed to solve these tasks. We also show results for GPT2.
Learning Rate Sensitivity: Figures 24, 25, 26, and 27 show the learning rate sensitivity for all agents in the Mujoco POMDP benchmark. Finally, we used the default hyper-parameters for PPO [40] indicated in Table for all agents.
30

---

200-
4000
7033
590-
Figure 23: Additional results on Mujoco-P, where we allow GRU and LRU to use a larger truncation length than needed to solve these tasks.
1800
Tord
UMA Teml
2000
700
2009
LAUUTH4
1530
0,011
200001
201
2.0091
0.01
Figure 24: Learning rate sweep for LRU in the control experiments.
Are
Choeaee
Wallore
Hppeot
4000
1.00
3110
1000
Lobur K0U
700
130
camen
2.001
00033
2.03
00901
0/0077
0.001
Figure 25: Learning rate sweep for RTUs in the control experiments.
Aree
Cheesee
Walker
Ppppa
4000-
3030
2100
2000
1000
1530
500
700
cacm
0.000
02003
20001
A3000
200033
0.00
Figure 26: Learning rate sweep for GRUs in the control experiments.
I Compute resources
We ran the Mujoco-P, Mujoco-V V on NVIDIA P100 GPU. Each of the Mujoco-P and Mujoco-V V trials took around 30 minutes complete on single GPU. For the POPGym experiments and animal learning experiments, we used large CPU cluster. Each trial of the POPGym experiments took around hours complete. While each run of animal learning took around 15 minutes to complete on single CPU with memory less than 1 GB.
31

---

320
Oretan?
Hepport
244
130
306
405
256
-19
6n0 Tridt 20:20
-300
199
20000
20000
10000
0,00
6.7001
0.0000
Figure 27: Learning rate sweep for GPT the control experiments.
Name
Value
Buffer size
2048
Num epochs
10
Number of Mini-batches
32
GAE,Ä
0.95
Discount factor, T
0.99
policy clip parameter
0.2
Value loss clip parameter
0.5
Gradient clip parameter
0.5
Optimizer
Adam
Optimizer step size
[le - 05, 3e - 05, 1e 04, 04, le 03
Table 1: Hyper Parameters for PPO.
NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification:
We provide detailed theoretical and experimental analysis that supports
claims.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Limitations are discussed in the conclusion section.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
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-specinf 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, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tes system might not be used reliably to provide closed captions for online lectures because it fails 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, worse outcome might be that reviewers discover limitations that aren't It acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full 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 include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide 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.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that 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.
Guidelines:
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 provided or not.
33

---

If the contribution is 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 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 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 large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome 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.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
A link for public repo is provided.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/ wariic/gguie///eee htt.g/de/o/eooeeemmiiiia
Polnic)lyy for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to reproduce the results. See the NeurIPS code and data submission guidelines (https: // //nips.co cctessccppppp wttpcccccppbbccccccccccccc http:/pulllllllleell wttpc//g/g///eeeeeeeeoocce http://dddee//eeeuuumieeeeeeeeee http///eeummmiiiiooooooooooo
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
34

---

Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: all hyper-paramedd 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.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Ys]]
Justification: we report enror bars/shaded standard enror 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, confi- dence 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, 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 the standard error of the mean.
It is OK to report 1-sigma eтror bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Compute details were provided in the appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
35

---

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).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics httpo::/eerris..cc/lll////////
Answer: [Yes]
Justification:
The research conducted follows the NeurIPS code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g.. if there is special consid- eration due to laws or regulations their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper presents work whose goal is to advance the field of Machine 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 no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended (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 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 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The research conducted here doesn't include large models or scraped datasets.
Guidelines:
The answer NA means that the paper poses such risks.
36

---

Released models that have high risk for misuse or dual use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: all related work has been properly cited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. dattattsess has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not introduce new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The research conducted here doesn't it include human subjects.
37

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The research conducted here doesn't 't include human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and 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 the institution conducting the review.
38


Reviews:
[{"review_id": "SpwUrmTQx1", "reviewer": "Reviewer_nsAS", "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. \n\nIn 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.", "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": "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. \n\nThe 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.\n\nIn 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.\n\nLastly, 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.\nGrammar:\n•\tLine 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?", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "•\tThe paper addresses an important problem in online reinforcement learning and proposes a computationally more efficient solution, RTUs, which outperform other architectures in various online reinforcement settings\n•\tHas comprehensive tests covering and analyzing a variety of cases, which empirically show the improvements from the proposed architecture, RTU.\n•\tProvides detailed explanations and derivations of the RTUs formulas, making it easier for readers to understand and replicate it\n•\tResults are explained well using figures\n•\tSupplementary section goes into detail on many areas which are not fully explained in the main text and answers many questions that may arise.", "weaknesses": "•\tMost 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\n•\tExplaining the importance and benefits of addressing problems with online reinforcement learning in partially observable environments would strengthen the motivation.\n•\tSome terminology not completely clear, such as difference between “Reinforcement learning” and “Online reinforcement learning”", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed and valuable feedback. Here, we respond to the points the reviewer mentioned. \n\n> Explaining the importance and benefits of addressing problems with online reinforcement learning in partially observable environments would strengthen the motivation.\n\nWe motivated the importance of learning under partial observability in the first paragraph of the introduction when talking about how in the real world, agents perceive their environments through imperfect (i.e., partial) sensory observations. Hence, for agents to be deployed and learn online in the real world, they need the ability to predict and control under partial observability. We agree that the motivation for learning under partial observability is core to our work, and we will expand our current motivation in the introduction to emphasize the importance of learning under partial observability. \n\n>Some terminology not completely clear, such as difference between “Reinforcement learning” and “Online reinforcement learning”\n\nThanks for pointing that out. When we say \"online reinforcement learning\" we mean that the agent learns while interacting with the environment. This contrasts with offline reinforcement learning or reinforcement learning with access to a simulator. We will add a couple of sentences to the intro to make this clear. \n\n>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?\n\nWe chose GRU as our main baseline since it’s widely used for reinforcement learning problems with partial observability and has been shown to outperform most of the other memory components in several domains, which makes it a strong baseline to compare against  [1][2][3][4]. Additionally, the authors of the original POPGym paper that we use in section 6 have extensively tested various recurrent and transformer architectures on the benchmark, and GRU was always the best-performing architecture [1]. Since we are reusing their benchmark and similar other tasks, it is natural to compare against GRU as they have already done extensive testing with the most known architectures. \nWhile LRU has not been explored in reinforcement learning domains, we added it to our experiments because it is similar to our proposed architecture. \n\n\n[1] Steven Morad, Ryan Kortvelesy, Matteo Bettini, Stephan Liwicki, and Amanda Prorok. POPGym:Benchmarking partially observable reinforcement learning. In International Conference on Learning Representations, 2023\n\n[2] Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhutdinov. Recurrent model-free rl can be a strong baseline for many pomdps. In International Conference on Machine Learning, 2022.\n\n[3] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through\nworld models. arXiv preprint arXiv:2301.04104, 2023.\n\n[4] Pleines, M., Pallasch, M., Zimmer, F., & Preuss, M. (2023). Memory gym: Partially observable challenges to memory-based agents. In The eleventh international conference on learning representations.\n\n> 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.\n\nWe are glad that you found our mathematical formulations easy to follow. We agree that mathematical formulations can be dense, but we think they were all needed to provide the reader with a thorough understanding of our approach and answer any question that might arise about it. We chose to delegate most of the mathematical analysis to the appendix to allow the reader to understand the main contributions easily and then dive into more details about the specifics in the appendix.\n\n>• 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?\n\nThat’s correct. The sentence should be “RTUs can be seen as a small generalization of LRUs, moving away from strict linearity”", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SElyjrKZzg", "reviewer": "Reviewer_EcqN", "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.", "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*?\n\n- 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": "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 ...", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well written overall. All the necessary background information is summarized in a clear fashion. The main idea is laid out concisely. The experiments used for evaluation are adequate and results are displayed in a visually appealing way. Finally, the direction of research is undoubtedly very significant.", "weaknesses": "A thorough ablation study of the proposed architecture is missing, i.e. taking discrete steps from LRUs towards RTUs and comparing all of them in one plot, including BPTT. \n\nThe work is not overly original, but rather an unavoidable next step after the success of LRUs.\n\nThe 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. \n\nThere are many typos and misplaced words.\n\nSomehow, the Appendix was more interesting that the actual paper. \n\nMinor issues:\n- L 29: \"... does not scalable well ...\"\n- L 32: \"... is well suited for update the state online ...\"\n- L 48: The sentence starting on this line looks like it needs a citation, please add something like \"as we will show in section...\"\n- L 143, missing word: diagonal \"matrix\"\n- L 151: \"... choices made in LRUs, that they showed ...\"\n- L 291: The sentence starting on this line is confusing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed and valuable feedback. Here, we respond to the points the reviewer mentioned.\n\n>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.\n\nThank you for raising this point. Currently, in Figure 1, we provide an ablation on several architectural choices and test them one at a time while fixing everything else. This means that the only difference between RTU and LRU lines in Figure 1 is the RTU parametrization, which is the main difference between LRU and RTU. \n\nIn the left plot of Figure 1 (titled: linear recurrence), we are testing the minimal architectural choice that could be made, a linear recurrence. Hence, this plot is only testing the utility of the RTU parametrization over LRU’s with the absence of any other factor. We then move to test different architectural choices that could be made for either of them, such as adding non-linearity or a projection layer. \n\nFor Figure 1, we chose to use RTRL as the learning algorithm rather than T-BPTT. We made this choice to prevent any other confounding factor (such as truncation length in T-BPTT) from affecting the performance of either architecture. We still tested RTUs with T-BPTT in the appendix, Figure 11, but as expected with T-BPTT, the performance of RTUs would reduce due to truncated gradients. \n\n>Somehow, the Appendix was more interesting that the actual paper.\n\nWe are glad that you found the appendix interesting. We tried to extensively answer any questions that might arise about our approach, but we still didn’t want to overwhelm the reader with a lot of mathematical details in the main body and rather focus on the core ideas, so we chose to delegate most of the mathematical work to the appendix. \n\n>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?\n\nThe discarded information from taking the real part also results in discarded gradient components. This makes the gradient estimation different from the true gradient and biased towards the information coming from the real components. We provide a simple derivation of that in Appendix D.2. It is worth noting, though, that the biased gradient issue happens when the recurrent part of LRU is not followed by a linear projection. Hence, for LRU to have an unbiased gradient estimate, it needs to always have a linear projection layer afterward, which is restrictive. \n\nAdditionally, while the original LRU paper proposes always having a linear projection after LRU, our empirical evaluations in section 4.1 show that this linearity restriction actually harms LRU's performance. \n\n>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?\n\nCurrently, Linear RTUs in Eq.2, where we apply non-linearity afterward, not on the recurrence, can still benefit from parallel scans, but the Non-Linear RTUs cannot since the recurrence part is non-linear. However, there are some recent works on parallelizing non-linear recurrence, which could be explored with Non-Linear RTUs as well [1][2].\n\n[1] Lim, Y. H., Zhu, Q., Selfridge, J., & Kasim, M. F. Parallelizing non-linear sequential models over the sequence length. In The Twelfth International Conference on Learning Representations.\n\n[2] Gonzalez, X., Warrington, A., Smith, J. T., & Linderman, S. W. (2024). Towards Scalable and Stable Parallelization of Nonlinear RNNs. arXiv preprint arXiv:2407.19115.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9zbxKZJoop", "reviewer": "Reviewer_PH1W", "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.\n\nRTU 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).\n\nLike LRU, RTU allows for tractable real-time recurrent learning (RTRL) through diagonal recurrence.\nThe focus of this work is on evaluating RTU-RTRL for reinforcement learning (RL) in partially observable environments.\n\nThe 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.", "questions": "Most of the questions have already been asked above. These are comments and suggestions.\n\n* Tractable RTRL with a diagonal recurrent matrix dates back to [R1] [R2] at least, which should be cited.\n\n[R1] Gori et al. IJCNN 1989. BPS: A learning algorithm for capturing the dynamic nature of speech.\n\n[R2] Mozer. Complex Systems 1989.  A focused backpropagation algorithm for temporal pattern recognition.\n\n* 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.\nThere are several prior works evaluating [R3] [R4] [R5] and discussing [R6] such models in the context of reinforcement learning.\n \n[R3] Irie et al. NeurIPS 2021. Going Beyond Linear Transformers with Recurrent Fast Weight Programmers.\n\n[R4] Irie et al. ICML 2022. A Modern Self-Referential Weight Matrix That Learns to Modify Itself.\n\n[R5] Pramanik et al. arXiv 2023. Recurrent Linear Transformers.\n\n[R6] Lu et al. ICML 2024. Rethinking Transformers in Solving POMDPs.\n\n* In the equation just below Line 118 (sec 3.1), *\"f(\"* is a typo at this stage (the non-linearity is not introduced yet).\n\n* 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.\n\n* 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.\n\n* The proposed RTU is not specific to RL. I'm wondering if the authors considered applying it to other supervised learning applications.", "limitations": "Comments are already provided above.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Evaluating variants of LRU using RTRL *for reinforcement learning* is novel and interesting.\n\n* Beyond the focus on the RL applications, several interesting details about LRU are discussed while developing RTU (Sec 3.2 and 3.4), which should be of interest to people interested in LRU in general.\n\n* Experiments are conducted on several relevant environments.\n\nOverall, I'm supportive of acceptance, provided that the authors will address the main issues described below.", "weaknesses": "**One major issue is the repeated claims about outperforming \"Transformers\" while no such experiment is provided, certainly not in the main text (not even in the appendix as far as I can tell).** See:\n- Abstract Line 11, \"RTUs significantly outperform GRUs and Transformers\"\n- Line 329: \"We also included a GPT2-transformer baseline\"\n- Line 330: \"We provide the results for GPT2 in Appendix G.\"\n- Conclusion: Line 356: \"performed better than the more computationally intensive GRUs and Transformers\"\n\nI could not find the Transformer/GPT2 results even in the indicated Appendix G, Figures 21 or 22 either. \n\nGiven 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.\n\nI 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.\n\n**There are also some clarity issues:**\n\n* Clearer explanations to explain the gap between RTU and LRU are expected.\n\nIn 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?\n\nSimilarly, some explanations should be provided to explain why Online LTU largely underperform RTU models in Figures 4, 5, and 6.\n\nThe 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\"*\nSo is LRU's problem just an implementation issue? Please clarify.\n\n* 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. \n\nInstead, 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.\n\n**Some of the experimental designs are not convincing:**\n\n* (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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed and valuable feedback. Here, we respond to the points the reviewer mentioned.\n\n> One major issue is the repeated claims about outperforming \"Transformers\" while no such experiment is provided.\n\nWe do provide the results for GPT-2 in Appendix G. Figure 21 shows the GPT-2 performance on Mujoco-P tasks, and Figure 26 has the learning rate sweep results for GPT-2. In Figure 21, the GPT-2 line is far below all the baselines we have, so it’s easy to miss. We changed it to a brighter color now, so it’s obvious. Please check the uploaded PDF for the updated figures.\n \nThe reason for not including GPT-2 in the main text is that the number of parameters in GPT-2 is way more than any of the architectures we have, so the comparison is a bit unfair. We will modify the abstract so as to not put a lot of emphasis on this comparison since it’s not a core result of the paper. \n\n>In Figure 1, some explanations are needed to help the readers understand why there is such a big gap between LRU and RTU \n\nOur goal from Figure 1 was twofold:\n 1. Understand the utility of RTU parametrization in contrast to LRU.\n 2. Understand the effect of different design choices, such as non-linearity and using a linear projection layer. \nBy controlling for other factors, we can ascertain that the gap in performance is due to the difference in parametrization choices, i.e., using explicit complex numbers in LRU versus using rotation matrices in RTUs. Using rotation matrices avoids the complications around taking the real part of the hidden states; we show in Appendix D.2 that taking the real part of the hidden states could result in a biased gradient estimate.\n\nFor Figure 1, we chose to use RTRL as using T-BPTT would complicate the comparison. Bad performance could be attributed to the truncation length rather than the architecture itself, while with RTRL, there is no confounding factor other than the parametrization.\n\n>Similarly, some explanations should be provided to explain why Online LTU largely underperform RTU models\n\nWhen comparing RTU with Online LRU, the core difference is RTU’s parametrization of complex numbers as rotation matrices. As we noted above, using rotation matrices avoids the complications around taking the real part of the output.  We also noticed that using complex numbers combined with auto-diff libraries resulted in unexpected results when taking the gradients. This appears to be an issue related to the use of complex numbers in general, not specific to LRU. We pointed to a similar discussion in the footnote on page 5. We also emailed Online LRU authors to verify that our implementation is correct, and we used their suggested implementation of Online LRU in our experiments.\n\n>I find Figure 2 misleading as they mix model architectures and learning algorithms within the same comparison.\n\nThe goal of Figure 2 was to highlight the trade-off between computational resources and performance when using T-BPTT versus RTRL. So, we could compare RTU-RTRL and RTU-TPBTT. However, we believe the reader would be more interested in seeing the performance of LRU-TBPTT and GRU-TBPTT to see how much benefit we get on top of these more widely used approaches when moving to RTU-RTRL. The plot becomes unreadable with too many choices. But your point is a good one, and we do provide a comparison of RTU-BPTT with LRU BPTT in Appendix E, Figure 11. As expected, when using T-BPTT with RTU, the truncation length affects the performance of RTU as well.  \n\n>Does the chosen \"computational budget of 15000 FLOPs\" (Line 237) correspond to something intuitive/useful? \n\nWe wanted to start with a computational budget under which all architectures can learn and produce relatively good performance so the comparison is meaningful. We looked at the sizes of the GRU architectures used in [1], where the animal learning experiments were introduced, and chose similar hidden dimensions to start with. That’s how we came to the choice of 15k flops specifically.\n\n[1] Banafsheh Rafiee, Zaheer Abbas, Sina Ghiassian, Raksha Kumaraswamy, Richard S Sutton, Elliot A Ludvig, and Adam White. From eye-blinks to state construction: Diagnostic benchmarks for online representation learning. Adaptive Behavior, 2020.\n\n>diagonal recurrent matrix dates back to [R1] [R2]\n\nThanks for pointing out these references. We will cite them in the introduction.\n\n> The authors' view on Transformers is restricted\n\nThanks for pointing out those relevant papers. We will modify the discussion on transformers to include them. While the approaches in [R3] [R4][R5] are all interesting and relevant, they face the same challenges as T-BPTT-based approaches in terms of the trade-off between sequence length, computational complexity, and performance. \n\n\n>Regarding the non-linearity \"f\", for the tractable RTRL to hold, \"f\" has to be a purely element-wise function. \n\nWe agree that the use of layer norm as an activation function is not conventional. The discussion we provided in Appendix A.2 was to mathematically state what conditions we need on f for the equivalence to hold. However, as we point out in the main body, we think of these theoretical results as negative results; as you pointed out, those functions are not element-wise and not conventional activation functions.\n\n> Based on [11], gating can be made compatible with RTRL. So some gating could have been a natural extension too.\n\nThank you for the suggestion. We agree that gating is a natural extension of this work. \n\n >The proposed RTU is not specific to RL.\n\nWe agree that RTUs are not specific to RL. However, we think RTRL approaches are more needed in the context of online RL than supervised learning. As in supervised learning, offline datasets are usually available, and having access to a long history is not as hard as in online RL. There are areas of online supervised learning where RTRL approaches could be useful, which could be explored in future work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "y8V9HmMrLx", "reviewer": "Reviewer_d6Pv", "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.", "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?\n- RTUs could be trained with BPTT. This would raise the question whether the increase in performance is exclusively in RTRL. \n- 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.\n\n\n[1] Time-Efficient Reinforcement Learning with Stochastic Stateful Policies", "limitations": "The authors made the main limitation of this method clear. The paper would benefit from answering the above questions. \n\n**Some minor points**:\n- line 118: where did the non-linearity f come from (compare left and right)?\n- line 148: “Each element of $h_{t-1}$” should be $h_t$\n- Why not keeping the $\\bar{h}_t$ notation from 3.1 in 3.2?\n- 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 [...]\" \n\n**Conclusion**\n\nThe 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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- the paper is sound\n- the contribution is clear \n- extensive ablations on toy experiments\n- interesting analysis", "weaknesses": "- minor contribution \n- some tasks are not very complex (Mujoco-P and Mujoco-V do not require much memory).\n- expensive in the multilayer setting, making it limited to simple tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the detailed and valuable feedback. Here, we respond to the points the reviewer mentioned.\n\n>some tasks are not very complex (Mujoco-P and Mujoco-V do not require much memory).\n\nWhile Mujoco-P and Mujoco-V could be considered short-term memory tasks, Reacher POMDP and some of the POPGym environments are long-term memory tasks. Table 1 in Ni et al. (2023) provides an estimate of the memory length needed to solve some tasks, and both Reacher-POMDP and Autoencode from POPGym have long memory requirements. Other environments we tested, such as CountRecall, Repeat First, and Concentration, also need a long history to be solved effectively. \n\n[1] Tianwei Ni, Michel Ma, Benjamin Eysenbach, and Pierre-Luc Bacon. When do transformers shine in rl? decoupling memory from credit assignment. In the Thirty-seventh Conference on Neural Information Processing Systems, 2023.\n\n>expensive in the multilayer setting, making it limited to simple tasks.\n\nIt’s true that getting the exact gradient for the multi-layer case is expensive. A common solution to this problem is to stop the gradient between the recurrent layers’ hidden states [1][2], this approach has been shown to work in practice and can be used when scaling our RTUs for multi-layers.  While this solution has been shown to work in practice, we think that a more principled approach is still needed.\n\n[1] Kazuki Irie, Anand Gopalakrishnan, and Jürgen Schmidhuber. Exploring the promise and limits of real-time recurrent learning, 2023\n\n[2]Nicolas Zucchet, Robert Meier, Simon Schug, Asier Mujika, and João Sacramento. Online learning of long range dependencies. In Advances in Neural Information Processing Systems, 2023.\n\n>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?\n\nThat’s correct, RTRL can still suffer from vanishing/exploding gradients. In our parametrization, we can prevent vanishing/exploding gradient by restricting the magnitude of the complex number to be $\\in (0,1]$. We discuss this point first in Appendix D.1, showing that different parametrization requires different restrictions to avoid vanishing/exploding gradients and that we choose the cosine representations as it has the least restrictions. In Appendix C.1, we discuss several ways to enforce the restriction on the magnitude of complex numbers and provide an ablation over these different approaches. \n\n\n>RTUs could be trained with BPTT. This would raise the question whether the increase in performance is exclusively in RTRL.\n\nThat’s true. RTUs can be trained with both BPTT and RTRL. We show an experiment in Appendix E.1 (Figure 11) comparing RTUs when trained with BPTT and RTRL. RTRL does play an important role in the improved performance of RTUs over other architectures; RTUs with BPTT will have the same issue of truncation length/computational complexity trade-off. However, the parametrization of RTUs was needed to get to a tractable RTRL.\n\n>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.\n\nThank you for pointing out this paper, we will add it to our discussion on previous work. While the paper introduces a way of calculating an unbiased gradient estimation of BPTT, they note that their estimation has  high variance since they introduce stochasticity to the policy's internal state. This is in contrast to our approach where there is no added stochasticity. Additionally, the proposed approach in [1] was not able to learn both the policy and the value function at the same time without providing additional privileged information to the value function. Our approach doesn’t have these limitations. \n\n>line 118: where did the non-linearity f come from (compare left and right)?\n\nThat’s a typo. The f shouldn’t be there at this point of the derivation. Thank you for pointing it out; we will fix it.\n\n> line 148: “Each element of ht−1” should be ht\n\nThat’s correct. Thank you for pointing it out.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "y69Y9TUIHd", "reviewer": "Reviewer_X1j7", "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.", "questions": "N/A", "limitations": "The article discusses the limitations of the method in detail, which I think is reasonable.", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-written.\n\n2. Method part has good theoretical basis and clear reasoning logic.\n\n3. The experimental results verify the effectiveness of the algorithm in a large number of scenes.", "weaknesses": "N/A", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to thank all of the reviewers for their valuable feedback on the paper. We’ve carefully considered each concern and suggestion and provided detailed responses.
While the reviewers had multiple concerns, there was no major common issue. To further address some of the reviewers' concerns regarding clarity, we added an updated version of Figures 21 and 26, changing the color of the GPT-2 baseline to be more obvious.

Meta-review: 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.
Decision: Accept (poster)