Recurrent Reinforcement Learning with Memoroids
Authors: Steven Morad
Chris Lu
Ryan Kortvelesy
Stephan Liwicki
Jakob Nicolaus Foerster
Amanda Prorok
Venue: NeurIPS 2024 poster (2024)
Keywords: POMDP
reinforcement learning
memory models
recurrent neural network
Primary area: reinforcement_learning
TL;DR: We propose a new memory model framework and batching method to improve time, space, and sample efficiency in POMDPs

Abstract: Memory models such as Recurrent Neural Networks (RNNs) and Transformers address Partially Observable Markov Decision Processes (POMDPs) by mapping trajectories to latent Markov states. Neither model scales particularly well to long sequences, especially compared to an emerging class of memory models called Linear Recurrent Models. We discover that the recurrent update of these models resembles a monoid, leading us to reformulate existing models using a novel monoid-based framework that we call memoroids. We revisit the traditional approach to batching in recurrent reinforcement learning, highlighting theoretical and empirical deficiencies. We leverage memoroids to propose a batching method that improves sample efficiency, increases the return, and simplifies the implementation of recurrent loss functions in reinforcement learning.

Recurrent Reinforcement Learning with Memoroids
Steven Morad 1.2, Chris Lu", Ryan Kortvelesy", Steeeen Liwicki', Jakob Foerster), Amanda Prorok2
Faculty of Science and Technology. University of Macau, China Coomuter Science and Technology, University of Cambridge, UK
Engineering Science, University of Oxford, UK "Toshiba Europe, UK
smorad@um.co edu christopher lueexeter ox ac uk, rk6270cst. cam ac uk,
Stephan www.ckittohhiba..ccommmomm jakob jakob. foerstereeng ox ac. uk, asp450cam. ac
Abstract
Memory models such Recurrent Neural Networks (RNNs) and Transformers address Partially Observable Markov Decision Processes (POMDPs) by mapping trajectories latent Markov states. Neither model scales particularly well to long sequences, especially compared to an emerging class of memory models called Linear Recurrent Models. We discover that the recurrent update of these mod- els resembles monoid, leading us to reformulate existing models using a novel monoid-based framework that we call memorvids. We revisit the traditional ap- proach to batching in recurrent reinforcement learning, highlighting theoretical and empirical deficiencies. We leverage memoroids to propose a batching method that improves sample efficiency, increases the return, and simplifies the implementation recurrent loss functions in reinforcement learning.
Introduction
Reinforcement learning (RL) traditionally focuses on solving Markov Decision Processes (MDPs), although for many interesting problems the Markov state is hidden. Instead, receive noisy or ambiguous observations, resulting in Partially Observable MDPs. The standard approach to RL under partial observability involves summarizing sequence of observations into latent Markov state using a memory model or sequence model. Commonly used models include RNNs and Transformers.
Training Transformers RNNs long sequences is computationally expensive. Instead, prior work often splits these sequences into shorter fixed-length subsequences called segments (Figure 1). Using segments adds implementation complexity, reduces efficiency, and introduces theoretical issues. Despite these drawbacks, most prior work and virtually all existing RL libraries follow this segment-based approach. A new class of sequence models, sometimes called Linear Recurrent Models, offers much greater efficiency over long sequences than Transformers or RNNs. We posit that we can utilize these efficient models to away with segments and their associated drawbacks.
Contributions We aim to remove the need for segments in RL. First, we discover that many efficient memory models share an underlying structure reminiscent of monoids, a concept from category theory. We propose to extend the monoid into a memoroid, a mathematical framework which can represent large class of efficient memory models. Armed with the memoroid, we propose a new batching method that eliminates the need for segments.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Episode Index
(o, a, r, d' Transition
Batch Dim (B)
0
0
O
1
1
1
One Rollout
Time Dim (L)
We visualize the Segment-Basee Batching approach often used prior literature. worker collects rollout of episodes, denoted by color. Each episode is split and zero-padded produce batch of segments, each with a constant, user-specified segment length L. Episodes exceeding the specified length are broken into multiple segments, preventing backpropagation through time from reaching earlier segments. Segments contain zero padding, reducing efficiency, biasing normalization methods, and necessitating padding-aware recurrent loss functions.
Figure 1:
In particular, we
Derive memoroids for existing sequence models, as well as the discounted return and advantage
Introduce method for inline resets, enabling any memoroid to efficiently process multiple episodes
Demonstrate that using segments degrades recurrent value functions
Propose memoroid-based batching method that eliminates the need for segments
Use this batching method to improve sample efficiency and simplify recurrent RL loss functions
Preliminaries
Consider an MDP (S, A, R, where at each timestep f, an agent produces a transition T = (s, a.a,r, r.)) from interaction with the environment. We let s' E denote the current and next states and state space, E A denote the action and action space, R S x A x $ - IR denote the reward function, and T S r AS denote the state transition matrix (A denotes a distribution). In RL, our goal is to learn policy parameterized by that maps states to action distributions Ty : S AA. The agent samples an action from the policy given the current state a ~ TT(s). and stochastically transitions to the next state s' T(s, a). receiving reward R(s, s'). The optimization objective to find the parameters maximize the expected return, discounted by y; 1t=0
E, D R(5 de, 8t41)]
2.1 Rollouts, Causality, and Episode Boundaries
It is often practical to model terminal states in MDPs, such game screen in video game. In a terminal state, all actions lead back to the terminal state and the discounted return after entering the terminal state is always done flag is stored the transition T = (s, a, r, 8', d). Transitions are often used in loss functions to train the policy. However, while navigating the MDP we do not have access to the full transition - juutt the state. We receive the done flag, reward, and next state (r, d, s') at the nexl timestep. This distinction between current and next timestep becomes important when execute memoroids multiple episodes.
We mark whether state is terminal using the done flag d E {0, 1}. The
We find that our paper is more clear if we introduce begin flag b {0, 1} that is emitted alongside each observation, available during both training and rollouts. The begin flag is at the initial timestep of episode and 0 otherwise. We differentiate between a transition T (s, a, b, d) available only in hindsight, and partial transition T (s, b) emitted while navigating the MDP. To reiterate, can access T any time, but can only access I during training.

---

2.2 Partial Observability
partially observable settings, we cannot directly measure the Markov state s. Instead, we indirectly measure via the observation
(ss), following the observation function O S - AO. With the observation replacing the state, interaction with the environment now produces a transition
P = (o, a, d', d, 6) and partial transition P [o. b). For certain tasks, the action from the previous timestep is also necessary, and is implicitly included in the observation.
A sequence of transitions starting where 1 and continuing until terminal state is known as an episode E. We use a memory model M to summarize the corresponding sequence of partial transitions into latent Markov state.
M:P'A T:P"+S' S".
(1)
If M is recurrent, we may alternatively write M single update batched update respectively
M M:
M M:HxP'HH'X
(2)
where is the set of recurrent states.
3 Background and Related Work
In deep learning, we often wish to train memory models over batches of sequences. For a single sequence, we can Backpropagation Through Time (BPTT) (Werbos, 1990). If the sequences differ in length, it is not immediately clear how to efficiently combine them into a batch. Williams & Peng (1990) propose Truncated BPTT (T-BPTT), which enables backpropagation over fixed-length sequences that we call segments. T-BPTT is the defacto standard for training memory models in both supervised learning and RL (Hausknecht & Stone. 2015; Kapturowski al., 2019; Hafner et al., 2023; Bauer et al., 2023; Liang et al., 2018; 88; Raffin et al., 2021; Huang et al., 2021; Serrano-Muice al., 2023; Lu et al., 2022; Ni et al., 2024).
In Segment-Bass Batching (SBB), we split and zero pad episodes so that they can be stacked into tensor with batch and sequence length dimensions B L. Each row in this tensor is segment a containing exactly I. transitions. Episodes longer than I, transitions will be split into multiple fragments, such that each is at most I, transitions. Fragments shorter than 7. transitions will be zero padded from the right, such that they become exactly length T.. We call these padded length T, fragments segments. We must also store mask m denoting which elements are zero-padding and which are data. The segments and masks are stacked along the batch dimension, creating B x L matrices for storage and training (Figure 1). We formally define SBB in Appendix C.
The Shortcomings of Segments SBB introduces number of shortcomings. (1) The zero padding and associated masks must be stored, taking up additional space. (2) The zero padding is fed to the memory model, wasting computation zero padding also prevents the use of BatchNorm (loffe & Szegedy, 2015) and other normalization methods by shifting the mean and variance of input data. (4) The extra time dimension and padding complicates RL loss functions. (5) Most importantly, when SBB splits episodes into distinct segments, approximate the true BPTT gradient with the T-BPTT gradient.
zeros that discarded during gradient computation. (3) The
Let us demonstrate the effect of SBB on the memory model gradient. Assume we have a loss function L. defined over model parameterized by O. We define the true gradient of the loss over an episode of length n as V.. In SBB, we split an episode into length L segments. We approximate the gradicnt over these segments V.L
[n/L-1] 8C(0, (P Pain(((t))-1n-1))) (3) V.C= =
VC AC(O, (Po, Pi, Pa-1))
Σ
00
00
Under SBB, we compute the gradient independently for each segment. The gradient across segment boundaries is therefore always zero. With zero gradient, it is unlikely that temporal dependencies greater than the segment length can be learned. In fact, our experiments show that V, often poor approximation of V.
3

---

Alternatives to Segments We are not the first to realize the drawbacks of SBB. Hausknecht & Stone (2015) store recurrent states the replay buffer, while Kapturowski et al. (2019) replay the previous segment to generate a "warm" initial recurrent state for the current segment. These methods improve the return, highlighting issues with zero-initialized states, but do not fix the underlying gradient truncation issue. Real Time Recurrent Learning (RTRL) is an alternative to BPTT, but it has O(n") time complexity and is thus much slower (Williams & Zipser, 1989). Irie et al. (2024) propose faster version of RTRL for RL, but the model must be at most one layer deep. Similar to our work, Lu et al. (2024) avoids truncating backpropagation entirely. They find that this results in greater returns, but do not explore why this occurs. Furthermore, their method is restricted to on-policy methods and the S5 memory model. Our method extends Lu et al. (2024) to off-policy algorithms and large majority of efficient memory models.
3.1 On the Efficiency of Sequence Models
SBB evolved alongside RNNs in RL (Hausknecht Stone, 2015), and Transformers to a lesser extent. Such models are only tractable when the sequence length L is small. RNNs rely on the previous recurrent state compute the following recurrent state, prohibiting parallelism over the time dimension. Thus, RNNs are unable to exploit the parallelism of modern GPUs over the time dimension. Transformers use pairwise attention on the sequence elements, scaling quadratically in space on the length of the sequence.
recent class of models espouse time-parallel execution while being either linear subquadratic in space complexity. These models, such as State Space Models, Linear Transformers, Fast Weight Programmers, RetNet, RWKV, Linear Recurrent Units, Gated Impulse Linear Recurrent Networks, and Fast and Forgetful Memory (Gu et al., 2021; Smith et al., 2022; Schlag et al., 2021; Anonymous, 2023; Peng et al., 2023; Orvieto al., 2023; Martin & Cundy, 2018; Morad al.. 2023b) are sometimes called Linear Recurrent Models because they usually (but not always) employ Linear Fime-Invarian (LTI) recurrent state update, which can be computed in parallel over the time axis (Gu & Dao, 2023).
Monoids Prior work efficient sequence modeling primarily updates the recurrent state using linear functions (Schlag et al., 2021; Gu et al., 2021: Smith et al., 2022; Orvicto et al., 2023). However, works like Blelloch (1990); Martin & Cundy (2018); Morad et al. (2023b) show that it is possible to create efficient models using nonlinear recurrent updates. The key to efficiency is not that updates linear, stated in Gu & Dao (2023), but rather that the recurrent update obeys the associative property. More formally, the recurrent update must be a monoid (Bourbaki, 1965). Hinze (2004) shows that all monoids have time-parallel implementations.
Definition 3.1. A tuple (H, er) is a monoid if:
(a c
b, H a,b, H
The binary operator is closed on H The binary operator associative
(4)
(a c (e7 a) (a eT)
(5)
a, E H
There exists identity element ey
(6)
where for single input a is defined as a) (er a).
Any monoid operator can be computed in parallel across the time dimension using a parallel scan (Appendix J) (Hinze, 2004; Dhulipala et al., 2021). Given sequence of length work-efficient parallel sean known as the Blelloch Scan executes O(n) calls to * in ((n) space to produce n outputs (Blelloch, 1990). With p parallel processors, the parallel time complexity of the scan is O(n/p + log For large GPUs where n = p, the parallel time complexity becomes O(logn).
Approach
While powerful, standard monoids are restrictive and cannot represent most Linear Recurrent Models in their entirety. Monoids require that the input, output, and recurrent space be identical. In memory models, we often decouple the input space, from the recurrent state space H, from the Markov state space S (Equation 2). Consider, for example, a navigation task where the input an image, the recurrent state H is latent map representation, and the Markov state S is a set of æ, y coordinates of the agent. In search of a more general memory model framework, we extend the monoid into a memory monoid, or memoroid for short.

---

Definition 4.1. oH)), constitute a memoroid if (H, , er)) defines a monoid and functions f.g are:
P:P-M
Mapping from a partial transition to the right argument of (7)
g: H P
Mapping a recurrent state and partial transition to a Markov state
(8)
Recall that partial transition consists of the observation and begin flag P = (o,b). The memoroid defines a recurrett memory model (Equation 2) over a sequence of partial transitions to produce recurrent states (ho, /ho,/1,--0 .i,... E H) and then compute Markov states (50; 81, ... E S)
ho h1 h2 " 80
El l(Po) gho. Po)
er f(Po) S(P)) er s(Po) f(P)) . f(22) (h2 P2)
P1h
...
Given n inputs, functions/ and g can each be split into concurrent threads. Recall that monoids have Ollgg n) parallel time and (n) space complexity. Consequently, all memoroids have Olog parallel time complexity and (n) space complexity on the length of the sequence'.
Reformulating Existing Sequence Models As an exercise in the flexibility of our memoroid, we rewrite the Linear Transformer, the Simplified State Space Model, the Linear Recurrent Unit, and Fast and Forgetful Memory (Katharopoulos et al., 2020; Lu al., 2024; Orvieto et al., 2023; Morad al., 2023b) as memoroids Appendix G. We note that our monoid reformulation of Morad et al. (2023b) improves upon the original, exhibiting better numerically stability by replacing exponentiated cumulative sums with a Bielloch scan.
Accelerated Discounted Returns Memoroids can model other recurrences well. For example, we can rewrite the discounted return and Generalized Advantage Estimate (GAE) (Schulman et al., 2016) as a memoroids. Reformulating the discounted return and GAE targets memoroids enables us to compute them in a GPU-efficient fashion, using a high-level framework like JAX (Bradbury al., 2018). We find that we can compute such quantities orders of magnitude more quickly than the standard approach. We provide the definitions and proof's of these formulations in Theorem D.I and Theorem E.
Inline Recurrent State Resets So far, we have assumed that we operate memoroids over single episode using the Blelloch Scan. To scan over batch of variable-length episodes, we could truneate and zero pad sequences such that each is fixed length (i.e., SBB). However, this introduces the issues explained Section 3.
Since memoroids are efficient over long sequences, we could consider concatenating individual episodes into one very long sequence, removing the need for padding and truncation. Unfortunately, as the scan crosses episode boundaries, it feeds information from all prior episodes into future episodes, and information from future episodes into preceding episodes.
To resolve this issue, we propose a resettable monoid transformation, which prevents information from leaking across episode boundaries. We can apply this transformation to any monoid (or memoroid), to produce new monoid that respects episode boundaries.
Theorem 4.2, All monoids (H, e,r) can be transformed into a resettable monoid (G, (G,) defined as
G ((A A H,D (0.1))
(10)
91 = (),,0) (A. o (A', = ((A ((A-(1-1)) - + el - b')
(11)
(12)
For a single episode, the term output by the operator equivalent the output of Over multiple contiguous episodes, e prevents information flow across episode boundaries.
Proof. See Appendix F.
Asssming (1) The binary operator and f. are constant- time and constant- space, which is the case for all Linear Recurrent Models listed thus far. (2) Our processor has parallel threads of execution.
5

---

Episode Index
E1
E2
D
0
O
0
1
3
3
I
7
Figure 2: A visualization of sampling TBB, with a batch size of = 4. Transitions from rollouts are stored in-order in D, with each color denoting a separate episodes. Associated episode begin indices are stored in I. We sample a train batch by randomly selecting from For example, we might sample 4 from I, corresponding to En in red. Next, we sample 7 from I, corresponding to Es in red. We concatenate B = concat B1, E2) and return the result as a train batch.
By transforming the monoid (H, e1) within memoroid, we no longer require separate time and batch dimensions during training. Now, memoroids can process long sequence comprised of many distinct episodes. Unlike Blelloch (1990); Lu et al. (2024) which provide a reset operator for specific model, our resettable transformation works for any monoid.
4.1 Tape-Based Batching
Training over the concatenation of episodes would be intractable using Transformers RNNs due poor sequence length scaling, while Linear Recurrent Models leak information between episodes. By combining the efficiency of memoroids with our resettable transform, resolve these issues, enabling us to fold the batch and time dimensions into a single dimension. We call this approach Tape-Base Based Batching (TBB), which consists two operations: insertion and sampling. We provide operations for both on-policy and off- off-policy RL algorithms. Furthermore, we design TBB in manner that greatly simplifies the implementation of recurrent loss functions.
Insertion In RL, we must store the training data (transitions) that collect during rollout. TBB stores them following Algorithm 1. We maintain ordered lists of transitions D, and begin indices I, corresponding to episode boundaries in D. Upon reaching the maximum capacity of D, discard old transitions by popping the episode begin index from the left of I, and discarding the resulting episode in D. This is guaranteed discard the oldest episode in D.
This method works both for rollouts that contain complete episodes (d,-1= = 1), and those that contain incomplete episodes (dz-1 * 1), where rollout might stop before finishing episode. When combining incomplete episodes with multiple rollout workers, we can experience race conditions. In this scenario, it is easiest keep one D, I per worker prevent conditions.
Sampling Once we have constructed D, I, we are ready to train policy or compute returns. We sample transitions from D.I following Algorithm 2. we are training on-policy, we can simply train on D. If we are training off-policy, we randomly sample training batch B from our dataset by slicing D using randomly-sample sampling approach to implement Prioritized Experience Replay (Schaul et al., 2015) by assigning each episode or index a priority.
sequential pairs of episode indices from One could extend our
Simplified Loss Functions With TBB, utilize unmodified, non-recurrent loss functions train recurrent policies, reducing the implementation complexity of recurrent RL algorithms. Unlike SBB, there is no need to mask outputs handle additional time dimensions like in SBB. With TBB, the only difference between recurrent and nonrecurrent update is precomputing the Markov states 8, s' before calling the loss function. We demonstrate this by writing the TBB Q learning update in Algorithm 3, highlighting departures from the standard, non-recurrent Q learning update in red. For posterity, we define the standard SBB Q learning update Algorithm 4. Note that the SBB update equation has an additional time dimension & and requires padding mask mi.j.

---

Algorithm Inserting transitions using TBB
Input: List transitions D, list indices Z, buffer size D Output: List of tramsitions D, list of indices p < (Po. Pi...., Pn-1)
D Collect rollout from env
if policy then
D
+ where(bn..... bn-1)
> Indices of new episodes
else
while (D card(p)) > D) do
I I D < D(7
Buffer full, pop oldest index b Pop transitions for oldest episode
end while
I 4- concat(I, card(D) + where(lo........ D con
D Update replay buffer indices
D Add new transitions buffer
end
Algorithm 2 Sampling transitions using TBB
Input: List transitions D, list indices I, batch size
Output: Batch of transitions
B t ()
> Empty list
while (B) < BBC do
i~ U(0 card(T)
Randomly sample index in Append episode to batch
B concat(B, D[I Iii 1]])
end while
B B|: B
Make batch exactly B transitions
Algorithm 3 TBB deep Q update
Input: params 0, target params o, Q function Q, sequence model M, train batch B, discount update rate B
Output: params 0,0 Ø
($1, 82, 1...........
Ma ..... PR)
Estimate Markov state
14,,8,.... sg) Ma ...... Pa]
Estimate next Markov state
91 = + maXa6A
vB
& Compute target
0 - min (s, a,) - 8,11
bB[j
D. Compute loss and update parameters > Update target network params
63 + (1 3)0
Algorithm 4 SBB deep Q update
Input: params O, target params di, Q function Q, sequence model M. train batch B, discount update rate 8 Output: params f,0 ø
$1.1 $1,4
Mo((Pi.a) (Pi,2))
> Estimate Markov state
14g1 #B.L
Mo((Pp,)) -. (Po.c)) Mc((P).)) (P1.2))
$1.1 $1.4
Estimate next Markov state
Mgl(Pm.) (Pn,)))
11.b = (Ty, max (s)
vB[j. k]
Compute target with extra time dimension
0 - ming m...k IQ (85. ag.4) 85.4 ]B
Compute loss and update params & Update target network params
á 63 (1 3)0
Experiments and Discussion
We begin our experiments by investigating the shortcomings of SBB, specifically the theoretical issues stemming from truncated BPTT. We then compare TBB to SBB across a variety of tasks and models. Finally, we examine the wall-clock efficiency of memoroids.
Our experiments utilize tasks from the POPGym benchmark (Morad al., 2023a), and all TBB to SBB comparisons use identical hyperparameters and random seeds. We validate our findings across Simplified State Space Models (S5), Linear Recurrent Units (LRU), Fast and Forgetful Memory (FFM), and the Linear Transformer (LinAttn) memoroids. We train policies using Double Dueling

