Title: Real-Time Recurrent Learning using Trace Units in Reinforcement Learning
Venue: NeurIPS 2024
Accepted: yes
Keywords: Reinforcement Learning

RNNs
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

Abstract

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

1 Introduction

Agents, animals, and people perceive their surrounding environment through imperfect sensory observations. When the state of the environment is partially observable, agents construct and maintain their own state from the stream of observations. The constructed agent state summarizes past environment-agent interactions in a form that is useful to predict and control future interactions [41]. Recurrent Neural Networks (RNNs) provide a flexible architecture for constructing agent state [19, 21, 13, 6, 10].

While standard RNN architectures have been mainly supplanted by Transformers [44], in online reinforcement learning settings where the agent learns while interacting with the environment, RNNs remain a promising direction to pursue [17, 11]. There are two main issues that limit the use of selfattention mechanisms from Transformers in online learning. First, calculating the similarity between each pair of points results in a computational complexity that is a function of k 2 , where k is the sequence length. Moreover, calculating the similarity between all pairs ignores the temporal order of the data points, which limits the usefulness of self-attention when the data is temporally correlated [47]. Second, we need access to the whole sequence of observations before taking an action or updating the learnable parameters, which is impractical in continual learning. While recent works have reduced the complexity of transformers from quadratic in the sequence length to linear [15, 16, 37], the entire sequence length is still needed to train such architectures. Gated Transformer-XL attempts to overcome this issue by keeping a moving window of previous observations [35]. A window of past observations does not scale well to long sequences-the computation is quadratic in the sequence length-and a window is one particular fixed function of history. The simpler recursive form in RNNs, on the other hand, can learn a great variety of functions of history from the data and is well

suited for updating the state online from a sequential data stream and have been shown to outperform transformers in such settings [23].

A key open question is how to efficiently train RNNs in online RL. We can divide the literature into methods that approximate Real-Time Recurrent Learning (RTRL) and those that restrict the recurrent architecture. RTRL [46] exploits the recursive nature of the gradient for RNNs, carrying forward the needed gradient information instead of unrolling the recurrent dynamics back in time like Truncated Backpropagation Through Time (T-BPTT) [45]. RTRL avoids storing past data but is so computationally expensive and is intractable for even moderately sized networks. Several methods approximate the RTRL gradient update, including NoBackTrack [33], Unbiased Online Recurrent Optimization [43, 5], Sparse N-Step Approximation (SnAp) [28]. All of these methods produce a biased gradient estimate. Other works have tried to approximate an unbiased gradient estimate of BPTT specifically for the case of policy gradient updates in RL. However, the approximation resulted in a high variance due to added stochasticity to the policy [1].

Methods in the second category usually restrict the RNN architecture to a diagonal RNN [9, 30], including Columnar Networks [18], the element-wise LSTM [17], and Independently Recurrent Neural Networks (IndRNNs) [20]. The RTRL algorithm is computationally efficient for such architectures. However, this approach sacrifices representation power and can perform poorly [18]. Recent work suggests overcoming the poor performance of diagonal RNNs with a small modification: having a complex-valued recurrent state instead of restricting it to real values [34]. In fact, as we will show in section 3.1, there exists an equivalence between using a dense linear recurrent layer and a diagonal recurrent layer with complex values, indicating no loss of representational capacity. LRUs have been combined with RTRL [48], though only empirically explored for supervised learning datasets.

In this work, we extend the insights from LRUs into the online RL setting. Our primary contribution are our experiments showing that such a lightweight recurrent architecture can outperform standard approaches like Gated Recurrent Units (GRUs) [4] in RL, with significantly less computation. To obtain this result, we propose a small extension on LRUs, which we call Recurrent Trace Units (RTUs). RTUs incorporate nonlinearity into the recurrence and use a slightly different parameterization than LRUs, but one we find is more amenable to the use of RTRL in online RL than LRUs. We extend Proximal Policy Optimization (PPO) [40] to use RTRL, ablating the decision choices we propose. We provide an in-depth study in an animal-learning prediction benchmark, showing that RTUs scale better than GRUs with increasing compute and number of parameters and that RTUs outperform alternative diagonal recurrent architectures trained with RTRL. We then show across numerous control environments that RTUs have comparable or better performance, compared to GRUs and LRUs.

2 Background

We formalize the problem setting as a Partially Observable Markov Decision Process (POMDP). At each time step t = 0 , 1 , 2 , . . . , the agent perceives an observation x t , a limited view of the state s t ∈ S , and takes an action A t ∈ A ( s t ) . Depending on the action taken, the agent finds itself in a new state s t +1 ∈ S , observes the corresponding observation x t +1 and a reward R t +1 ∈ R . In the online control setting, the agent's goal is to maximize the discounted sum of the received rewards. It may also make predictions about its environment, such as future observations' outcomes.

For prediction and control in a partially observable environment, the agent should use the stream of observations to construct its agent state . The agent state summarizes information from the history of the agent-environment interactions that are useful for prediction and control [41]. We could use the whole history up to t , namely ( x 0 , A 1 , R 1 , x 1 , A 2 , R 2 , . . . x t ) , as the agent state. Though the history preserves all the information, it is not feasible to use directly. We want the agent to have constant memory and computation per time step and storing the whole history causes the memory and the computation to grow with time. Instead, the agent needs to compress this history into a concise representation. We refer to the agent's internal representation of the history at time t as its agent state or its hidden state h t . The agent constructs its current agent state h t ∈ R n from its previous agent state h t -1 ∈ R n and the recent observation x t ∈ R d using a state-update function g : R n × R d → R n : h t = g ( h t -1 , x t ) .

One way to learn this state-update function g is with a recurrent neural network (RNN). A simple form is a linear recurrent layer, where g ( h t -1 , x t ) = W x x t + W h h t -1 for weight matrices W x and W h . We can also add a nonlinear activation, such as ReLU.

In general, we will write

where ψ are the learnable parameters in the network. The agent maps the agent state h t to an output ˆ y t and then receives a loss L t . = L (ˆ y t , y t ) indicating how far the output is from a target y t . The agent updates ψ to minimize this loss over time.

Two main gradient-based algorithms are widely used to train RNNs: Truncated Backpropagation Through Time (TBPTT) and Real-Time Recurrent Learning (RTRL). T-BPTT specifies a truncation length T , which controls the number of steps considered when calculating the gradient [45]. As a result, the computation and memory complexities of T-BPTT are functions of the truncation length. Learning with T-BPTT involves a trade-off between the network's ability to look further back in time and its compute and memory requirements. Picking a large T can be expensive, or require us to limit the network size, but picking too small of a T can cause the agent to miss important relationships and so result in poor performance.

Williams and Zipser (1989) introduced the Real-time Recurrent Learning algorithm (RTRL) as a learning algorithm for continual recurrent learning. Instead of unrolling the recurrent dynamics back in time, RTRL computes the gradient using the most recent observation, and the gradient is calculated and carried from the last step [46]. Assuming the network parameters have not changed, this recursive form gives the exact gradient and does not suffer from the truncation bias inherent to T-BPTT. We provide a more detailed background on the BPTT and RTRL in Appendix A. In reality, the agent updates its parameters frequently, so the gradient information saved from previous time steps is stale, i.e., calculated w.r.t old parameters; yet, under the assumption of small learning rates, RTRL is known to converge [46]. These properties make RTRL ideal for online learning, but unfortunately, there is a catch: its computational complexity is quartic, of fourth order, in the size of h t , which can be prohibitively expensive. For this reason, we pursue a restricted diagonal form in this work, for which RTRL is efficient and linear in h t .

3 Recurrent Trace Units

In this section, we introduce Recurrent Trace Units (RTUs). We start by revisiting why complexvalued diagonals represent dense recurrent layers, and why using real-valued diagonals is insufficient. We then introduce the specific form for RTUs that leverages this relationship. We then provide the RTRL update for RTUs, highlighting that it is simple to implement and linear in the hidden dimension. We finally contrast RTUs to LRUs and motivate why this small extension beyond LRUs is worthwhile.

3.1 Revisiting Complex-valued Diagonal Recurrence

Assume we have the recurrence relationship, with learnable parameters W h ∈ R n × n and W x ∈ R n × d , h t . = W h h t -1 + W x u ( x t ) , where u can be any transformation of the inputs x t before they are inputted into the recurrent layer. We can rewrite the square matrix W h using an eigenvalue decomposition W h = P Λ P -1 , where P contains the n linearly independent eigenvectors and Λ ∈ C n × n is a diagonal matrix with the corresponding eigenvalues. Then we have that