---

Batch Mode S88, L-10
Batch TBB.
Nom. Cumulelive K2to,
1.0
Model
LinAttn
S5
-102
-82
-22 42 -22
-2
-102 -82
-62 -42 Age
-22
2
Observation Age
Figure 3: We demonstrate that SBB can hurt Q learning through truncated BPTT. We examine the Repeat Previous task, with RMI. 10, comparing SBB (left) to TBB (right). For SBB, we set == RMI. 10 capture necessary information, After training. we plot the cumulative partial derivative with respect to the observations the y-axis. This partial derivative determines the VML- how much each prior observation contributes to the Q value. We draw a vertical red line at L RML = 10. We see that across models, majority of the Q value is not learnable when using SBB. Even when we set L = oc using TBB, we see that the VML still spans far beyond the RML. This surprising finding shows that truncated BPTT degrades recurrent value estimators.
DQN (Van Hasselt al., 2016; Wang al., 2016). See Appendix H for architecture and training details.
What are the Consequences of Truncating BPTT? In Section 3, we discussed how the estimated (truncated) gradient used in SBB differs from the true gradient. We aim to determine whether the estimated gradient used in SBB is sufficient approximation of the true gradient. We note that if all episodes are a fixed length, and we set L to be this length, both SBB and TBB produce identical results although this is rare in practice.
We utilize the Reward Memory Length (RML) and Value Memory Length (VML.) metrics from Ni et al. (2024). The RML refers to the maximum temporal dependency j required to predict the expected reward, while the VML determines at which point & prior observations stop affecting the Q value. In other words, the environment defines the RML while the memory model defines the VML. We hope to find that VML RML; that the Q value only depends on the necessary history.
E [R(ss a, ou:n] = E [R(s, a, s') Oj:nl Q(M(oo a) ((((s))))
(RML)
(13)
(VML))
(14)
We examine the VML for the Repeat Previous task with a known RML of ten timesteps. We measure the VML as the impact each observation has on the terminal Q value of an episode (i.c., Q(5n, an,) R[(s,, ans 8n+))). Any observations older than ten timesteps are not necessary to predict the reward, and given a relative-time policy, should have little to no impact on the terminal Q value. Recall that we can write a memory model as Sn, (oo.....o on). We explicitly compute
0Q(5n, an) 0Q( an) 8sn
(15)
Oo:
Üsn 3o:
and plot the results for FFM, S5, LinAttn, and LRU models in Figure 3. We examine the VML and RML of other model and task combinations in Appendix B.
Surprisingly, the VML differs significantly from the RML. The RML is fixed at ten, while the VML appears span the entire episode. This means that the recurrent models unable to ignore uninformative prior observations, suggesting that truncated BPTT degrades recurrent value estimators. Learned recurrent Q functions do not generalize well over time, although we find that policies trained with TBB tend to generalize better.
In the case of FFM, roughly 90% of the Q value is produced outside the segment boundaries. where truncated BPTT cannot reach. This means that these unnecessary prior observations have a significant impact our Q values, yet we unable forget ignore such observations. Our findings suggest that SBB could be a major contributor to the increased difficulty and reduced sample efficiency of recurrent RL, as we demonstrate in the next experiment.
Is TBB More Sample Efficient? For our second experiment, we measure the difference in sample efficiency between TBB and SBB. There are two reasons that TBB could improve upon SBB sample
8

---

efficiency: (1) As previously discovered, the truncated gradient used by SBB is often poor estimate of the true gradient (2) SBB decreases the effective batch size through zero padding. We note that the cost of zero padding in SBB is equívalent to the cost of real data it takes up equivalent space in the replay buffer and takes just as much compute to process as real data. We report some combinations of model and task in Figure 4 and present the full results in Appendix A.
S6/Count Recall
LinAltrvPos. Caripole
RRUiRepeal First
FFMRRepeal Previous
Batching
Return
-0.50
TBB
Se8 L=10
0.75
-0.5
SB8 L=50
5000
10000
5000 10000
2000 4000
2000
4000
SBB L-100
Train Epoch
Train Epoch
Train Epoch
Train Epoch
Figure 4: We compare TBB (ours) SBB across POPGym tasks and memory models, reporting the mean and 95% bootstrapped confidence interval of the evaluation return over ten seeds. We find that TBB significantly improves sample efficiency. See Appendix A for more experiments.
We find that TBB produces noticeable improvement in sample efficiency over SBB, across nearly all configurations of memory model and environment. Even for large segments lengths = 100, find significant gap between SBB and TBB. SBB must make inherent tradeoff it can long segments to improve gradient estimates at the cost of smaller effective batch sizes, or shorter segments improve the effective batch size the expense of worse gradient estimate. TBB does not need to make this tradeoff. In our experiments, SBB with larger always outperforms shorter L, suggesting that the gradient estimation error is larger contributor SBB's lackluster performance than reduced effective batch sizes.
Wall-Clock Efficiency In Figure 5, we investigate the wall-clock efficiency of memoroids. We find that memoroids compute the discounted return and GAE roughly three orders of magnitude faster than a standard implementation. Next, we compare the wall clock time TBB and SBB take to train a policy from start finish. For SBB, the parallel time complexity O(log L) while TBB has Olog B) complexity where B > L, but in practice there is no perceivable difference in wall-clock time. One possible reason for this discrepancy that SBB applies expensive saii--nd-pat operations to the trajectories. while TBB does not. Each fragment contains a varying number of transitions, which corresponds to varying amount of padding we need to add. Variable-size operations are generally slow and difficult to batch efficiently.
Retur Rundine
Method
Train Time (s)
Std. Dev. (8)
Thes oo
Mehhod
Seral
SBB L=10
886.39
54.47
Moreld
SBB L=20
886.30
49.71
Max Ep. Length
SBB L=50
887.58
50.29
100
1000
SBB L=100
886.25
49.87
Num, Eplcodes
TBB
886.87
53.21
Figure 5: (Left) We compare how long it takes to compute the discounted return using memoroid, compared to the standard way of iterating through a batch. Computing the discounted return is orders of magnitude faster when using memoroid implementation. (Right) we compare the total time to train a policy on Repeat First. For both experiments, we evaluate ten random seeds on a RTX 2080Ti GPU.
Limitations and Future Work According sensitivity analysis, old observations unexpectedly impacted the Q value across models and tasks. Moving forward, we suggest testing newly- dewly-designed memory models to see whether VML RML, determine whether such models truly generalize over time.
In our experiments, we focused on long-term memory tasks from the POPGym benchmark, each of which tests a specific aspect long-term memory. We did not experiment on environments like Atari, primarily because it is unclear to what extent Atari tasks require long-term memory.

---

Although memoroids scale well to long sequences, TBB still pays an increased log B time cost compared with SBB's log cost. There was no perceptible difference in our experiments, but very long sequences such as those used for in-context RI could incur more noticeable training costs. TBB does not strietly require memoroids, but would likely intractable for RNN or Transformer-based memory models.
6 Conclusion
We introduced memoroids as unifying framework for efficient sequence modeling. We found that memoroids represent large number of efficient recurrent models, well as the discounted return and the advantage. Using our resettable transformation, we extended our approach to encompass batching across variable length sequences. Given the efficiency of memoroids over long sequences, we questioned whether the standard split-and-pad approach to POMDPs was still necessary. We found that said approach causes issues, with shorter segment lengths hampering sample efficiency and ultimately converging to lower returns. We proposed a simple change to batching methodology. that when combined with memoroids, improves sample efficiency at a negligible cost.
Acknowledgements
We gratefully acknowledge the support of Toshiba Europe Ltd. This work was also supported in part by ARL DCIST CRA W911NF-17-2-0181 and European Research Council (ERC) Project 949940
(gAla). We thank Matteo Bettini for suggesting the term "memoroid".
10

---

References
Anonymous. Retentive Network: A Successor to Transformer for Large Language Models. In Submitted to The Twelfth International Conference on Learning Representations, 2023. URL
ittps://openreww
Jakob Bauer, Kate Baumli, Feryal Behbahani, Avishkar Bhoopchand, Nathalic Bradley-Schmice Michael Chang, Natalie Clay, Adrian Collister, Vibhavari Dasagi, Lucy Gonzalez, Karol Gregor, Edward Hughes, Sheleem Kashem, Maria Loks-Thompson, Hannah Openshaw, Jack Parker- Holder, Shreya Pathak. Nicolas Perez-Nieves, Nemanja Rakicevic, Tim Rocktäschel, Yannick Schroecker, Satinder Singh, Jakub Sygnowski. Karl Tuyls, Sarah York, Alexander Zacherl, and Lei M Zhang. Human-timess timescale adaptation in open-ended task space. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp- 1887-1935. PMLR, 23-29 Jul 2023. URL https::/pproceddinn
mlr peress/v202/ htt.02//bauee2aaaa html.
M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling. The Arcade Learning Environment: An Evaluation Platform for General Agents. Journal of Artificial Intelligence Research, 47:253-22 279,, June 2013. ISSN 1076-9757. doi: 10. 10.1613/jair.912.. URL https: htops:://www.. jair wacar.org/innee /indexx.
pptp//aai//atticll//ee/////
Guy E Blelloch. Prefix Sums and Their Applications. Technical report, School of Computer Science, Carnegie Mellon University, November 1990.
Nicolas Bourbaki. éléments de mathématique. Integration, Livre 1, Livre 5, 1965. Publisher:
Hermann Paris.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milnn and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
httt://githur com/ google/jax.
Laxman Dhulipala, Guy E. Blelloch, and Julian Shun. Theoretically Efficient Parallel Graph Algo- rithms Can Be Fast and Scalable. ACM Trans. Parallel Comput., 8(1), April 2021. ISSN 2329-4949. doi: 10. 10/3444/343 URL stttps://dol org/10 10 11445/34334393, Place: New York, NY,
USA Publisher: Association for Computing Machinery.
Albert Gu and Tri Dao. Mamba: Linear-tine sequence modeling with selective state spaces, arXiv preprint arXiv:2312.00752, 2023.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining Recurrent, Convolutional, and Continuous-tima Models with Linear State Space Layers. In Advances in Neural Information Processing Systems, volume 34, pp. 572-585, Cur- Associates, Inc., 2021. URL https: htL.httss://proceennn www..pp.cc.pppprr////... 35546bessaa117ccs0taaaccGGbo Absttaaattttttrtar html.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering Diverse Domains through World Models. arXiv preprint arXiv:2301.041104 2023.
Matthew Hausknecht and Peter Stone. Deep Recurrent Q-Learning for Partially Observable MDPs. In 2015 AAAI Fall Symposium Series, September 2015. URL https:: //w::///wwwww aaai org/oca/ index, pttpppp/sss....... hwtps//aape//v//ii pappe//iese/11733
Ralf Hinze. An algebra of scans. In Mathematics of Program Construction: 7th International Conference, MPC 2004, Stirling, Scotland, UK, July 12-14, 2004. Proceedings 7, pp. 186-210. Springer, 2004.
Shengyi Huang, Rousslan Fernand Julien Dossa, Chang Ye, and Jeff Braga. CleanRL: High- quality Single-file Implementations of Deep Reinforcement Learning Algorithms. 2021. _eprint: 2111.08819.
11

---

Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. pmlr, 2015.
Kazuki Irie, Anand Gopalakrishnan, and Jürgen Schmidhuber. Exploring the Promise and Limits of Real-Time Recurrent Learning, February 2024. URL http://arxiv orggabs//////eeooo arXiv:2305. 19044 [cs].
aba//b///// 2305 19044.
Steven Kapturowski, Georg Ostrovski, John Quan, Remi Munos, and Will Dabney, RECURRENT EXPERIENCE REPLAY IN DISTRIBUTED REINFORCEMENT LEARNING. pp. 19, 2019.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers
RNNs: fast autoregressive transformers with linear attention. In Proceedings of the 37th Interna- tional Conference on Machine Learning, ICML '20, pp. 5156-5165. IMLR.org. July 2020.
Bric Liang, Richard Liaw, Robert Nishihara, Philipp Moritz, Roy Fox, Ken Goldberg, Joseph Gonzalez, Michael Jordan, and Ion Stoica. RLlib: Abstractions for distributed reinforcement learning. In International Conference on Machine Learning, pp. 3053-3062. PMLR, 2018
Chris Lu, Jakub Kuba, Alistair Letcher, Luke Metz, Christian Schroeder de Witt, and Jakob Foerster. Discovered policy optimisation. Advances in Neural Information Processing Systems, 35:16455- 16468, 2022.
Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, and Feryal Behbahani. Structured state space models for in-context reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Eric Martin and Chris Cundy. Parallelizing Linear Recurrent Neural Nets Over Sequence Length. In International Conference on Learning Representations, 2018
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, and others. Human-level control through deep reinforcement learning. nature, 518(7540):529-33 533, 2015. Publisher: Nature Publishing Group.
Steven Morad, Ryan Kortvelesy, Matteo Bettini, Stephan Liwicki, and Amanda Prorok. POPGym: Benchmarking Partially Observable Reinforcement Learning. In The Eleventh International Conference Learning Representations, 2023a. URL chDrutUTsOK.
net/ httemformmmccc
Steven Morad, Ryan Kortvelesy, Stephan Liwicki, and Amanda Prorok. Reinforcement Learning with Fast and Forgetful Memory. In Thirty-seventh Conference Neural Information Processing Systems, 2023b. URL attps://openrevies wttp://wwww..eeeooee.eeeeettett.
Tianwei Ni, Michel Ma, Benjamin Eysenbach, and Pierre-Luc Bacon. When do transformers shine in ri? decoupling memory from credit assignment. Advances in Neural Information Processing Systems, 36, 2024.
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting Recurrent Neural Networks for Long Sequences. In Proceedings of the 40th International Conference on Machine Learning, ICMI.'23. JMIR.org. 2023. Place: Honolulu, Hawaii, USA.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Jiaju Lin, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song. Xiangru Tang, Bolun Wang, Johan S. Wind, Stanislaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Qinghua Zhou, Jian Zhu, and Rui-Jic Zhu. RWKV: Reinventing RNNs for the Transformer Era, December 2023, URL http: /octt::/arssee org. wttp://./.t.a.....///////// 200s.//305..0 130488 arXiv:2305.1 1044 [cs].
Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Emesstu, and Noah Dor- mann. Stable- Baselines3: Reliable Reinforcement Learning Implementations. Journal of Machine Learning Research, 22(268):1-83 2021. URL http:// org. papers/v22/20 1364. html.
12

---

Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay. arXiv preprint arXiv:1511.059522 2015.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear Transformers Are Secretly Fast Weight Programmers. In Proceedings of the 38th International Conference Machine Learning, pp. 9355- 9366. PMLR, July 2021. URL ISSN: 2640-3498.
mlr press/vi3e schasg/chhlaaaaaaaa html.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High- High-Dimensional Continuous Control Using Generalized Advantage Estimation. In Proceedings of the International Conference Learning Representations (ICLR), 2016.
Antonio Serrano-Muñoz, Dimitrios Chrysostomou, Simon Bøgh, and Nestor Arana-Arexolaleibba skrl: Modular and Flexible Library for Reinforcement Learning. Journal of Machine Learning Research, 24(254):1 -9, 2023. URL http: jmwrr///m/rr.////// www.org///a... http///pprr/////2/1111111 0112. html.
Jimmy TH Smith, Andrew Warrington, and Scott Linderman. Simplified State Space Layers for Sequence Modeling, In The Eleventh International Conference Learning Representations, 2022.
Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q- learning. In Proceedings of the AAAI conference on artificial intelligence, volume 30, 2016. Issue:
Ziyu Wang, Tom Schaul, Matteo Hessel, Hado Hasselt, Marc Lanctot, and Nando Freitas. Dueling network architectures for deep reinforcement learning. In International conference on machine learning. pp. 1995-2003. PMLR, 2016.
Paul Werbos. Backpropagation through time: what does and how do it. Proceedings of the IEEE, 18010): 1990. Publisher: IEEE.
Ronald J. Williams and Jing Peng. An Efficient Gradient- Algorithm for On-Line Training of Recurrent Network Trajectories. Neural Computation, 2(4):490-501, December 1990. ISSN 0899-7667. doi: 4.490.
URL https: https: org/oogg////oo 1112/1eco. 2.
Ronald J. Williams and David Zipser. Experimental Analysis of the Real-time Recurrent Learning Algorithm. Connection Science, 1(1):87-111, January 1989. ISSN 0954-0091. doi: 1080/ 09540098908115631. URL http:p/s//oa..ooo///.....0000000 Taylor & Francis _eprint: http:///do.uggg//....00000000009
Publisher:
13

---

A Return Comparison Between TBB and SBB
FFM
LRU
Linkttn
96
0.3
0.5
-0.2
Retur
-0.4
0.50
-0.5
-0.5
0.5
5000
10000
5002
10003
5000
10002
5000
10000
Rehum
0.5
0.4
0.4
0.2
0.0
-0.2
-0.2
6003
10000
6000
10003
5200
10000
10000
-0.4
-04
Ratum
0.0
0.8
-0.5
-0.5
-0.8
-0.8
-1.0
-1.0
10000
Robum
0.0
0.0
-0.5
-0.5
-0.5
20200
40000
20000
40000
20000
40000
20000
40000
1.0
1.0
1.0
Method
Ratum
TBB
0.5
0.5
0.5
0.5
L-VB-110
L281-20
0.0
0.0
0.0
0.0
L-100
5003
10000
5000
10003
5000
10000
5000
10000
1.0
1.0
Rah.
0.5
0.5
0.5
0.5
0.0
0.0
5003
10000
5000
10003
5000
10000
5000
10000
1.0
0.25
Return
0.5
0.5
0.00
-0.25
-0.5
-0.5
-0.50
2000
4000
2000
4003
2000
4000
2000
4000
0.0
1.0
-0.40
1.0
Ralbi
0.5
-0.45
0.5
-0.4
-0.5
0.50
-0.5
2000
4000
2000
4000
2000
4000
2000
4000
1.0
1.0
1.0
19h
Return
0.5
0.5
0.5
0.0
0.0
0.0
5002
10000
5000
10000
5000
10000
5000
10000
Train Epoch
Train Epoch
Train Epoch
Train Epoch
14