By defining h t . = P -1 h t ∈ C n and W x . = P -1 W x ∈ C n × d , we get a new recurrence h t = Λh t -1 + W x u ( x t ) .

We can see h t and h t are representationally equivalent: they are linearly weighted for downstream predictions, and so the linear transformation on h t can fold into this downstream linear weighting. But it is more computationally efficient to use h t with a diagonal weight matrix Λ , meaning each hidden unit only has one recurrent relation instead of n. LRUs precisely leverage this equivalence [34]. Specifically, they learn a complex-valued h t , and use Re ( Wh t ) as an input to an MLP for downstream nonlinearity.

Since we did not impose constraints on the matrix W h , other than being diagonalizable, the eigenvalues of W h can be complex or real numbers. Previous diagonal RNNs such as eLSTM [17], Columnar networks [18], and IndRNN [20] use only real-valued diagonal matrices. Having only real-valued diagonals assumes that the matrix W h is symmetric. We provide a small experiment in Appendix B.2 showing that this assumption does not hold even in the simplest setting and that complex eigenvalues do arise. We also investigate whether this result can be extended beyond linear recurrence, and largely obtain a negative theortical result (see Appendix B.1 and C ).

3.2 The RTU Parameterization

A complex number can be represented in three ways: a + bi (the real representation), r exp( iθ ) (the exponential representation), and r (cos( θ ) + i sin( θ )) (the cosine representation). Mathematically, these three representations are equivalent, but do they affect learning differently? Orvieto et al. [34] empirically showed that using the exponential representation resulted in a better-behaved loss function than the real representation on a simple task; we provide some discussion in Appendix D.1 further motivating why the real representation is less stable. We chose instead to pursue the cosine representation, because it allows us to represent the complex hidden vector as two real-valued vectors. The remainder of this section outlines RTUs, with and without nonlinearity in the recurrence.

Our goal is to learn a complex-valued diagonal matrix with weights λ k = r k (cos( θ k ) + i sin( θ k )) on the diagonal, for k = 1 , . . . , n . Multiplying by a complex number is equivalent to multiplying by a 2x2 block matrix with a rescaling. We can use this rotational form to avoid explicitly using complex numbers, and instead use two real-values for each complex-valued hidden node. We write this real-valued matrix Λ ∈ R 2 n × 2 n as blocks of rotation matrices 2

Each element of h t = Λh t -1 + W x x t ∈ R 2 n has two components h c 1 t , h c 2 t , updated recursively:

We finally combine the new recurrent states into one state h t . = [ f ( h c 1 t ); f ( h c 2 t )] , potentially using a non-linearity f after the recurrence.

We also adopt two parameterization choices made in LRUs that showed improved performance. The first is learning logarithmic representations of the parameters rather than learning them directly: instead of learning r and θ , the network learns ν log and θ log , where r . = exp( -ν ) , ν = exp( ν log ) , and θ log . = log( θ ) . This re-parametrization restricts the r to be ∈ (0 , 1] , required for stability. We found these modifications to improve stability of RTUs (see Appendix E). The second parameterization choice we adopt from LRUs is to multiply the input ( W x x t ) k by a normalization factor of γ k = (1 -r 2 k ) 1 / 2 . Putting this all together, the final formulation of RTUs is:

where γ ∈ R n is the vector composed of γ k = (1 -exp( -exp( ν log k )) 2 ) 1 / 2 and

Note that γ can be absorbed by W , and so does not change representation capacity.

There are two ways to incorporate non-linearity into RTUs: inside the recurrence or after the recurrence. In the above, in Equation (2), the non-linearity is after the recurrence. These RTUs

maintain the equivalence to a dense linear RNN, because the recurrence itself remains linear. We refer to this definition of RTUs as Linear RTUs , because the recurrence is linear, even though we have the ability to represent nonlinear functions by allowing for any nonlinear activation after the recurrence. We also evaluated a different variation of RTUs where the non-linearity is added to the recurrence directly. These Nonlinear RTUs are written as:

Notice now f -a nonlinear activation like ReLU-is used in the update to h c 1 t and h c 2 t , and the final h t simply stacks these two components. Nonlinear RTUs lose the equivalence to dense RNNs, though in our experiments, we find they perform as well or better than Linear RTUs.

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 θ ; the updates for ν log and θ log are easily obtained then using the chain rule. The full derivation is in Appendix E.2.

Consider the partial derivative with respect to r 1 for the first RTU with input ¯ x 1 . = ( W c 1 x x t ) 1 :

Then

Since r 1 only impacts the two units in the first RTU, and derivatives w.r.t. the remaining hidden units are zero. Therefore, we just need to keep track of the vector of partial derivatives for these two components, e r,c 1 t . = [ ∂h c 1 t, 1 ∂r 1 , . . . , ∂h c 1 t,n ∂r n ] and e r,c 2 t . = [ ∂h c 2 t, 1 ∂r 1 , . . . , ∂h c 2 t,n ∂r n ] with recursive formulas:

We can similarly derive such traces for θ . The update to r involves first computing ∂ L t ∂h c 1 t , using backpropagation to compute gradients back from the output layer to the hidden layer; this step involves no gradients back-in-time. Then r is updated using the gradient ∂ L t ∂ h c 1 t ⊙ e r,c 1 t + ∂ L t ∂ h c 2 t ⊙ e r,c 2 t , which is linear in the size of r ∈ R n , as the vectors e r,c 1 t , e r,c 2 t ∈ R n can be updated with linear computation in the above recursion. This update is the RTRL update, with no approximation.

3.4 Contrasting to 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 a small generalization of LRUs, moving away from strict linearity-and thus motivating the name change-but nonetheless a generalization we find performs notably better in practice.

Let us now motivate the utility of moving to a cosine representation and real-valued traces. LRUs parameterize each hidden unit with λ k = r k exp( iθ k ) = exp( -exp( ν log k )+ i exp( θ log k )) and directly work with complex numbers. Consequently, the hidden layer cannot be directly used to predict realvalues. It would be biased to take Re ( h t ) (see Appendix D.2), and instead an additional weight matrix W ∈ C n × n must be learned, to get Re ( Wh t ) . To understand why this works, assume that we took the original h t from the dense NN, and handed it to an MLP. This would involve multiplying Wh t for some W . If we set W = WP , then Wh t = WPP -1 h t = Wh t and we did not introduce any bias. In fact, if W is set this way, we do not need to take the real-valued part, because the output of Wh t is real-valued. Of course, learning does not force this equivalence-in fact this parameterization is more flexible than the original-and so it is necessary to take the real-part.

RTUs avoid some of these complications by explicitly writing the recurrence and updates with realvalued hidden states. Implicitly, the relationship between the two real-valued hidden vectors forces them to behave like complex numbers (as rotations), but all equations and learning stay in real-valued space. RTUs consequently avoid the need to post multiply by the matrix, removing a small number of learnable parameters, allowing the use of a nonlinear activation directly on the output, and allowing the hidden state to be immediately passed to a downstream MLP. We acknowledge that others may argue that working directly with complex numbers is preferable. The preference for real-valued hidden layers may simply be our own limitations, but we suspect much of the reinforcement learning community is similarly more comfortable to work in real-valued space. We found small choices in our implementation for LRU did not always behave as expected, partially due to how auto-differentiation is implemented in packages for complex numbers 3 .In the end, our goal is to make these simple recurrent traces easy to use, and providing updates with real numbers may remove some barriers.

4 Online Prediction Learning

In this section, we explore different architectural variants of RTUs and LRUs in a online prediction task and then move on to study the tradeoffs between computational resources and performance when using RTUs with RTRL compared to GRUs and LRUs with T-BPTT.

4.1 Ablation Study on Architectural Choices for RTUs and LRUs

In this first experiment, we investigate the impact of several architectural choices on the performance of RTUs and LRUs varying where nonlinearity is applied. We use a simple multi-step prediction task called Trace conditioning [39] inspired by experiments in animal learning. The agent's objective is to predict a signal-called the Unconditional Stimulus (US)-conditioned on an earlier signal-the Conditional Stimulus (CS). The prediction is formulated as a return, G t . = ∑ ∞ k =0 γ k US t + k +1 , where the agent's goal is to estimate the value function for this return. More details on this environment and experimental settings are in Appendix F. Figure 1 summarizes the results.

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 a recurrent layer followed by a linear layer generating the prediction. We measure performance of the agents online by calculating the Mean Square Return Errors (MSRE), which is the mean squared error between the agent's prediction at time t and G t . In all the experiments, we used the Adam optimizer. We first ran each agent with different step sizes for 5 runs and 2 million steps. We then averaged the MSRE over the 2 million steps and selected each agent's best step size value. Finally, we ran the agents with the best step size value for 10 runs, which we report here. We also report all agents' step size sensitivity curves in Appendix F.3.

15

30

60

Truncation Length

0.00

0.05

0.10

0.15

0.20

0.25

0.30

MSRE

Learning w. Computational

Constraints (15K FLOPS)

GRU

Linear RTUs

Non-Linear RTUs

LRU

15k

30k

60k

FLOPs

(a)

(b)

(c)

(d)

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Scaling w. Computation

Option 1

GRU

Linear RTUs

Non-Linear RTUs

LRU

2

4

8

16

32

64

128

Approximate Compute in FLOPs (×10 3 )

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Scaling w. Computation

Option 2

GRU, T=45

Linear RTUs

Non-Linear RTUs

LRU, T=45

2

4

8

16

32

64

128

Number of Parameters (×10 3 )

0.00

0.05

0.10

0.15

0.20

0.25

0.30

Scaling w. Parameters

GRU, T=45

Linear RTUs

Non-Linear RTUs

LRU

Learning under computational constraints: We first investigate: how well do different agents exploit the available computational resources? We specified a fixed computational budget of 15000 FLOPs. Since RTUs are learning using RTRL and have a linear computational complexity, the computational budget only determines the number of hidden units in the architecture. For GRU and LRU, both the truncation length and the hidden dimension contribute to the budget. We tested several configurations of truncation lengths and parameters such that the overall computations fit the computational budget. Figure 2.a shows the results of this experiment. As we move along the horizontal axis, the number of parameters for GRU and LRU decreases as T increases to fit the computational constraints. However, the RTU agents do not depend on T , so their performance and computation is constant.

Scaling with computation: The computational complexity of T-BPTT depends on the truncation length and the number of parameters in the neural network. Thus, the agent can use the additional resources in two ways: (1) Increasing the truncation length, and (2) Increasing the number of parameters. On the other hand, RTUs use all the computations to have more parameters.

Now, we move to our second question: how well do different methods scale with increasing the computational budget? We answer this question in two stages: Firstly, we study T-BPTT with increasing T and a fixed number of parameters. For RTU, the computation increases by adding more parameters such that all corresponding points from GRU and RTU use the same amount of computation. Secondly, we fixed the truncation length for GRU to 45, which is more than the maximum distance between the CS and the US, and increased the computation by increasing the number of parameters for GRU. Again, for RTU, we increased the computation by increasing the number of parameters.

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.

Scaling With Parameters: Finally, we study the performance of RTU and GRU when given the same number of parameters and allow the GRU agent to use more computation. We fixed the truncation length for GRU to 45 as before and used the same number of parameters for both agents. Figure 2.d shows the results of this experiment. For RTU, we see the same consistent performance improvement as we increase the number of parameters. For GRU, the performance improvement is also consistent, though it degrades slightly towards the end. The RTU agent outperforms the GRU agent even though the GRU uses more computation.

We provide additional experiments comparing RTUs to two other approaches that use RTRL: online LRUs and a real-valued diagonal RNN in Appendix F.1.

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 on the number of parameters. In contrast, T-BPTT methods have a complexity proportional to the truncation length T since T-BPTT methods require storing a sequence of past T activations to perform one gradient update. Figure 3 shows the time it takes to make one update with linear RTRL and T-BPTT given the same number of parameters. For T-BPTT, the time to make one update scales with the truncation length T, while for linear RTRL, it is constant.

Figure 3: Contrasting runtime in incremental and batch settings. In the incremental setting, evaluated in the animal-learning prediction task, T-BPTT updates scale with truncation length, whereas linear RTRL is constant. With batch updates, evaluated in Ant-P with PPO, linear RTRL remains linear and T-BPTT is slightly more efficient.

The computational analysis for the batch setting is different than the incremental setting. In the batch setting, linear RTRL still have a constant cost per update and provide an untruncated yet stale gradient for all the samples. When using T-BPTT in the batch setting, there are two possibilities for the gradient updates. The first option, the typical strategy, is to divide the batch into non-overlapping sequences, each of length T, and perform T-BPTT on each sequence. In this case, the cost of one gradient update per sequence is a function of T, resulting in an effective update cost per sample independent of T. As a result, T-BPTT is computationally efficient in this case, albeit at the expense of a worse gradient estimate; in each sequence, only the last sample has a gradient estimate with T steps [25]. Figure 3 shows the time it takes to make one batch update with linear RTRL and T-BPTT given the same number of parameters. In this case, both methods use similar time per update. The second option is to divide the batch into overlapping sequences, where each gradient uses a sequence of T steps [25]. This approach increases the cost of updates per sample to be proportional to T, as in the incremental setting, with the benefit of better gradient estimates. However, all standard implementations of RL methods with T-BPTT use the computationally efficient option [38, 14, 24].

Integrating Linear RTRL Methods with PPO When performing batch updates, as with PPO, the RTRL gradients used to update the recurrent parameters will be stale, as they were calculated during the interaction with the environment w.r.t old policy and value parameters. One solution to mitigate the gradient staleness is to go through the whole trajectory after each epoch update and re-compute the gradient traces. However, this can be computationally expensive. In Appendix G, Algorithm 1, we provide the pseudocode for integrating RTRL methods with PPO with optional steps for re-running the network to update the RTRL gradient traces, the value targets, and the advantage estimates. We also performed an ablation study to investigate the effect of the gradient staleness in RTRL when combined with PPO, Appendix G. The results from the ablation study show that using a stale gradient results in better performance with RTUs and suggests that the staleness might help PPO maintain the trust region.

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

accumulate and integrate the sensory observations over time to account for the missing information. 2) Remembering important cues, where the agent must remember an essential cue about the environment that happened many steps in advance.

Integrating Sensory Observations: We use the standard Mujoco POMDP benchmark widely used in prior work for evaluating memory-based RL agents [31, 12, 27, 32]. The benchmark consists of several challenging tasks where the agent controls a multi-joint dynamical body while only observing the joints' positional (Mujoco-P) or velocity information (Mujoco-V). To increase experiment throughput, we use the Jax implementation of Mujoco from the Brax library [7] and implemented wrappers to mask either the velocity (Mujoco-P) or positional information (Mujoco-V).

We evaluated our Linear and Non-linear RTUs against GRU, LRU, and Online LRU. All agents use PPO [40] as the control algorithm, and the difference between the agents is the recurrent component. For all agents, we fixed the number of parameters for the recurrent part to be ∼ 24 k. We tuned the learning rate for all agents in all environments and selected the best learning rate for each agent per environment. We also included a GPT2-transformer baseline. We followed the implementation details in previous work [32], and used a GPT2 variant with 200 k parameters. We provide the results for GPT2 in Appendix H.

When given the same number of parameters, RTU agents outperform other baselines in all environments in Figure 4. Furthermore, we show in Appendix H that even when increasing the truncation length of both GRU and LRU agents to use significantly longer history, they outperform RTUs in only one task. Of particular note is again that RTUs outperform online LRUs, highlighting again that our simple modifications have a large impact on performance in this online RL setting.

Remembering Important Cues:

Next, we test the agents' ability to remember essential environmental cues. We use several tasks from the POPGym benchmark [29] in addition to the Reacher POMDP task, a modified version of Mujoco Reacher where the agent observes the target position only at the beginning of the episode.

0.0

0.2

0.4

0.6

0.8

1.0

140

120

100

80

60

40

20

Undiscounted

Return

(30 runs)

Reacher POMDP

GRU

Linear RTUs

NonLinear RTUs

Online LRU

The POPGym tasks we consider along with the Reacher POMDP are all long-term memory tasks [32] as the agent must remember and carry the information for a long time.

Figure 5 summarizes the results for the reacher POMDP task and the POPGym results can be found in Figure 6. In both cases, we can see that RTUs outperform the other approaches. Non-linear RTUs achieve a better performance than linear RTUs in reacher POMDP, and both achieve a better performance in all tasks than online LRUs. In