---

Figure 6: We run four memoroids nine different POPGym environments ten seeds and report the mean and 95% bootstrapped confidence interval. The POPGym environments have a minimum episodic return of -1.0 and a maximimum of 1.0. In virtually all experiments, Tape-Based Based Batching provides improved sample efficiency over all tested segments length using Segment-Based Batching. The Count Recall and Autoencoder environments have temporal dependencies that span the entire sequence, demonstrating the importance of TBB for long range dependencies. On the other hand, Positional Caripole has temporal dependency of two timesteps, and so policies trained via SBB can still do reasonably well. Like Count Recall, Repeat First has long term temporal dependencies, however, 1BB-trained methods do better than in Count Recall because Repeat First requires storing and recalling only a single observation. All policies perform poorly on bandits because we use deterministic policy, and this environment benefits from a stochastic policy.
2030
1400
1390
1991
Mothod
300
Tral: 3p:th
Tral Spoth
Tral: Tpoch
Figure 7: We examine three memoroids on Atari environments from the Arcade Learning Environment (ALE) Bellemare al. (2013), plotting the mean and 95% confidence interval over three random seeds. In all environments, we see that TBB outperforms SBB.
B Observation Sensitivity Analysis
Mode
Batch Mode TBB, Lsw
Norm Cunuletive 40%
1.0
Model
FFM
LinAttn
LRU
0.0
50
40
30
10
30
20
-10
Observation Age
Observation Age
Figure 8: We follow similar approach to Figure 3 for the Repeat First environment from the POPGym benchmark. In Repeat First, each the agent receives reward for outputting the initial observation each timestep. Here, would expect the initial observation, and zero elsewhere. Again, we see that the gradient is distributed throughout the episode, suggesting that VML will always span the episode. The Linear Transformer (LinAttn) with SBB does very poorly on this task, and so its gradient distribution is not informative.
virtually all of the probability mass
15

---