Reacher POMDP, GRU was able to achieve a similar performance to that of linear RTUs.

7 Conclusion and Limitations

In this work, we investigated using complex-valued diagonal RNNs for online RL. We built on LRUs, to provide a small modification (RTUs) that we found performed significantly better in online RL across various partially observable prediction and control settings. We also found RTUs performed better than the more computationally intensive GRUs. Overall, RTUs are a promising, lightweight approach to online learning in partially observable RL environments.

A primary limitation of RTUs is the extension to multilayer recurrence. This limitation is inherent to all RTRL approaches; with multilayers, we need to save the gradient traces of the hidden state w.r.t the weights from all the preceding layers [17]. Previous work [17, 48] showed that using stop gradient operations between the layers and not tracing the time dependencies across layers is a viable choice. However, we need a more principled approach for tracing the gradient across layers.

One advantage of the linearity restriction in LRUs is that it allows the use of parallel scans for training [26]. However, recent works have shown the possibility of employing parallel scans to non-linear RNNs [8, 22]. A future direction is to investigate the use of parallel scans for training RTUs.

8 Acknowledgments

We would like to 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

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 .

Atilim Gunes Baydin, Barak A Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: a survey. Journal of machine learning research , 2018.

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.

Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. arXiv preprint arXiv:1409.1259 , 2014.

Tim Cooijmans and James Martens. On the variance of unbiased online recurrent optimization. arXiv preprint arXiv:1902.02405 , 2019.

Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, Iain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In International conference on machine learning , 2018.

C. Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax a differentiable physics engine for large scale rigid body simulation, 2021.

Xavier Gonzalez, Andrew Warrington, Jimmy TH Smith, and Scott W Linderman. Towards scalable and stable parallelization of nonlinear rnns. arXiv e-prints , pages arXiv-2407, 2024.

Gori, Bengio, and De Mori. Bps: a learning algorithm for capturing the dynamic nature of speech. In International Joint Conference on Neural Networks , 1989.

Audrunas Gruslys, Will Dabney, Mohammad Gheshlaghi Azar, Bilal Piot, Marc Bellemare, and Remi Munos. The reactor: A fast and sample-efficient actor-critic agent for reinforcement learning. In International Conference on Learning Representations , 2018.

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

Dongqi Han, Kenji Doya, and Jun Tani. Variational recurrent models for solving partially observable control tasks. In International Conference on Learning Representations , 2019.

Matthew Hausknecht and Peter Stone. Deep recurrent q-learning for partially observable mdps. In 2015 aaai fall symposium series , 2015.

Shengyi Huang, Rousslan Fernand Julien Dossa, Chang Ye, Jeff Braga, Dipam Chakraborty, Kinal Mehta, and JoÃG , o GM AraÃšjo. Cleanrl: High-quality single-file implementations of deep reinforcement learning algorithms. Journal of Machine Learning Research , 2022.

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.

Kazuki Irie, Imanol Schlag, Róbert Csordás, and Jürgen Schmidhuber. A modern self-referential weight matrix that learns to modify itself. In International Conference on Machine Learning , 2022.

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

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.

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.

Shuai Li, Wanqing Li, Chris Cook, Ce Zhu, and Yanbo Gao. Independently Recurrent Neural Network (IndRNN): Building A Longer and Deeper RNN. In Conference on Computer Vision and Pattern Recognition , 2018.

Xiujun Li, Lihong Li, Jianfeng Gao, Xiaodong He, Jianshu Chen, Li Deng, and Ji He. Recurrent reinforcement learning: a hybrid approach. arXiv preprint arXiv:1509.03044 , 2015.

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.

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 .

Chris Lu, Jakub Kuba, Alistair Letcher, Luke Metz, Christian Schroeder de Witt, and Jakob Foerster. Discovered policy optimisation. Advances in Neural Information Processing Systems , 2022.

Owen Marschall, Kyunghyun Cho, and Cristina Savin. A unified framework of online learning algorithms for training recurrent neural networks. Journal of Machine Learning Research , 2020.

Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. In International Conference on Learning Representations , 2018.

Lingheng Meng, Rob Gorbet, and Dana Kuli´ c. Memory-based deep reinforcement learning for pomdps. In International Conference on Intelligent Robots and Systems (IROS) , 2021.

Jacob Menick, Erich Elsen, Utku Evci, Simon Osindero, Karen Simonyan, and Alex Graves. Practical real time recurrent learning with a sparse approximation. In International Conference on Learning Representations , 2021.

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

Michael Mozer. A focused backpropagation algorithm for temporal pattern recognition. Complex Systems , 1989.

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.

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

Yann Ollivier, Corentin Tallec, and Guillaume Charpiat. Training recurrent networks online without backtracking, 2015.

Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences, 2023.

Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, et al. Stabilizing transformers for reinforcement learning. In International conference on machine learning , 2020.

P Ivan Pavlov. Conditioned reflexes: an investigation of the physiological activity of the cerebral cortex, 1927.

Subhojeet Pramanik, Esraa Elelimy, Marlos C Machado, and Adam White. Recurrent linear transformers. arXiv preprint arXiv:2310.15719 , 2023.

Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research , 2021.

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.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017.

Sutton. Toward a new approach to model-based reinforcement learning, 2020.

Sutton. Markov and agent state, 2020.

Corentin Tallec and Yann Ollivier. Unbiased online recurrent optimization, 2017.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 2017.

Williams and Peng. An efficient gradient-based algorithm for on-line training of recurrent network trajectories. Neural computation , 1990.

Williams and Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural computation , 1989.

Ailing Zeng, Muxi Chen, Lei Zhang, and Qiang Xu. Are transformers effective for time series forecasting?, 2022.

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.

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 BackPropagation Through Time

BPTT calculates the gradient, ∇ ψ L , by unfolding the recurrent dynamics through time and incorporating the impact of the parameters on the loss from all observed time steps. Formally, we can write ∇ ψ L as:

Applying the chain rule, we re-write Eq.5 as:

When calculating ∂ h i ∂ψ , we need to consider the effect of ψ from all the time steps. To illustrate this effect, consider unrolling the last 2 steps of the RNN dynamics:

Equation 7 shows that the network parameters ψ affect the construction of the recurrent state h t through two pathways: a direct pathway, i.e., using ψ to evaluate f ( h t -1 , x t , ψ ) , and an implicit pathway, i.e., ψ affected constructing all previous recurrent states, h t -1 , . . . , h 1 , and all those recurrent states affected h t construction. Thus, to calculate ∂ h t ∂ψ , we need to consider those two pathways:

Once again, we need to consider the two pathways when evaluating ∂ h t -1 ∂ψ in 8. For simplicity, let J t . = ∂ h t ∂ψ , B t = ∂ f ( h t -1 , x t ,ψ ) ∂ψ , C t = ∂ f ( h t -1 , x t ,ψ ) ∂ h t -1 , and re-write 8:

Writing ∂ h t ∂ψ using the results from 9:

According to Eq. 10, the agent needs to store all the previous inputs to calculate ∂ h t ∂ψ which is impractical; the computation and memory complexity will be increasing with t .

A.1.1 Truncated-BackPropagation Through Time

Williams and Peng (1990) introduced Truncated-BackPropagation Through Time (T-BPTT) which solves the issue of increasing memory and computational complexities of BPTT. In T-BPTT, we specify a truncation length T , which controls the number of steps considered when calculating the gradient in 10. We now write the truncated version of 9 which takes into consideration the gradient from the last T steps only:

Using results from 11, we then write the approximated gradient of the loss w.r.t the learnable parameters:

A.2 Real-Time Recurrent Learning

Williams and Zipser (1989) introduced the Real-time Recurrent Learning algorithm (RTRL) as a learning algorithm for continual recurrent learning. RTRL employs the recurrent formulation of the gradient in 8; instead of unrolling ∂ h t -1 ∂ψ further back in time, RTRL saves its calculated value from the previous time step and use it later when needed. It is worth emphasizing that after the agent updates its parameters, the gradient information saved from previous time steps would be stale, i.e., calculated w.r.t old parameters, however, under the assumption of small learning rates, RTRL is known to converge. The gradient formulation of RTRL can be written as:

B More Details on Representability with Complex-valued Diagonal Recurrence

This section explains why we need complex-valued diagonals to represent dense recurrent layers. We first show when it is equivalent to use complex-valued diagonal and a dense recurrent layer. We highlight that using a real-valued diagonal is like restricting the weights to be symmetric-because

the (complex) diagonal corresponds to the eigenvalues of the weight matrix-which can severely limit representability. We provide a small experiment to show that complex eigenvalues naturally arise when training both a dense linear and nonlinear RNN, further motivating the utility of moving towards complex-valued diagonals.

B.1 Representability with Complex-valued Diagonals

Let us first consider when we can perfectly represent a dense, linear recurrent layer with a complexvalued diagonal recurrent layer. Assume we have the recurrence relationship, with learnable parameters W h ∈ R n × n and W x ∈ R n × d ,

where f is a potentially nonlinear function that inputs a vector and outputs the same-sized vector and u can be any transformation of the inputs x t before they are inputted into the recurrent layer. The following equivalence result is straightforward but worthwhile formalizing.

Proposition B.1. Assume f ◦ P = P ◦ f for any full rank, potentially complex-valued P ∈ C n × n with unit-length column vectors. Then given any W h and W x for Equation (14) , there is a corresponding complex-valued diagonal weight matrix Λ ∈ C n × n and W x ∈ C n × d

where h t ∈ C n is a linear transformation of h t ∈ R n .

Proof. We can rewrite the square matrix W h using an eigenvalue decomposition W h = PΛP -1 , where P contains the n linearly independent eigenvectors and Λ is a diagonal matrix with the corresponding eigenvalues. Then, we can re-write (14) as:

where P came outside of f under our assumption that it commutes with such matrices of eigenvectors. By defining h t . = P -1 h t and W x . = P -1 W x , we get Eq. (15).

We can see h t and h t are representationally equivalent: they are linearly weighted for downstream predictions, and so the linear transformation on h t can fold into this downstream linear weighting. But it is more computationally efficient to use h t with a diagonal weight matrix Λ , meaning each hidden unit only has one recurrent relation instead of n.

Since we did not impose constraints on the matrix W h , other than being diagonalizable, the eigenvalues of W h can be complex or real numbers. Previous diagonal RNNs such as eLSTM [17], Columnar networks [18], and IndRNN [20] use only real-valued diagonal matrices. Having only real-valued diagonals implicitly assumes that the matrix W h is a symmetric matrix. [34] suggested using complex-valued diagonal matrices for better performance.

The above equivalence has only been used without any activation, namely in linear-recurrent units (LRUs) [34]. A natural question is if only the identity f ( x ) = x (and linear functions) satisfy this property of commuting with eigenvector matrices. Intuitively, this seems like the only option, as imagining a nonlinearity that commutes is hard. Surprisingly, for the more restricted case of symmetric W h , we can show a slightly more general class of activations can be used, proving an if-and-only-if relationship (see Appendix C). However, even for this restricted setting, this generalized class is limited and such activations unlikely to be preferable to a linear recurrence. We see this as a negative result, that suggests this equivalence only holds for the linear setting.

B.2 Complex Eigenvalues in Vanilla RNNs

We empirically investigate whether complex eigenvalues appear when training dense RNNs in a simple task. The goal is to show that the weight matrix, W h , is not a symmetric matrix, even in the simplest tasks. Hence, having only real-valued diagonals is too restrictive.

We used the Three State POMDP [42], depicted in Figure 7, for this experiment. In this task, the agent needs to remember one cue from the previous time-step ago, to make a prediction about the next time-step. The MDP has three states, s 1 , s 2 , and s 3 , and no actions. If the agent is in either s 1 or s 2 , it transitions to any of the three states with equal probability. However, if the agent is in s 3 , it transitions to the state preceded by s 3 . A sequence of observations would look like 1 , 3 , 1 , 2 , 2 , 3 , 2 , · · · , and we ask the agent to predict the next observation.

We trained a vanilla RNN with 3 hidden states with T-BPTT with truncation length 2 , which is a sufficient history length in this problem to predict the next observation. Since we have 3 hidden states, the matrix W h is ∈ R 3 × 3 and could have at most 2 complex eigenvalues.

We measured the performance in terms of the percentage of correct predictions made in S 3 . We recorded the number of complex eigenvalues of W h after each parameter update, shown in Figure 8. This agent reaches 100% accuracy in this problem relatively quickly. We can also see that the agent oscillates between having two complex eigenvalues and zero eigenvalues. The average number of complex eigenvalues across 30 run is above 1.5, which means that on more than 3 4 of the steps, the RNN has two complex eigenvalues. The primary point is that we see complex eigenvalues appear frequently.

0

200

400

Steps (×1K)

0.0

0.2

0.4

0.6

0.8

1.0

Prediction

0

200

400

Steps (×1K)

0.0

0.5

1.0

1.5

Eigenvalues

C More on the Equivalence of Non-Linear RTUs and Dense RNNs

As discussed in the main body, we likely only have an equivalence between using a full dense weight matrix and a complex-valued diagonal matrix for linear recurrent layers. However, we can obtain a slightly more general equivalence in the restricted setting where the weight matrix for the recurrence is symmetric. This restricted setting is not of general interest, but we include the result here because it could be of interest to a few.

For the case where we have a symmetric weight matrix, we need the activation to commute with orthonormal matrices. Consider again the form

where f is a potentially nonlinear function that inputs a vector and outputs the same-sized vector. To obtain the equivalence, assume that for any orthonormal matrix A , f ◦ A = A ◦ f . We can rewrite the square and symmetric matrix W h using an eigenvalue decomposition W h = AΛA ⊤ , where A contains the n linearly independent eigenvectors and is an orthonormal matrix and Λ is a diagonal matrix with the corresponding eigenvalues. Then, we can re-write (14) as:

where A came outside of f because commutes with orthonormal matrices. By defining h t . = A ⊤ h t and W x . = A ⊤ W x , we get:

Each hidden unit now has one recurrent relation instead of n, because our weight matrix Λ is diagonal.

A natural question is if only the identity f ( x ) = x -namely linear recurrence-satisfies this property of commuting with orthonormal matrices. We show below that it holds for a slightly more general class of recurrent layers, proving an if-and-only-if relationship. We see the below result as a negative result, highlighting that this equivalence largely only holds for the linear setting and does not generalize to other activations of interest. It provides even further evidence that likely the only setting of interest for the general non-symmetric case is also with a linear recurrence.

Nonetheless, let us obtain the if-and-only-if for completeness. A simple extension that continues to satisfies this property is f ( x ) = x c ( || x || 2 ) for any c : R → R . We can see that for any orthonormal A , we have

This means that we can have activations that rescale the input x depending on the norm of that input.

More generally, the activation can involve matrices and rotations. We can define θ ( x ) = [ x U ( x )] ⊤ where U ( x ) ∈ R n × n -1 is a matrix where the columns are orthogonal vectors to each other and to x . Then for any vector-valued g : R → R n , we have that f ( x ) = θ ( x ) ⊤ g ( || x || 2 ) also satisfies this property:

The last line follows because AU ( x ) = U ( Ax ) . Namely, for any orthogonal vector u with u ⊤ x = 0 , we have that ˜ u . = Au satisfies ˜ u ⊤ Ax = u ⊤ A ⊤ Ax = u ⊤ x = 0 because A ⊤ A = I .

Now we show this formally. Denote O n ⊂ R n × n : A ∈ O n ⇐⇒ A T A = AA T = I . Denote

Definition C.1. θ ( x ) is a matrix such that:

Lemma C.2. θ ( x ) exits and ∀ A ∈ O n , θ ( Ax ) = θ ( x ) A T

Proof. Let θ ( x ) be ∈ O n . Then, by definition of O n , θ ( x ) T θ ( x ) = θ ( x ) θ ( x ) T = I . Let θ ( x )[1] = x , and θ ( x )[2 :] be any orthogonal vectors, also orthogonal to x . Where θ ( x )[ i ] be the i th row of θ ( x ) . Then, for x ∈ R n , θ ( x ) x = ∥ x ∥ 2 e 1 .

Theorem C.3. Let f : R n → R n and A ∈ O n . Then, f ◦ A = A ◦ f ⇐⇒ ∃ g : R → R n , f ( x ) = θ ( x ) T g ( ∥ x ∥ 2 )

Proof. Define g ( α ) = f ( α 2 e 1 ) .

where h 0 is the initial hidden state.

Real Representation a + bi : Substituting λ in (22) with the real representation, we get:

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 are not preferable.

D.1 Stability