C
Segment-Based Batching
After collecting episodes E during a rollout, we split E into fragments such that each has a maximum length of L. Fragments
zero padded from the right until they precisely length L, turning them into segments a and padding masks m. The segments are stacked into dataset D, enabling easy batching, storage, and training (Figure 1). We define this approach more accurately in the following paragraphs.
We define segment as a length L sequence of transitions. During collection, episodes E longer than L transitions are splif into fragments F. Fragments are then zero-padded to be length L, resulting in fixed-size segments and associated masks The resulting segments and masks stacked into a dataset D. enabling easy batching, storage. and training (Figure 1).
The split function splits single episode E' into one or more fragments F, each of size except for the final fragment.
Fu
To.
Tı, TE+1s
TE-1 T2(1-1)
Fi
TE,
(16)
Fk
Tek:
Ta
The pad function zero pads a fragment F' into fixed size segment and associated mask m denoting the padding elements
o, m = pad(F,,F L) concat F, ( concat(1sun((r) (4--cun((F)
(17)
(18)
Using our split and pad operators, we split and pad each incoming episode, producing one or more segments and associated masks for each episode
du, mo
pad L), VF; E split(E, L).
(19)
Ok, mi
represent our training dataset DD as concatenation of segments and masks
dg.mo
Thyma o
OLimE 1K11-
D = concat
(20)
ajims
During training, we randomly sample rows from D for minibatching (on-policy) or experience replay (off-policy).
16

---

D The Discounted Return as a Memoroid
Theorem D.I. The discounted cumulative return given by
G
r
(21)
1=
is equivalent to compuling the following memoroid over ro, .........
= [(a,r) a E [0, 1), e ER)
(22)
er (1,0)
(23)
ra. (a', r') (aa',at ar +r)
(24)
f(o,b) b) = (7,0)
(25)
s((a, (o, (o,
(26)
Proof. We prove the correctness of our discounted return memoroid by showing the expansion is equivalent the discounted return.
(1,0) (7) (7, 0) (7,
(27)
(1, (n,ro) (7,r1) -11-7 7, To + m1) (72 + to mi
(28)
(1,0) 0) (7, (7) =
7..
(29)
(30)
T=U
If we let n - oo, we see that the second element in the monoid tuple approaches the discounted return
lim r
r
(31)
2m0 0
0
17

---

E The Generalized Advantage Estimate as a Memoroid
Let us define Generalized Advantage Estimation (GAE) in memoroid form: Theorem E.1. The GAE target given by
A
Σ
( & = (V (((s)
(32)
1=0
is equivalent to computing the following memoroid over 8,, S2+1,...
= ((a.g) a E [0, (
(33)
H (1,0)
(34)
(a. 9) (a', g') ) (aa', ag' + g)
(35)
f(o,b, b) = (7),0)
(36)
g g), b)) g.
(37)
Proof. We prove the correctness of our GAE memoroid by showing the expansion is equivalent to the GAE target. This proof is very similar the proof of the discounted return.
(1,0)
(7), 6e) (7 de + 0) (7)
(38)
(1, 0) (7),
(gd,o(+1) (1 7
(184+1) ((72)),8 716++1)
(39)
(1,0) (7 6,)
((4m)) = (1.7)) 7A.... ,,, 8 770++1 +... +...((2))6btt)))
(40)
(7A)", Σ (72))
(41)
1=0
we let n - co, we see that the second element the monoid tuple approaches the GAE target
lim
E(12)/b444
- (27)260 8444
(42)
2-+00
1=0
1=0
18

---

Resettable Monoid Transformation Proof
Proof of Theorem 4.2. First, let us compute all possible pairs of inputs, as we will use them to simplify the rest of the proof.
(A,0) (A', 0) (A- (1 0) + 0 A), = (A A',0)
(43)
(A, o (A', 0) -1- (4- A', 0) - (A 1)
(44)
(A, (A', (A-(1-1)))))) 11 H1 A', 1) = (Hr 1'
(45)
(46)
Now, we must demonstrate that associativity holds ((A (A', (A", (") (A, ((A',
(A", (")) for all possibilities of A, A', A" and b, D'.". That is, we must ensure that the episode boundaries correctly handled for all possibilities that information does not leak across episode boundaries and that prior information otherwise propagates forward in time.
(4 "
=((A- 1') 0) H2 AV.0V
=A A".
(47)
(A 1) (A".0) ((A A') - 0) + H, 0 1 (V))
= -A A' A", 1)
(48)
(H, 1', (A",
((Hr A') (1 H1, 0 A", 1V0)
(Hr A' A", 1).
(49)
And for 6" = 1, we have
(A A', 0)
(A A",1) (50)
(A A', 1) (A",1)
=((A -1
= =(H, A",1) (51)
(H, 1) 1A =((H (ar::)))-1-110 (1 1,
=(Hr A", (52)
We that resets correctly remove the impact of any terms that occur before D' = 1, while correctly propagating state when 6' = 0.
19

---

G Rewriting Sequence Models as memoroids
In this section, we reformulate existing models used in our experiments as memoroids. This reformu- lation necessary to use inline resets for these models.
G.1 Linear Transformer
The Linear Transformer from Katharopoulos et al. (2020) written as
X = 0 E 1
(53)
To = 0 E R'
(54)
X, Xn-1 1 + (WWoo))(Wonn)
(55)
Tn In-1 + H(W)
(56)
Xn (()))) + On I
Sn MLP
(57)
can be reformulated as the following memoroid
H = ((X,)) X E Rixk x E ER
(58)
es (0,0)
(59)
xx (X' x4) = (X + XX', XX+
(60)
[(o (W())))(W,D) .d(WE)))
(61)
xx, MLP x)(Wo))
x (WWo) to
(62)
where (0x) 1 + (LU(r).
G.2 Simplified State Space Models
Prior work (Lu al., 2024) defines associate scan operator for the S5 variant of State Space Models. Little work is required to rewrite this in memoroid form:
H {(X,x) X C
x
(63)
H1 (Jm,0)
(64)
(X,x)) (X', x²) =(X'X, X'x +'')
(65)
f(o,b) b) = -(Wx, Wg
(66)
(((,,)), b)) GeL,((W..)) + b)) o sigmoid(W, GeL) + b2)
(67)
where Wx, Wz, W. are learnable weights, b,, b2 are learnable biases, and Jm the square identity matrix of size m.
G.3 Linear Recurrent Unit
The Linear Recurrent Unit (Orvieto et al., 2023) could be roughly described as a theoretical simplifi- cation of S5, bringing it closer to classical RNNs. Writing it out as a memoroid, we see that it nearly identical to S5, however the weight initialization is different
H = {(X, xx) X Cm 2 E C
(68)
H, = (Im; 0)
(69)
(X,z) (X','' (X'z, X'x +x)
(70)
f(o = (Wx, W20)
(71)
g((X, x), (o,b)) MLP((a)
(72)
G.4 Fast and Forgetful Memory
Finally, we can rewrite Fast and Forgetful Memory (FFM) as memoroid, with the parallel scans simplifying its implementation and fixing numerical instabilities caused by large positive exponentials over long sequences, discussed in Morad al. (2023b). The original formulation is written as an
20

---

aggregator and cell. First, let us write down the I' term used in the aggregator.
r(t) exp (-t|a=) exp (-tiw)
(73)
exp -l( | + iwi) : exp - t( am + )wr)
exp-i(lox) + iwo)
(74)
exp =((1am + iwe)
We then write the aggregator
(E)= [(t- -j) O (ontsle))
T(1
Sk-
T(-1] ((o)
Skin
+
@
(75)
T(t
Sk-1
(Ei= l't - jj o (ox+s
where is the Hadamard product (or power), m the trace size, c the context size, and a I are trainable parameters representing decay and context respectively, Multiplying column vector by 1 "broadcasts" or repeats the column vector times. The cell is defined as
R
fkm & (ok:a) @ o (E2(2k)))
(76)
Skm Agg(#km S
(77)
2km Es (Flatten [R[Skm]] 3(Skml))
(78)
Mkm (Z O (C1(0):)) (5(0))) o (1 (la((o))))
(79)
Agg represents the aggregator (Equation 75) and represents linear layers with mappings & &
R + R", la R=x2 - Rd La.2z Rd + Rd. R, extract the real and imaginary components of complex number as reals, Flatten reshapes matrix (m c me) and || is the concatenation operator. LN is nonparametric layer norm, and a is sigmoid activation. We reformulate and simplify
r, the FFM aggregator, and the cell single memoroid
H ((X,0) X Cm EE
(80)
H1 (0, 0)
(81)
(X,1) - (X' (X o exp (t @ iw)) X ,t++t)
(82)
2 0 [Wio+
bz)
T
[(o,b)
,1
(83)
(Wi b.) @ o ( + b2)
(((x,t), (o. = MLP(IN(Wss (R(X) || ((X)) + ba)) R oWao ba) + (1 aW + ba)) Co o.
(84)
where W, b are learnable weights and biases, R, 3 extract the real and imaginary part of a complex number, O is the elementwise product, @ is an outer sum, and E R" w E Rm are learnable parameters. Note that the original FFM formulation requires distributing ((-t) into the sum. Since is learned, the real component can grow very large and cause numerical instabilities as overflows even a double precision float. This discussed in the limitations section of the original paper. Since our formulation utilizes Blelloch scan, we can do away with the negative exponent, removing the numerical instability. We note that unlike the other memory models we implemented, FFM is time-varying recurrence because the recurrent updates depends on t.
expé(la) + iw)
21

---

H Experiment Setup
The code necessary to reproduce all of our experiments is available at https:///iithn com/ proroklab aroroklabmmemory-ooo mhemorymmoooide We used the same model hyperparameters across all experiments. Training hyperparameters, such as number of epochs, varied across tasks. To find hyperparameters, we simply ran many experiments using SBB, the approach used prior literature. Once we arrived at good set of hyperparameters, we simply reused them for our TBB method.
H.1 Compute Used
We ran out of GPU credits early in the paper. We estimate roughly 70% of experiments were run on CPU only, across number of hardware configurations. Thus, it is not straightforward to arrive at single number. Users should be able to run at least one seed for each experiment we did, on reasonable laptop, approximately one week.
H.2
Model Setup
We construct our model using blocks. A block contains a linear layer with nonparametric layer normalization and leaky ReLU activation. Observations feed into block, followed by a memory model, followed by two more blocks. The hidden size all blocks 256 dimensions. For the S5 and LRU models, stacked two $5 and LRU layers, resulting in (256 per layer). The Linear Transformer and Fast and Forgetful Memory models use just a single layer with 256 dimensions of recurrent state. We use the ADAM optimizer without weight decay.
of 512 dimensions of recurrent state
H.3 Task Setup
For each task, we selected replay buffer large enough such that no old observations ever needed to be discarded. Epochs Rand, Train describes the number of episodes we collect randomly, and then the number of training epochs. Polyak r determines the target network update rate. Batch Size measures the batch size in transitions for each model update. I.R the learning rate with linear warmup over specified number of model updates. The ratio describes the number of episodes collected at each epoch, compared the number of model updates per epoch. 1 means we would perform gradient updates for each I episode collected. V Clip corresponds to gradient clipping, where the gradient magnitude is rescaled to at most Clip. y is the decay term used in MDPs. We use a linear learning rate warmup of 200 updates for all tasks.
Task
Epochs Rand, Train
Polyak r
Batch Size
LR
Ratio
V Clip
7
RepeatFirst
5,000, 5,000
0.995
1,000
0.0001
1:1
0.01
0.99
RepeatPrevious
5,000, 5,000
0.995
1,000
0.0001
1:1
0.01
0.5
CountRecall
10,000, 10,000
0.995
1,000
0.0001
0.99
PosOnlyCartPole
10,000, 10,000
0.995
1,000
0.0001
1:1
0.01
0.99
VelOnlyCartPole
10,000, 10.000
0.995
1,000
0.0001
1:1
0.01
0.99
NoisyCarPPol
10,000, 10,000
0.995
1,000
0.0001
1:1
0.01
0.99
AutoEncode
10,000, 10,000
0.995
1,000
0.0001
1:4
0.01
0.99
MultiarmedBandit
10,000, 10,000
0.995
1,000
0.0001
1:1
0.01
MineSweeper
10,000, 40,000
0.9975
1,000
0.0001
1:1
0.01
0.99
H.4 Wall-Clock Experiment Details
In the Figure 5 plot, standard approach of iterating over episodes in a batch. Both the monoid and standard approach are just-in-time compiled on GPU, however the standard approach requires for loop when the episode lengths are not fixed. We sample batch of episodes, where each episode length is sampled from discrete uniform distribution between one and maximum episode length. We find that our memoroid computes the discounted return between three orders of magnitude faster.
the wall-clock efficiency of our discounted return monoid against the
Next, we compare TBB and SBB scaling. TBB scales worse than SBB (O B) and Olog L) respectively, where B is the batch size and L is segment length). We question how this overhead translates wall-clock training time. the Figure table, we examine the total time spent training, finding that the time difference is negligible. The memory model forward pass is only fraction of the time spent each epoch, with environment sampling, replay buffer sampling (and in the case of SBB, splitting, truncating, and padding sequences) all taking nontrivial amount of time.
22

---

H.S Atari Experiment Details
We describe the model and training configuration for the Atari experiments below. We use a CNN similar to that of Mnih al. (2015), with filter sizes 8. 4, 3 and filter channels 32, 64, 64, and layemorm. The CNN is followed by the recurrent model with recurrent states of size 512, and two-layer MLP of width 512. We collect one episode per training epoch, and perform 5 gradient updates per epoch. We use a batch size of 16,000 transitions for each update, and evaluate our policy every 100 epochs.
23

---

I Non-Recurrent Q Learning
Algorithm 5 Non-recurrent Q learning update
Input: params 0, target params 6, function Q, train batch B, discount r3 maX.eA ( (s',.), a), vB[]
Target Q update Target update
0 mino [o a,) Usll, vB[
6 Φ 8)0
A Primer on Scans
In this section, we briefly review scans and associative scans. Generally speaking, we express classical RNNs using scans, and linear recurrent models using associative scans which tend to be more efficient.
J.1 Scans
scan is an operation over a sequence elements, often used in tensor processing. We write scans as some function defined over elements 21, 22, .
h 21 22, Jn
(85)
In deep learning, often formulate an RNN using scan. Consider, for example, the following simple recurrent network
hn (WWhn-1 Wz
(86)
yn = Wyh
(87)
where represents some nonlinearity, and the W terms learned parameters. We can define
hn-1 Zyn = oWinh-W Wz
(88)
Thus, execute scan hg 21 22 Fn compute the recurrent state hú and output Yn-
Note that in doing so, we must also compute all intermediate recurrent states h,...., hn- This is due the dependence of hn on hn-1. Let us write out the formula for h3 to demonstrate this dependence
h1 (W + Wzæ1)
(89)
h2 (W W2+1) + W222) h3 (W + W2 + W + W223).
(90)
(91)
Due to their sequential nature, standard scans tend be slow on GPU, since all computations must be executed in sequence.
J.2 Associative Scans
Certain operators may exhibit the associative property
(xt 33 (x2 x3))
(92)
When this is the case, we may use associative scans instead of scans. Associative scans (alternatively called parallel scans) are generally much faster to execute GPU than standard scans, While scan is (n) parallel time complexity, a work-cfficicnn parallel scan is (log n).
The key idea behind parallel scans is that the operator is associative, there is no explicit dependency that requires execute in series. Rather, we can we can parallelize computation. Consider the following expression
21 T2
33
Ft
(93)
If e exhibits the associative property. then we can compute the expression as
(x) Ø22
(I'
(94)
That is, we can compute the first term 22 (m1 222) independently of 24 (Z3 .TA). Then, we can compute the resulting operator h4
22
Z4.
This is a naiive associative scan executes the binary operator O(n logz ni) times. The Blelloch Scan (Blelloch, 1990) produces equivalent outputs the naiive associative scan, but does so in ((n) calls to The Blelloch Scan algorithm is relative complex implement or explain, but fortunately exists in the CUDA and JAX libraries.
24

---

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 highlight the claims in the abstract and contribution paragraph, which we provide evidence for throughout the paper.
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 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.
is fine include aspirational goals as motivation long 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:
We have a section on limitations and future work.
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 encouraged to create 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 few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, 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 impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
25

---

Justification: We provide proofs for the reset mechanisms and the discounted return and generalized advantage estimate memoroids in the appendix.
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 cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch 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
Does the paper fully disclose all the information needed to reproduce the main ex- perimental 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)?
Question:
Answer: [Yes]
Justification: We provide all hyperparameters the appendix. We also include all source code with yaml files denoting the configuration of each experiment we ran.
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 dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe 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 possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
26

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We include the code used to run all experiments in this paper. The supplemen- tary material upload has a readme describing the commands to install the library and run the experiments.
Guidelines:
The answer NA means that paper does not include experiments requiring code. Simply run 'python tape_ _epee_dnyyy PATH_ LULLTO_COOOOO CONFIG to run the tape-based RL experiments.
Please see the NeurIPS code and data submission guidelines (https://nipp cc/ mublic/guides http///SSuuiisiiiooooooo.ooo
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
cttpssccccbbllc/ wwtiiicgguieeeeeeee Sewieddeee/oeSuummmiiie ismiissioonnn ssionPolicy) 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. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We describe mostly everything, such as all hyperparameters and how we chose them, in the Appendix. Anything missing should be readily available in the code.
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 appreciate the results and make sense of them.
The full details can provided either with the code, 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 bootstrapped 95% confidence interval for all experiments except for the wall-clock time, which uses standard deviation instead.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
27

---

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 the standard error of the mean.
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 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 range (e.g. negative error rates).
If error bars 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
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We detail this in the appendix. We do not have hard number, but explain roughly how much compute would be required.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums 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: / http//nnurrrp.ppeeee wwtppccwcccccccccccccccccccccclt Etthcc///iiieeee dew.iieeeeeeeeeeeee
Answer: [Yes]
Justification: Yes, we have adhered to ethics guidelines.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws 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 research focuses improving efficiency in RL, as demonstrated on toy problems. It could potentially decrease energy usage in this manner, but we not want to jump any conclusions. Further research is required.
Guidelines:
28

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymentt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The worst thing our models can do is learn to play MineSweeper.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere 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 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: Yes, credit the various model authors for their memory models and the POPGym authors for their dataset.
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 URL.
The name of the license (c.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 souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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 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:
We just include methods, not new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
There were no 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 contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS 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:
There were no human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
30

---

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


Reviews:
[{"review_id": "dz9eG5YjV1", "reviewer": "Reviewer_KMv5", "summary": "The authors present a novel approach to recurrent reinforcement learning aimed at improving efficiency and performance. They introduce the concept of memory monoids, algebraic structures used to represent and manipulate the memory of RL agents with recurrent components. They also rewrite Simplified State Space Models (S5), Linear Recurrent Units (LRU), Fast and Forgetful Memory (FFM), and the Linear Transformer (LinAttn) as memoroids, demonstrating that memoroids can represent a large class of sequence models. While the method is slightly more computationally expensive compared to existing methods, it promises higher sample efficiency and improved performance.", "questions": "- The authors mention that \"We did not experiment on environments like Atari, primarily because it is unclear to what extent Atari tasks require long-term memory.\" I agree with this argument. I'd still be curious about games requiring counting or tracking objects across frames. For example, *Frostbite: The agent needs to keep track of ice floes and build igloos over time or Asteroids: Remembering the positions and trajectories of multiple asteroids can be crucial.*  Do the authors have any guesses for how their approach would be beneficial here?\n\n- \"TBB does not strictly require memoroids, but would likely be intractable for RNN or Transformer-based memory models.\" This is mentioned in the limitations by the authors. I'm curious about how they'd handle tasks that would require RNNs or Transformers.", "limitations": "I agree with the limitations listed by the authors.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is very well-written. It was a pleasure reading this work!\n- Clear experiments that demonstrate the superiority of tape-based batching to segment-based batching\n- Fig.3 discusses the consequences of truncating BPTT, which is quite interesting\n- Memorids can offer GPU efficient computation for existing recurrent variables like $\\lambda-$return. Its simple enough to reformulate them", "weaknesses": "- The experiments are carefully chosen to showcase certain features of memoroids. That said, they are fairly simple tasks which do not provide any insight about how well memoroids would perform on more challenging tasks\n- The authors might want to consider more complex envs like Memory MAze (https://github.com/jurgisp/memory-maze) to really showcase the strength of their method", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for reviewing our paper, we are happy to hear that it was a pleasure to read!\n\n> The experiments are carefully chosen to showcase certain features of memoroids. That said, they are fairly simple tasks which do not provide any insight about how well memoroids would perform on more challenging tasks\nThe authors might want to consider more complex envs like Memory MAze (https://github.com/jurgisp/memory-maze) to really showcase the strength of their method\n\n> The authors mention that \"We did not experiment on environments like Atari, primarily because it is unclear to what extent Atari tasks require long-term memory.\" I agree with this argument. I'd still be curious about games requiring counting or tracking objects across frames. For example, Frostbite: The agent needs to keep track of ice floes and build igloos over time or Asteroids: Remembering the positions and trajectories of multiple asteroids can be crucial. Do the authors have any guesses for how their approach would be beneficial here?\n\nIn our rebuttal PDF, we have added results for Atari Asteroids. We were surprised to find that TBB provides a significant improvement over SBB, even on tasks with weak partial observability. We plan to test Frostbite next.\n\nIf accepted, we will add a Memory Maze task as well.\n\n> \"TBB does not strictly require memoroids, but would likely be intractable for RNN or Transformer-based memory models.\" This is mentioned in the limitations by the authors. I'm curious about how they'd handle tasks that would require RNNs or Transformers.\n\nFor RNN or Transformers, one could just use more GPUs. Transformers in particular are strictly limited to shorter sequences due to their $O(n^2)$ space complexity. One solution could be to use strong regularization, so that we can significantly reduce the batch size/sequence length.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VZmMRQbc4G", "reviewer": "Reviewer_WNca", "summary": "This paper proposes a new interpretation on how to sample from a buffer of data to avoid the well known trade-offs of truncated BPTT. The final proposed method  interprets a recurrent networks as a monoid, and re-uses ideas developed for linear recurrent networks. This interpretation is time-invariant and enables the parallelization across timesteps better taking advantage of how current computer architectures work (i.e. GPUs). The new method for sampling from the experience replay is then compared to a method which segments the episodes into set truncated sequences (i.e. the sequences, termed \"segments\" here, are decided when data is put into the experience replay buffer). The comparison uses the POPGym benchmark.", "questions": "- Q-1. According to Algorithm 2, if the length of an episode is longer than the chosen batch size, wouldn't TBB also truncate the gradients, or not ever reach the end of the episode? Or is there some way I'm misinterpreting the algorithm?", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper addresses an important problem in partially observable domains, i.e. how do we deal with the limitations of truncated BPTT without incurring significant computational costs and variance costs of full BPTT.\n- Restricting the recurrent network to be a time-invariant kernel is an interesting constraint not yet explored in the context of reinforcement learning. This is novel, and the motivation for such a study is well received.\n- I think the overall approach is well founded, and doesn't have any technical flaws.", "weaknesses": "While I think the paper's novel contributions have interest for the RL community looking into partial observability, the language is overly dower towards previous works and the interpretation of prior methods is flawed. I will justify these claims below. I encourage the authors to reconsider their tone for subsequent submissions.\n\n\n- 1. **Segment Based Batching**\nThe assertion that reinforcement learning using deep recurrent networks (i.e. DRQN and following results) use what is termed as Segment-Based Batching (SBB) is fundamentally flawed. From (Hausknecht and Stone, 2015) to later works, \"segments\" are not pre-determined when putting episodes into the buffer. As stated in the **Stable Recurrent Updates** section of (Hausknecht and Stone, 2015) they worked on two approaches. The method they chose to follow randomly samples a starting point in the episode to generate a sequence, initializing the hidden state to zero. Later (Kapturowski et al, 2019) study this as well, albeit in the distributed RL setting. They also randomly sample starting episodes and studied whether to store hidden states from when the sample was gathered or allow for a \"burn-in\" period. I am unaware of which paper the SBB interpretation comes from as I understand it, and have not seen it in the literature.\n\nBecause of this interpretation, and not testing other replay sampling methods (i.e. those presented by (Kapturowski et al, 2019), it is difficult to say for certainty where TBB fits into the literature in terms of performance. I don't even care if TBB beats other approaches, but it should be compared with what the rest of the field actually uses.\n\n- 2. **Tone**\n\nThe issue with the tone of the **Background and Related Work** section continues, and wrongfully attributes the core assumption of truncated BPTT (i.e. $\\nabla \\approx \\nabla_\\sigma$) to reinforcement learning papers. This algorithm as far as I know was officially first discussed in (Williams and Peng, 1990), and the biased nature of the algorithm (i.e. due to truncating sequences) has been a consistent focus of study since then. While RL papers should be clearer about this assumption they are taking advantage of (as most recurrent architectures have before them), the burden for them to justify the usefullness of this algorithm for training recurrent networks at the time is not really their concern. It has only been recently, with the structured state spaces and linear recurrent networks lines of research, that there have been reasonable alternatives to truncated BPTT.\n\n\n\n\n[(Williams and Peng, 1990)]( https://direct.mit.edu/neco/article-abstract/2/4/490/5561/An-Efficient-Gradient-Based-Algorithm-for-On-Line)\n\n\n- 3. **Some missing Related Works**\n\nWhile coming from a different direction, there has been activity in removing the need for truncated BPTT--and sometimes BPTT entirely--in the RL setting. I'll only list a few papers which might be of interest:\n- Restricts the state of an RNN to be predictions and mostly removes the need for BPTT in simple domains. This paper also discusses the history of predictive representations as an alternative to traditional recurrent architectures extensively: https://www.jair.org/index.php/jair/article/view/12105. There are several works in that area to choose from.\n- From the angle of classical conditioning, this paper looks directly at the effects of truncation on prediction in a sequential decision making process: https://journals.sagepub.com/doi/full/10.1177/10597123221085039", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to read our paper and provide useful feedback. We are glad to hear that our work \"addresses an important problem\" and that our \"overall approach is well founded, and doesn't have any technical flaws.\" \n\nLet us address your concerns and questions below.\n\n### Weaknesses\n> Segment Based Batching The assertion that reinforcement learning using deep recurrent networks (i.e. DRQN and following results) use what is termed as Segment-Based Batching (SBB) is fundamentally flawed. From (Hausknecht and Stone, 2015) to later works, \"segments\" are not pre-determined when putting episodes into the buffer... Later (Kapturowski et al, 2019) study this as well, albeit in the distributed RL setting. They also randomly sample starting episodes and studied whether to store hidden states from when the sample was gathered or allow for a \"burn-in\" period. I am unaware of which paper the SBB interpretation comes from as I understand it, and have not seen it in the literature.\n\nSBB is equivalent to the approach from Kapturowski et al. with a burn-in length of zero and a zero-initialized recurrent state (Figure 1c in Kapturowski et al.). **Kapturowski et al. does not store full trajectories**. They **do not randomly sample a start point** within the sequence. From Kapturowski et al.:\n\n> Instead of regular (s, a, r, s0) transition tuples, we store fixed-length (m = 80) sequences of (s, a, r) in replay, with adjacent sequences overlapping each other by 40 time steps\n\nIn Kapturowski et al., Fig 1c they:\n1. Split a collected sequence into length 80 segments\n2. Place these segments in the buffer\n3. Randomly sample full segments from this buffer\n5. Initialize the RNN with a recurrent state of zeros\n6. Burn-in for 0, 20, or 40 timesteps, starting from the beginning of the segment\n7. Train on the remaining 80, 60, or 40 timesteps\n\nSee the R2D2 implementation from Google Research [1], where an \"unroll\" refers to the fixed 80-timestep sequence.\n\nIndeed, in practice many follow the SBB approach, forming $B \\times T$ subsequences so that subsequences can be stacked along the $B$ axis. Examples include\n\n- Seed RL [1]\n- SKRL [2]\n- TorchRL [3]\n- POMDP-Baselines [4]\n- R2L [5]\n\nFew libraries actually implement a nonzero burn-in length. Of the above listed libraries, only TorchRL and SeedRL implement burn-in.\n\n> it is difficult to say for certainty where TBB fits into the literature in terms of performance\n\nTBB is both simpler to implement and computes a strictly more accurate gradient than prior works using burn-in or stored recurrent states. The wall-clock time difference between SBB and TBB is imperceptible. If using a memoroid, we cannot see any reason to choose burn-in or stored-state approaches over TBB.\n\n> The issue with the tone of the Background and Related Work section continues, and wrongfully attributes the core assumption of truncated BPTT (i.e. $\\nabla \\approx \\nabla_\\sigma$ ) to reinforcement learning papers. This algorithm as far as I know was officially first discussed in (Williams and Peng, 1990), and the biased nature of the algorithm (i.e. due to truncating sequences) has been a consistent focus of study since then. \n\nThe shortcomings of truncation are not unique to RL, as you state. We have added the (Williams and Peng, 1990) citation and a few sentences explaining that this was common in supervised learning as well.\n\n> While RL papers should be clearer about this assumption they are taking advantage of (as most recurrent architectures have before them), the burden for them to justify the usefullness of this algorithm for training recurrent networks at the time is not really their concern. It has only been recently, with the structured state spaces and linear recurrent networks lines of research, that there have been reasonable alternatives to truncated BPTT.\n\nWe argue that authors are responsible for how they train their models. With classical RNNs, there is no good alternative to truncated BPTT, hence its widespread use. The authors made the approximations they needed for tractability -- that is perfectly fine. But are you suggesting we should not discuss the shortcomings of prior work? How are we to motivate our contributions without exploring the shortcomings of prior work?\n\n> Some missing Related Works\n> While coming from a different direction, there has been activity in removing the need for truncated BPTT–and sometimes BPTT entirely–in the RL setting. I'll only list a few papers which might be of interest:...\n\nThank you for the references, we have added both works to the \"Alternative to Segments\" subsection.\n\n### Questions\n> Q-1. According to Algorithm 2, if the length of an episode is longer than the chosen batch size, wouldn't TBB also truncate the gradients, or not ever reach the end of the episode? Or is there some way I'm misinterpreting the algorithm?\n\nYes, we guarantee at most one truncated episode per batch. The key point is that the user need not make the tradeoff between $B$ and $T$ dimensions. For comparison, we are currently using a batch size of 16,000 on Atari which fits in 20GB of GPU memory. We did not need to truncate any episodes during training.\n\n## References\n\n[1] https://github.com/google-research/seed_rl/blob/master/agents/r2d2/learner.py#L391-L392\n\n[2] https://github.com/Toni-SM/skrl/blob/636936f3ac49c6d2260bd130d72b789ca6dfe42b/skrl/agents/torch/td3/td3_rnn.py#L226\n\n[3] https://github.com/pytorch/rl/blob/da898261ba18cb221a0c3b54a679b692c3610f06/torchrl/objectives/value/utils.py#L214\n\n[4] https://github.com/twni2016/pomdp-baselines/blob/28cb5dda93d8182fbe3e3055b39f839b1274aa94/policies/models/recurrent_actor.py#L112\n\n[5] https://github.com/siekmanj/r2l/blob/247e1cbde3e5fdc2c1c72e31c52b3d08c9d68cf1/algos/off_policy.py#L77", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mgpmQ1823W", "reviewer": "Reviewer_M71j", "summary": "The authors identify a shared structure in the update rules of linear recurrent models, analogous to monoids. Leveraging this insight, they introduce a mathematical framework called memoroids, which unifies the recurrent update rules of these models. They also derive an inline resettable memoroid, which eliminates the need for segment-based batching in training recurrent reinforcement learning (RL) systems, while preventing information leakage across multiple episodes of varying lengths in long training sequences.\n\nThe paper demonstrates that segment-based batching fails to accurately compute the true gradient of multi-episode sequences during training, resulting in degraded algorithm performance. Using their resettable memoroid formalism, the authors propose tape-based batching, which simplifies and improves the accuracy of gradient calculation for recurrent loss functions. They recast four linear recurrent models in their framework and show that this approach, combined with tape-based batching, leads to better convergence and improved sample efficiency in POPGym environments that require memory for problem-solving.", "questions": "1. Did the authors ever consider the case where the training segments in segment-based batching is a full-episode rollout? What are the effects of doing this on performance?\n2. The authors show that the wallclock time it takes to train a tape-based batching system is the same wallclock time it takes to train a segment-based batching system with varying zero-padded segment lengths. Why do you think this is the case? \n3. Could the authors please elaborate on the ‘What are the Consequences of Truncating BPTT’ experiment?", "limitations": "Yes, the authors mention the limitations of their method.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The introduction of the memoroid formalism, the extension to inline resettable memoroids, and the proposal of tape-based batching are all novel contributions.\n2. The formalism is clearly explained and the first sections of the paper are clear to follow.\nSegment-based batching, the default for training recurrent RL policies, where a sequence of multi-episode experience is zero-padded and split into smaller segments for training, is clearly introduced. The way that tape-based batching simplifies the loss calculations for recurrent systems is also well-defined. \n3. The benefits of recasting four existing linear recurrent models within the memoroid formalism and using tape-based batching are empirically shown by the improvement in sample efficiency and converged episode returns in the experimental section.", "weaknesses": "1. The sample efficiency and wall clock efficiency experiments are easy to follow and clearly demonstrate the advantages of the formalism. However, the experiment demonstrating that segment-based batching leads to poor recurrent value estimators is unclear. In general, I feel this experiment subtracts from the flow of the paper. The conclusions drawn from this experiment are also not clearly explained. \n2. Only benchmarking on one environment is quite limiting, although I understand that POPGym is specifically made to test memory. Is there another environment in which you expect your method to lead to improved performance due to being able to train on very long sequences with more accurate gradients? Perhaps some tasks in the Atari suite? If so, could you please run some experiments to illustrate this?\n3. The paper only shows how off-policy algorithms may be improved and focuses only on discrete action space environments. It is mentioned that other work considers on-policy methods. It would be good to include one set of on-policy experiments as well.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for spending time to read and critique our paper. We are happy to hear that you consider our contributions novel. Below, we will respond to your questions and concerns.\n\n### Weaknesses\n> ...the experiment demonstrating that segment-based batching leads to poor recurrent value estimators is unclear. In general, I feel this experiment subtracts from the flow of the paper. The conclusions drawn from this experiment are also not clearly explained.\n\nWe have written a more intuitive tutorial under the \"Questions\" heading at the bottom of our response. Please let us know if this is helpful, and if so, which parts you think would aid in understanding. We will integrate these parts into the text.\n\n> Only benchmarking on one environment is quite limiting...Is there another environment in which you expect your method to lead to improved performance due to being able to train on very long sequences with more accurate gradients? Perhaps some tasks in the Atari suite?\n\nPOPGym is not one environment but a collection of environments. Each one has entirely different observation spaces, action spaces, and transition dynamics which makes them quite diverse. We note that we ran over 1800 experiments.\n\nWe have added results from Atari Asteroids to the rebuttal PDF. We find that TBB still shows a noticeable benefit over SBB. We plan to add one more Atari environment as well. We do not have enough compute to run the entire Atari suite.\n\n> The paper only shows how off-policy algorithms may be improved and focuses only on discrete action space environments...It would be good to include one set of on-policy experiments as well.\n\nAlgorithms 1,2 work on-policy as well, see line 4 of Algorithm 1. Frankly, many other papers just focus on a single algorithm. For example, the inspiration for our work [1] only evaluates PPO. We believe we have already made significant contributions, and while more experiments are always beneficial, focusing on off-policy algorithms is not an inherent weakness.\n\n### Questions\n> Did the authors ever consider the case where the training segments in segment-based batching is a full-episode rollout? What are the effects of doing this on performance?\n\nYes, see lines 265-269, quoted below\n\n> Even for large segments lengths L = 100, we find a significant gap between SBB and TBB. SBB must make an inherent tradeoff – it can use long segments to improve gradient estimates at the cost of smaller effective batch sizes, or shorter segments to improve the effective batch size at the expense of a worse gradient estimate. TBB does not need to make this tradeoff.\n\ncertain tasks like `CountRecall` have a maximum episode length less than $100$ timesteps, and we evaluated SBB with $L = 100$. \n\nIf both\n1. All episodes are a fixed length\n2. The SBB truncation length $L$ is set to precisely this length\n\nthen SBB and TBB produce identical outputs.\n\n> The authors show that the wallclock time it takes to train a tape-based batching system is the same wallclock time it takes to train a segment-based batching system with varying zero-padded segment lengths. Why do you think this is the case?\n\nSBB applies additional split-and-pad operations to the trajectories, which are quite expensive. Each subsequence contains a varying number of transitions, which corresponds to a varying amount of padding we need to add. Variable-size operations are generally slow and difficult to batch efficiently.\n\nFurthermore, SBB has $\\log L$ complexity while TBB has $\\log B$. Generally speaking, $B$ is usually only one or two orders of magnitude larger than $L$, which in log space is not very significant. \n\n> Could the authors please elaborate on the ‘What are the Consequences of Truncating BPTT’ experiment?\n\nYes, it is a bit complex, but hopefully will clarify an important finding.\n\nConsider a simple recurrent network with some weights $W$ that update recurrent state $h$\n\n$$ h_{t} = h_{t-1} + W x_t$$\n\nWe compute the error at $h_t$, which backpropagates through $W$ at each timestep\n\n$$ E = \\textrm{error}(h_t, W x_1) +  \\textrm{error}(h_t, W x_2), \\dots $$\n\nNow, what happens when we truncate BPTT, with a truncation length of $L$? Consider the following:\n\n$$ h_{1} = h_{0} + W x_1 $$\n\n$$ \\vdots $$\n\n$$ h_{L + 1} = h_{L} + W x_{L+1} $$\n\nNow, let us compute the error again\n\n$$ E = \\textrm{error}(h_{L+1}, W x_2) +  \\textrm{error}(h_{L+1}, W x_3), \\dots $$\n\nBy truncating BPTT, we do not consider the $\\textrm{error}(h_{L+1}, W x_1)$ term. $W x_1$ is influencing $h_{L+1}$ and we cannot do anything about it!\n\nIn Figure 3, we plot how much $x_t, x_{t-1}, \\dots, x_1$ affects the Q value $q_t$. In this task, at each timestep $t$, our policy needs to output $x_{t-9}$. In this case, it makes sense to set $L = 10$ right? If $L = 10$, we will backpropagate through all the necessary observations $x_t, x_{t-1}, \\dots x_{t-9}$. \n\nSince we cannot change how $x_{t-10}, x_{t-11}, \\dots$ impact $q_t$ (and given that $x_{t-10}, x_{t-11}, \\dots$ are not necessary to remember $x_{t-9}$) we would hope that the contribution of $x_{t-10}, x_{t-11}, \\dots$ to $q_t$ is zero. That no matter what $x_{t-10}, x_{t-11}, \\dots$ are, $q_t$ should not change.\n\n**We find that this is not the case at all**. Even if we use TBB which effectively sets $L=\\infty$, the $x_{t-10}, x_{t-11}, \\dots$ terms all affect $q_t$. This is a huge issue for SBB and truncated BPTT because we cannot change how $x_{t-10}, x_{t-11}, \\dots$ contribute to the recurrent state, while simultaneously showing they have a big contribution to $q_t$. Figure 3 shows that even changing $x_{t-80}$ has a noticable impact on $q_t$. At least with TBB, we can in theory learn to not use the $x_{t-10}, x_{t-11}, \\dots$ terms. With SBB, we cannot learn to ignore these terms.\n\nWe think that this could be a major reason why training recurrent policies in RL is so difficult.\n\n## References\n[1] https://openreview.net/forum?id=4W9FVg1j6I", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lmm2YIImji", "reviewer": "Reviewer_hZNa", "summary": "This paper introduces the concept of *memoroids*, formalizing a method for computing reset-able episodes with recurrent models that use an associative operation for the latent state update. The authors argue that memoroids can simplify the computation of the loss function and further forego the error in gradient computation introduced by truncation. They compare the method to state-of-the-art segment-based batching and find that it leads to improved performance across a range o sequence models, on RL tasks requiring memory.", "questions": "- The tradeoff between L and B for SBB is nowhere given formally (did i miss it)?\n\n- Shouldn't TBB and SBB with L larger than the max episode length be comparable in performance since the only difference is the zero padding?\n\n- Figure 4: It seems like the plots are truncated - not showing the final performance of all methods? I was hoping to, at least in the Appendix, find the full plots.", "limitations": "The authors are upfront about limitations concerning an increased time cost of the approach compared to segment-based batching.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "A intuitive and simple, yet powerful generalization of a range of different computations is presented. The limitations of segment-based batching are introduced in a clear fashion. Investigating the effects of the truncation in BPTT is an interesting and worthwhile endeavor!", "weaknesses": "The introduction of the method (TBB) itself is not very clear and somehow confusing. Parallel scans, the sole reason for the frameworks efficiency, are only mentioned. A thorough introduction would be appropriate here. The experimental evaluation should be improved. For example, when investigating the effects of truncation in SBB, investigating more than one environment, and also comparing a range of different truncation horizons would be interesting. It seems like the plots demonstrating the sample efficiency of TBB are truncated - not showing the final performance of all methods. Finally, the explanation of the sample efficiency experiment is missing some crucial details: it seems like L and B are chosen in a way to ensure a fair comparison between TBB and SBB, how was it done?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to read our paper and provide feedback. We are happy to hear that we present an \"intuitive and simple, yet powerful generalization\", and that our research is an \"interesting and worthwhile endeavor\". \n\nWe think your reject rating is quite harsh, given your comments. Hopefully we can clarify some points and reach a shared understanding. \n\nFirst, we would like to bring up two theoretical contributions that were not listed in your review, which we think could be useful even if you are not sold on TBB!\n\n- We derive memoroids for the discounted return and advantage, achieving $O(\\log n)$ complexity compared to standard $O(n)$ implementations. As far as we know, we are the first to model these quantites using a parallel scan.\n- We introduce a method for inline resets, enabling any memoroid to efficiently span multiple episodes\n\nNow, let us address your questions and concerns.\n\n### Weaknesses\n> The introduction of the method (TBB) itself is not very clear and somehow confusing.\n\nCould you elaborate? We are certainly willing to improve the text if you can help us pinpoint which passages are unclear. Would it be helpful to move Algorithms 1, 2, 3 before their descriptions in the text?\n\n> Parallel scans, the sole reason for the frameworks efficiency, are only mentioned. A thorough introduction would be appropriate here.\n\nWe provide a summary of parallel scans on lines 138-142, highlighting [1,2] for further reading. \n\nWe already devoted almost half the paper (4 / 9 pages) to background and related work, so we cannot fit an introduction to scans in the main paper. We have added a primer on parallel scans to the Appendix.\n\n> The experimental evaluation should be improved. For example, when investigating the effects of truncation in SBB, investigating more than one environment, and also comparing a range of different truncation horizons would be interesting.\n\nFigure 6 in the Appendix compares TBB to four different truncation horizons, across four models, nine environments, and ten random seeds for a grand total of 1800 experiments.\n\nIf you are referring to the VML/RML plots, we also investigate another environment in Figure 7 in the Appendix. The nature of these experiments limits them to environments with a fixed episode length and known RML, which are rare. That said, we believe Figure 6 provides sufficient evidence that TBB is beneficial.\n\nAs for different truncation horizons of the RML/VML experiments, we show that even with TBB (infinite horizon), we experience similar behavior to SBB. This implies that varying SBB with larger $L$ would not change much.\n\n> Finally, the explanation of the sample efficiency experiment is missing some crucial details: it seems like L and B are chosen in a way to ensure a fair comparison between TBB and SBB, how was it done?\n\nWe feed TBB and SBB the same number of transitions for each update. We have updated the text to clarify this.\n\nFor SBB with large $L$, a portion of these transitions will be zero padding. This comparison makes sense given our paper's focus on efficiency, because:\n\n> the cost of zero padding in SBB is equivalent to the cost of real data – it takes up equivalent space in the replay buffer and takes just as much compute to process as real data (lines 260-262)\n\n### Questions\n> The tradeoff between L and B for SBB is nowhere given formally (did i miss it)?\n\nWe address this in lines 257-271. We will quote some of this in answering your next question.\n\n> Shouldn't TBB and SBB with L larger than the max episode length be comparable in performance since the only difference is the zero padding?\n\nNot exactly, here is an excerpt from lines 264-271:\n\n> Even for large segments lengths $L = 100$ we find a significant gap between SBB and TBB. SBB must make an inherent tradeoff – it can use long segments to improve gradient estimates at the cost of smaller effective batch sizes, or shorter segments to improve the effective batch size at the expense of a worse gradient estimate. TBB does not need to make this tradeoff.\n\n> It seems like the plots demonstrating the sample efficiency of TBB are truncated - not showing the final performance of all methods.\n\nSorry, what do you mean by \"final performance\"? Are you asking for us to train to convergence? If so, this is not standard in RL. Even the papers that introduce the SAC [3] and PPO [4] algorithms do not train to convergence.\n\nThe min/max return in POPGym environments is bounded in $[-1, 1]$, so TBB is already solving many of the tasks in Figure 6. Exceptions include `CountRecall` and `Autoencode`. We have rerun `CountRecall` experiments for longer, and plotted the results in the rebuttal PDF. You can see that two TBB models converge to the optimal return. We are also rerunning the `Autoencode` experiments for longer.\n\nWe do not truncate any plots -- we plot the full training duration for all experiments. We trained all tasks for either 5k, 10k, or 40k epochs.\n\n### References\n[1] Hinze. _An algebra of scans._\n\n[2] Blelloch. _Prefix Sums and Their Applications._\n\n[3] Haarnoja et al. _Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor._\n\n[4] Schulman et al. _Proximal Policy Optimization Algorithms._", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the AC and all the reviewers for taking time to read our paper and provide useful feedback. In general, the reviewers consider our contributions beneficial:

`KMv5` writes
> The paper is very well-written. It was a pleasure reading this work!

> Clear experiments that demonstrate the superiority of tape-based batching

`WNca` writes
> The paper addresses an important problem

> This is novel, and the motivation for such a study is well received

> the overall approach is well founded, and doesn't have any technical flaws

`M71j` writes
> The introduction of the memoroid formalism, the extension to inline resettable memoroids, and the proposal of tape-based batching are all novel contributions

> The formalism is clearly explained and the first sections of the paper are clear to follow

> The benefits...are empirically shown by the improvement in sample efficiency and converged episode returns

`hZNa` writes
> A intuitive and simple, yet powerful generalization of a range of different computations is presented

> Investigating the effects of the truncation in BPTT is an interesting and worthwhile endeavor!

We have attempted to address each reviewer's concerns individually. In particular, we have added an Atari pixel task and POPGym task to the rebuttal PDF for reviewers `KMv5`, `M71j`, and `hZNa`. This brings the total number of experiments to over 2,000.

We have also agreed to add an appendix section introducing parallel scans for reviewer `hZNa`, and additional citations for reviewer `WNca`.

Meta-review: Dealing with partial observability is one of the most importance challenges in reinforcement learning. This paper proposes a memory-based architecture that makes memory computation efficient while simplifying the loss function for partially observable environments. The proposed idea can be combined with advanced sequence models such as State Space Model. The empirical result on POPGym, a benchmark for memory capability, shows that the proposed method outperforms the widely-used truncated backpropagation through time (BPTT) approach. 
Some reviewers initially raised concerns regarding clarity, as they were initially confused or even misunderstood the paper. However, most of the questions were clarified during the rebuttal period, and all of the reviewers ended up leaning towards acceptance. The reviewers acknowledged that the proposed method is novel and promising as a new approach to deal with partial observability in RL. Thus, I recommend to accept the paper. I suggest the authors to revise the paper to make it easier to read for the camera-ready version.
Decision: Accept (poster)