We look at the gradient when using each complex representation to understand how different representations affect learning stability. Since each hidden unit has only one recurrent relation in diagonal RNNs, it is sufficient to consider one unit, in isolation. To keep the below intuition simple, we also omit the input of x , and consider

Then it follows that the gradient w.r.t the learnable parameters a and b is:

To prevent the gradient from vanishing/exploding, we need to restrict both | a | and | b | to be ∈ (0 , 1] .

Exponential Representation r exp( iθ ) : Substituting λ in (22) with the exponential representation, we get:

and the gradient w.r.t the learnable parameters r and θ is:

To prevent the gradient from vanishing/exploding, we need to restrict r ∈ (0 , 1] .

Cosine Representation r (cos( θ ) + i sin( θ )) : Substituting λ in (22) with the cosine representation, we get:

and the gradient w.r.t the learnable parameters r and θ is:

To prevent the gradient from vanishing/exploding, we need to restrict r ∈ (0 , 1] . It is simpler to maintain stability with the exponential and cosine representations, since we only need to constrain r ∈ (0 , 1] . , whereas the real representation requires us to restrict both the complex number's magnitude and phase.

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 to a real-valued one within the LRU. However, taking only the real part results in a biased gradient, as we show in this section.

Consider again the one recurrent unit example, but now also consider the output obtained by taking only the real part of the recurrent state: y t = wRe { h t } , where w is a learnable parameter. The gradient w.r.t r and θ is:

Multiplying by a complex number z is equivalent to a rotation by the matrix [ Re { z } -Img { z } Img { z } Re { z } ]

and a scale by √ Re { z } 2 + Img { z } 2

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:

Comparing Eq. 25 and Eq. 23, we can see that using only the real part of the recurrent state leads to a loss of information in the gradient.

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 r and θ parameterizations:

Since r represents the magnitude of a complex number, then r ∈ R + , it is preferred to have r ∈ (0 , 1] to avoid vanishing/exploding gradients as discussed in the previous section. Let w r be a learnable parameter which could directly represent r or represent a function of r , we can enforce the constraints on r in several ways:

Direct learning: Learn r directly, and clip it after each parameter's update to be in (0 , 1] .

Enforcing r ∈ R + : Learn ν such that r . = exp( -ν ) . This parameterization enforces r to be positive. However, additional clipping is needed to enforce r ∈ (0 , 1] .

Enforcing stability on r : We can enforce r to be ∈ (0 , 1] by using a positive non-linear function. For example, learn ν log such that r . = exp( -exp( ν log )) , this parameterization ensures that r ∈ (0 , 1] and is suggested by Orvieto et al. [34]. Another example is learning r = σ ( ν ) , which also ensures r ∈ (0 , 1] .

Finally, we can also enforce stability on θ by learning θ log such that θ . = exp( θ log ) to ensure that θ is always positive.

9

10

11

12

13

14

15

16

17

18

0.10

0.15

0.20

0.25

0.30

0.35

MSRE

Linear RTU with 80 hidden units

r

=exp(

exp(

log

)),

=exp(

log

)

r

=exp(

exp(

log

)),

r

=exp(

),

We empirically compare the different parameterizations of r and θ . In our experiments, learning r directly resulted in unstable training where the MSRE diverges. We plot the learning rate sensitivity

9

10

11

12

13

14

15

16

17

18

0.0

0.5

1.0

1.5

MSRE

Linear RTU with 512 hidden units

r

=exp(

exp(

log

)),

=exp(

log

)

r

=exp(

exp(

log

)),

r

=exp(

),

for the other parameterization for two different sizes of RTUs in Figures 9 and 10. While all the parametrizations produce similar performance, we notice that learning θ log and ν log has better learning rate sensitivity.

E.2 Real-Time Recurrent Learning for a Single Layer Linear RTUs

We now outline the details for the RTRL update rules for RTUs. The set of learnable parameters for RTUs is ψ . = { ν log , θ log , W c 1 x , W c 2 x } . At each time step t , the learner receives a loss L t (ˆ y t , y t ; ψ ) where y t is the network output at time t , then the gradient of the loss w.r.t the parameters is:

We can derive the following gradients:

where

To efficiently compute the gradient w.r.t W c 1 x and W c 2 x , we look at the influence of each when considering a single element from each recurrent state, h c 1 t and h c 2 t :

We then get:

We see that each h c 1 t,i gets affected by weights from only one row of W c 1 x , thus, ∂ h c 1 t W c 1 x can be written as a matrix of the same dimension as W c 1 x . The same is true for ∂ h c 2 t W c 2 x , ∂ h c 1 t W c 2 x , and ∂ h c 2 t W c 1 x .

E.3 Single Layer Non-Linear RTUs Formulation:

We extend the linear RTUs to non-linear RTUs by adding a non-linear activation function f to the recurrent states. We can write the non-linear RTUs as follows:

where f : R n → R n . Following the same procedure as in the linear case, we can derive RTRL update rules for the non-linear RTUs.

We then get:

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 t ∈ R d and hidden states h t = [ f ( h c 1 t ); f ( h c 2 t )] ∈ R 2 n ,we have g ( ν,θ ) , ϕ ( ν,θ ) , γ ∈ R n and W c 1 x , W c 2 x ∈ R d × n .

An agent using the RTU with RTRL needs to store the gradient information, ∂ h c 1 t -1 ∂ψ and ∂ h c 2 t -1 ∂ψ , from one step to the next. We denote the set of saved gradient information as:

The saved gradient information has the following dimensions:

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, a forward pass according to 2 has a computational complexity of O ( n + nd ) . Additionally, after doing the forward pass, the learner needs to update the saved gradient information according to equations 27 through 31 which has a computational complexity of O ( n + nd ) . To summarize, using Real-Time Recurrent Learning with RTUs has linear computational and memory complexities.

E.5 Multi-Layers Recurrent Trace Units

We now extend RTUs to a multilayer setting. We show that in the multilayer case, we lose the computational advantages. However, prior work suggested that treating each recurrent layer independently is a sensible choice and allows us to gain a computational advantage[18]. Consider an RTU with n layers, we refer to the hidden dimension of a layer i where 0 < i ≤ n as d i . The network has the following set of parameters:

To update the network parameters, we need to calculate the gradient of the loss w.r.t the parameters from all the layers:

where

Take as an example the gradient of h c 1 ,n t w.r.t ν log ,n , ν log ,n -1 , . . . , ν log , 1 . Unrolling the last two layers of the network, we get:

The gradient of h c 1 ,n t w.r.t ν log ,n can be calculated in linear complexity as indicated in the previous section.

Calculating the gradient w.r.t the parameters from the earlier layer:

Let's define J i,i -1 = { ∂ h c 1 ,i t -1 ∂ν log ,i -1 , ∂ h c 2 ,i t -1 ∂ν log ,i -1 } . Then, to calculate the gradient of the hidden units from layer n w.r.t the parameters from layer n -1 , we need save J n,n -1 , and to calculate the gradient of the hidden units from layer n w.r.t the parameters of layer n -2 , we need to save J n,n -2 and J n -1 ,n -2 . If we keep going, to calculate the gradient of the hidden units of layer n w.r.t the parameters of the first layer, we need to save J n, 1 , J n -1 , 1 , J n -2 , 1 , . . . , J 2 , 1 .

E.6 Implementing RTRL within the reverse-mode automatic differentiation

For a function f : R n → R m , we have the Jacobian ∂f ( x ) ∈ R m × n and we can calculate this Jacobian in two ways: forward-mode or reverse-mode differentiation. In the forward-mode differentiation, the chain rule is applied to each operation while traversing the computational graph in the forward pass [2]. While computing the derivatives during the forward pass is appealing, we need to do n forward passes to get the full Jacobian, where each forward pass would allow us to compute the derivative w.r.t only one of the inputs. i.e., with forward mode differentiation, we evaluate the Jacobian one column at a time. As a result, forward-mode differentiation is inefficient for neural networks; neural networks map from learnable parameters, which can be in millions, to a loss function, hence, have very wide jacobians, n ≫ m .

Reverse-mode differentiation, on the other hand, offers a more efficient approach. It allows us to evaluate the Jacobian one row at a time, which is particularly advantageous for neural networks. This efficiency comes at the cost of two passes through the network: a forward pass for function evaluation and a backward pass for derivative evaluation [2].

RTRL is an instance of forward-mode differentiation; during the forward pass, the gradient information is evaluated along with the recurrent function computation. As a result, there is no need to perform a backward pass for the recurrent component. To efficiently use a recurrent layer with RTRL within a larger neural network, we combine RTRL for the recurrent layer with the reverse mode for the rest of the network. We use a stop gradient operation on the recurrent layer hidden state and do a normal reverse-mode differentiation. Due to the stop gradient operation, the gradient from the reverse mode assumes no time dependencies between the recurrent states. We then use the gradient traces calculated during the forward pass of the recurrent layer to correct the gradient from the reverse mode and account for the time dependencies between the recurrent states [3]. 4

F Additional Details on Trace Conditioning Experiments

In animal learning, Trace conditioning is a type of experiment where animals predict the occurrence of a stimulus (e.g., food), based on the occurrence of another stimulus like a tone. There is no prior connection between the two stimuli. However, after enough repetitions of pairing them togetherplaying the tone and then serving the food-the animal learns to anticipate food arrival when it hears the tone [36]. We use an open-source trace conditioning benchmark introduced in prior work [39]. Two signals appear sequentially: the Conditional Stimulus (CS) and the Unconditional Stimulus (US). The CS is the trigger signal, similar to the tone, and the US is the signal of interest and appears several time steps after the CS, similar to the food. The agent also observes several distractor signals which are uncorrelated with the CS and US; the agent must learn ignore them and focus only on predicting the US.

The agent's objective is to predict the onset of the US, which we model as a prediction of the discounted sum of the future US, G t :

where γ is a discount factor determining the prediction horizon. This problem is challenging because the CS appears, then disappears, and sometime later the US appears; the agent must construct an internal state that represents the time period between the two signals.

4K

8K

64K

Number of Parameters

0.0

0.1

0.2

0.3

0.4

0.5

MSRE

Block Diagonal

Linear RTUs

Online LRU

F.1 Additional Trace Conditioning Experiments

Comparison to Other RTRL-based Architectures: We now compare RTUs to other RTRL-based approaches with similar architectures: an online version of LRU [48] and a vanilla block diagonal RNN. The block diagonal is a recurrent formulation similar to RTU but ignores the relation between the learnable parameters. i.e., replaces c k in 3.2 with c k = [ a k b k c k d k ] .

The results in figure 11 indicate that these seemingly small differences between the diagonal RNNs can result in significantly different behavior. RTUs outperform online LRUs, with the differences discussed in-depth in Section 3.4. RTUs also outperform the block diagonal RNN. We emphasized using real-valued diagonals implicitly assumes symmetric matrices, but that is for a single real-value. This block diagonal has more representational capacity than the RTU. This result suggests it is beneficial for learning to enforce these constraints on the learnable parameters, that they correspond to the rotational representation of complex numbers.

Onthe role of RTRL in RTUs: To highlight the role of RTRL in RTUs, we evaluated the performance of both linear and non-linear RTUs with T-BPTT. In this experiment, we all agents use the same number of parameters; the only difference is whether they use RTRL or BPTT.

Figure 12 summarises the results of this experiment. We can see that the performance of T-BPTT approaches the performance of RTRL as the truncation length increases to cover the whole context of the task.

15

30

60

120

Truncation Length

0.05

0.10

0.15

0.20

0.25

MSRE

Evaluating T-BPTT in RTUs

Linear RTUs w. T-BPTT

LRU

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 t , the agent passes the observation o t to the recurrent layer, which outputs the recurrent state, the agent state. The recurrent state is then passed to the linear layer generating the prediction. For each agent, we swept over the learning rate α used to update the network parameters, α ∈ { 10 -1 , 10 -2 , 10 -3 , 10 -4 , 10 -5 , 10 -6 } , and averaged the performance for each learning rate over 5 independent runs. We then selected the best-performing learning rate for each agent and ran 30 independent runs using it. For all the experiments, we ran the agents for 2 million steps, and the performance was the mean squared prediction error averaged over the 2 million steps.

F.3 Learning Rate Sensitivity

We show the learning rate sensitivity for all agents in the animal learning benchmark.

0.1

0.01

0.001

0.0001

1e-05

1e-06

Learning Rate

0.5

1.0

1.5

2.0

2.5

3.0

MSRE

(5 runs)

GRU in Learning w. Computational

Constraints

(H:T)13:15

(H:T)5:60

0.1

0.01

0.001

0.0001

1e-05

1e-06

Learning Rate

0.5

1.0

1.5

2.0

2.5

3.0

GRU in Scaling w. Computation

Option 1

(H:T)13:15

(H:T)13:30

(H:T)13:60

0.01

0.001

0.0001

1e-05

1e-06

Learning Rate

0.2

0.4

0.6

0.8

1.0

GRU in Scaling w. Computation

Option 2

(H:T) 2:45

(H:T) 3:45

(H:T) 5:45

(H:T) 8:45

(H:T) 13:45

(H:T) 20:45

(H:T) 30:45

0.01

0.001

0.0001

1e-05

1e-06

Learning Rate

0.0

0.2

0.4

0.6

0.8

1.0

1.2

GRU in Scaling w. Parameters

(H:T) 20:45

(H:T) 30:45

(H:T) 48:45

(H:T) 67:45

(H:T) 97:45

(H:T) 140:45

(H:T) 200:45

-10

-11

-12

-13

-14

-15

-16

-17

-18

Learning Rate

0.16

0.18

0.20

0.22

0.24

0.26

MSRE

(5 runs)

LRU in Learning w. Computational

Constraints

(H:T)22:15

(H:T)10:30

(H:T)6:60

-10

-11

-12

-13

-14

-15

-16

-17

-18

Learning Rate

0.15

0.20

0.25

0.30

0.35

LRU in Scaling w. Computation

Option 1

(H:T)22,10:15

(H:T)22,10:30

(H:T)22,10:60

-10

-11

-12

-13

-14

-15

-16

-17

-18

Learning Rate

0.05

0.10

0.15

0.20

0.25

0.30

0.35

0.40

LRU in Scaling w. Computation

Option 2

(H:T) 2:45

(H:T) 3:45

(H:T) 5:45

(H:T) 10:45 (H:T) 20:45

(H:T) 40:45 (H:T) 80:45

-12

-13

-14

-15

-16

-17

-18

Learning Rate

0.05

0.10

0.15

0.20

0.25

LRU in Scaling w. Parameters

(H:T) 40:45

(H:T) 80:45

(H:T) 160:45 (H:T) 320:45

(H:T) 640:45

(H:T) 1290:45

(H:T) 2560:45

-8

-9

-10

-11

-12

-13

-14

-15

-16

-17

-18

log2 (Learning Rate)

0.00

0.05

0.10

0.15

0.20

MSRE

(5 runs)

RTUs in Learning w. Computational

Constraints

-8

-9

-10

-11

-12

-13

-14

-15

-16

-17

-18

log2 (Learning Rate)

0.00

0.05

0.10

0.15

0.20

Linear RTUs in Scaling w. Computation

Option 1

Linear RTUs:550

Linear RTUs:1110

Linear RTUs:2220

-8

-9

-10

-11

-12

-13

-14

-15

-16

-17

-18

log2 (Learning Rate)

0.00

0.05

0.10

0.15

0.20

NonLinear RTUs in Scaling w. Computation

Option 1

NonLinear RTUs:550

NonLinear RTUs:1110

NonLinear RTUs:2220

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 is to go through the whole trajectory after each epoch update and re-compute the gradient traces. However, this can

-10

-11

-12

-13

-14

-15

-16

-17

-18

log2 (Learning Rate)

0.00

0.05

0.10

0.15

0.20

MSRE

(5 runs)

Linear RTUs in in Scaling w. Computation

Option 2

h: 160

h: 230

h: 450

h: 800

h: 6100 h: 3200 h: 1650

-9

-10

-11

-12

-13

-14

-15

-16

-17

-18

log2 (Learning Rate)

0.00

0.05

0.10

0.15

0.20

NonLinear RTUs in in Scaling w. Computation

Option 2

h: 160

h: 230

h: 450

h: 800

h: 6100 h: 3200 h: 1650

-9

-10

-11

-12

-13

-14

-15

-16

-17

-18

log2 (Learning Rate)

0.00

0.05

0.10

0.15

0.20

MSRE

(5 runs)

Linear RTUs in Scaling w. Parameters

h:80

h:160

h:300

h:600

h:1150

h:2300

h:4600

-8

-9

-10

-11

-12

-13

-14

-15

log2 (Learning Rate)

0.00

0.05

0.10

0.15

0.20

NonLinear RTUs in Scaling w. Parameters

h: 80

h: 160

h: 300

h: 600

h: 1150

h: 2300

h: 4600

be computationally expensive. In Algorithm 1, we provide the pseudocode for integrating RTRL methods with PPO with optional steps for re-running the network to update the RTRL gradient traces, the value targets, and the advantage estimates.

In the next experiment, we investigate the effect of the gradient staleness in RTRL when combined with PPO and how this staleness interacts with the targets and advantage estimates. To understand this interaction, we evaluate all combinations of stale gradient and stale targets with increasing the number of epoch updates. We perform this analysis on the Ant-P environment from the Mujoco POMDP benchmark [31, 12, 27, 32]. Surprisingly, Figure 18 shows that using a stale gradient results in better performance with RTUs than re-computing the gradient traces. This performance improvement is also consistent when we increase the number of epochs from 4 to 8. It also shows that re-computing the value targets and advantage estimates has a minimal effect on the performance. We repeated the same experiments for NonLinear RTUs and Online LRU with consistent results in figures 19 and 20.

One hypothesis for the superior performance of stale gradients is that the staleness is helping PPO maintain the trust region. We investigate this hypothesis by measuring the KL divergence between the policy used to collect the trajectory and the most recent policy. We use the following estimate for

the KL divergence between the two policies as ( r -1) -log( r ) , where r = ( π θ new ) / ( π θ old ) . The rightmost subplot of Figure 18 shows that at the beginning of learning, agents with stale gradients move away from the old policy more slowly than agents with fresh gradients; perhaps stale gradients help with maintaining the trust region. However, this hypothesis still needs more investigation in future work.

H More Details on the Memory-Based Control Experiments

In our implementations, we use a shared representation learning network followed by two MLPs for the actor and the critic heads, as illustrated in Figure 21. The shared representations consist of a feedforward layer with 64 hidden units and memory components: an RTU, LRU, or a GRU. The actor and the critic's heads consist of two feedforward layers with tanh activation function.

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 to solve the Mujoco POMDP tasks. These results show that even when the truncation length is larger than needed, RTUs still outperform T-BPTT baselines. We also show that the transformerbased models, GPT2, perform worse than RNN-based models. This is consistent with previous work suggesting that transformers might not be suitable for RL tasks [32].

0

1

2

3

4

5

Env Steps

1e6

1000

0

1000

2000

3000

4000

5000

Undiscounted

Return

(30 runs)

Ant-P

GRU

LRU

Linear RTU

NonLinear RTU

GPT

Online LRU

0

1

2

3

4

5

Env Steps

1e6

200

400

600

800

1000

1200

Walker-P

0

1

2

3

4

5

Env Steps

1e6

200

400

600

800

1000

1200

1400

Hopper-P

0

1

2

3

4

5

Env Steps

1e6

0

1000

2000

3000

4000

Cheetah-P

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 1 for all agents.

0

1

2

3

4

5

Env Steps

1e6

1000

0

1000

2000

3000

4000

5000

Undiscounted

Return

(30 runs)

LRU

GRU

Linear RTU

NonLinear RTU

Online LRU

0

1

2

3

4

5

Env Steps

1e6

200

300

400

500

600

700

Walker-V

0

1

2

3

4

5

Env Steps

1e6

200

400

600

800

1000

1200

Hopper-V

0

1

2

3

4

5

Env Steps

1e6

500

0

500

1000

1500

2000

2500

Cheetah-V

0.00001

0.00003

0.0001

0.0003

0.001

Learning rate

0

1000

2000

3000

4000

Undiscounted

Return

(5 runs)

Ant-P

LRU, T=4

LRU, T=64

Online LRU

0.00001

0.00003

0.0001

0.0003

0.001

Learning rate

1000

1500

2000

2500

3000

3500

Cheetah-P

LRU, T=4

LRU, T=64

Online LRU

0.00001

0.00003

0.0001

0.0003

0.001

Learning rate

400

500

600

700

800

900

1000

Walker-P

LRU, T=4

LRU, T=64

Online LRU

0.000001

0.000003

0.00001

0.00003

0.0001

0.0003

0.001

Learning rate

400

600

800

1000

1200

Hopper-P

LRU, T=4

LRU, T=64

Online LRU

0.00001

0.00003

0.0001

0.0003

0.001

Learning rate

0

1000

2000

3000

Undiscounted  4000

Return

(5 runs)

Ant-P

Linear RTU

NonLinear RTU

0.00001

0.00003

0.0001

0.0003

0.001

Learning rate

1500

2000

2500

3000

Cheetah-P

Linear RTU

NonLinear RTU

0.00001

0.00003

0.0001

0.0003

0.001

Learning rate

600

700

800

900

1000

1100

Walker-P

Linear RTU

NonLinear RTU

0.00001

0.00003

0.0001

0.0003

0.001

Learning rate

800

900

1000

1100

1200

1300

Hopper-P

Linear RTU

NonLinear RTU

0.00001

0.00003

0.0001

0.0003

0.001

Learning rate

0

1000

2000

3000

4000

Undiscounted

Return

(5 runs)

Ant-P

GRU, T=4

GRU, T=64

0.00001

0.00003

0.0001

0.0003

0.001

Learning rate

1000

1500

2000

2500

3000

3500

4000

Cheetah-P

GRU, T=4

GRU, T=64

0.00001

0.00003

0.0001

0.0003

0.001

Learning rate

500

600

700

Walker-P

GRU, T=4

GRU, T=64

0.00001

0.00003

0.0001

0.0003

0.001

Learning rate

700

800

900

1000

1100

1200

Hopper-P

GRU, T=4

GRU, T=64

I Compute resources

We ran the Mujoco-P, Mujoco-V on NVIDIA P100 GPU. Each of the Mujoco-P and Mujoco-V trials took around 30 minutes to complete on a single GPU. For the POPGym experiments and animal learning experiments, we used a large CPU cluster. Each trial of the POPGym experiments took around 2 hours to complete. While each run of animal learning took around 15 minutes to complete on a single CPU with memory less than 1 GB.

0.00001

0.00003

0.0001

0.0003

0.001

Learning rate

750

500

250

0

250

500

750

Undiscounted

Return

(30 runs)

Ant-P

GPT2, T=64, h=128

GPT2, T=1024, h=768

0.00001

0.00003

0.0001

0.0003

0.001

Learning rate

300

200

100

0

100

200

300

Cheetah-P

0.00001

0.00003

0.0001

0.0003

0.001

Learning rate

200

250

300

Walker-P

0.00001

0.00003

0.0001

0.0003

0.001

Learning rate

100

200

300

400

500

Hopper-P

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We provide detailed theoretical and experimental analysis that supports our 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 in 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 as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of 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 no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Details of all proofs are provided in the appendix and referenced in the main paper.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: all hyper-parameters and experiment details are shared in the appendix and referenced in the main paper.

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 are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: A link for a 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/ public/guides/CodeSubmissionPolicy ) 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 a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) 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 a 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 in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: all hyper-parameters and experiment details are shared in the appendix and referenced in the main paper.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: we report error bars/shaded standard error regions in all our experiments.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Compute details were provided in the appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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 circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in 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 is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The research conducted here doesn't include large models or scraped datasets.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or 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?

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

The authors should state which version of the asset is used and, if possible, include a 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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are 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 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 was 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 as details about compensation (if any)?

Answer: [NA]

Justification: The research conducted here doesn't include human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or 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 to 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 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 we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

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

STRENGTHS:
1. The paper is well-written.

2. Method part has good theoretical basis and clear reasoning logic.

3. The experimental results verify the effectiveness of the algorithm in a large number of scenes.

WEAKNESSES:
N/A

QUESTIONS:
N/A

LIMITATIONS:
The article discusses the limitations of the method in detail, which I think is reasonable.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2024-07-01

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- the paper is sound\n- the contribution is clear \n- extensive ablations on toy experiments\n- interesting analysis\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\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- 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\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n**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?\n\nQUESTIONS:\nMost 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.\n\nLIMITATIONS:\nComments are already provided above.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nA 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.\n\nQUESTIONS:\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\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\nLIMITATIONS:\nThe 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 ...", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n•\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.\n\nWEAKNESSES:\n•\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”\n\nQUESTIONS:\nHowever, 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\nLIMITATIONS:\nPaper 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?", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}