Recurrent Reinforcement Learning with Memoroids

Authors: steven  morad
chris  lu
ryan  kortvelesy
stephan  liwicki
jakob  foerster
amanda  prorok
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-12 14:01:15

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.

Main paper: 
 introduction: 
 Reinforcement learning (RL) traditionally focuses on solving Markov Decision Processes (MDPs), although for many interesting problems the Markov state is hidden. Instead, we receive noisy or ambiguous observations, resulting in Partially Observable MDPs. The standard approach to RL under partial observability involves summarizing a sequence of observations into a latent Markov state using a memory model or sequence model. Commonly used models include RNNs and Transformers.Training Transformers or RNNs over 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 do 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 a large class of efficient memory models. Armed with the memoroid, we propose a new batching method that eliminates the need for segments. Figure 1 : We visualize the Segment-Based Batching approach often used in prior literature. A worker collects a rollout of episodes, denoted by color. Each episode is split and zero-padded to produce a 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.In particular, we• Derive memoroids for existing sequence models, as well as the discounted return and advantage • Introduce a method for inline resets, enabling any memoroid to efficiently process multiple episodes • Demonstrate that using segments degrades recurrent value functions • Propose a new 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, T , γ), where at each timestep t, an agent produces a transition T = (s, a, r, s ′ ) from interaction with the environment. We let s, s ′ ∈ S denote the current and next states and state space, a ∈ A denote the action and action space, R : S × A × S → R denote the reward function, and T : S × A → ∆S denote the state transition matrix (∆ denotes a distribution).In RL, our goal is to learn a policy parameterized by θ that maps states to action distributions π θ : S → ∆A. The agent samples an action from the policy given the current state a ∼ π θ (s), and stochastically transitions to the next state s ′ ∼ T (s, a), receiving a reward r = R(s, a, s ′ ). The optimization objective is to find the parameters θ that maximize the expected return, discounted by γ:E π [ ∞ t=0 γ t R(s t , a t , s t+1 )].It is often practical to model terminal states in MDPs, such as a game over screen in a 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 zero. We mark whether a state is terminal using the done flag d ∈ {0, 1}. The done flag is stored in the transition T = (s, a, r, s ′ , 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 -just the state. We receive the done flag, reward, and next state (r, d, s ′ ) at the next timestep. This distinction between current and next timestep becomes important when we execute memoroids over multiple episodes.We find that our paper is more clear if we introduce a begin flag b ∈ {0, 1} that is emitted alongside each observation, available during both training and rollouts. The begin flag is 1 at the initial timestep of an episode and 0 otherwise. We differentiate between a transition T = (s, a, r, s ′ , b, d) available only in hindsight, and a partial transition T = (s, b) as emitted while navigating the MDP. To reiterate, we can access T at any time, but we can only access T during training.In partially observable settings, we cannot directly measure the Markov state s. Instead, we indirectly measure s via the observation o ∼ O(s), following the observation function O : S → ∆O. With the observation replacing the state, interaction with the environment now produces a transition P = (o, a, r, o ′ , d, b) 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 b = 1 and continuing until a terminal state is known as an episode E. We use a memory model M to summarize the corresponding sequence of partial transitions into a latent Markov state.EQUATION)If M is recurrent, we may alternatively write M as a single update or batched update respectivelyEQUATION)where H is the set of recurrent states. 
background and related work: 
 In deep learning, we often wish to train memory models over batches of sequences. For a single sequence, we can use 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 et al., 2019; Hafner et al., 2023; Bauer et al., 2023; Liang et al., 2018; Raffin et al., 2021; Huang et al., 2021; Serrano-Muñoz et al., 2023; Lu et al., 2022; Ni et al., 2024) .In Segment-Based Batching (SBB), we split and zero pad episodes so that they can be stacked into a tensor with batch and sequence length dimensions B × L. Each row in this tensor is a segment σ containing exactly L transitions. Episodes longer than L transitions will be split into multiple fragments, such that each is at most L transitions. Fragments shorter than L transitions will be zero padded from the right, such that they become exactly length L. We call these padded length L fragments segments. We must also store a mask m denoting which elements are zero-padding and which are data. The segments and masks are stacked along the batch dimension, creating B × L matrices for storage and training (Figure 1 ). We formally define SBB in Appendix C.The Shortcomings of Segments SBB introduces a 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 on zeros that are discarded during gradient computation.(3) The zero padding also prevents the use of BatchNorm (Ioffe & 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, we approximate the true BPTT gradient with the T-BPTT gradient.Let us demonstrate the effect of SBB on the memory model gradient. Assume we have a loss function L defined over a model parameterized by θ. We define the true gradient of the loss over an episode of length n as ∇L. In SBB, we split an episode into length L segments. We approximate the gradient over these segments asEQUATIONUnder 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 L can be learned. In fact, our experiments show that ∇ σ is often a poor approximation of ∇.We are not the first to realize the drawbacks of SBB. Hausknecht & Stone (2015) store recurrent states in 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 4 ) time complexity and is thus much slower (Williams & Zipser, 1989) . Irie et al. (2024) propose a 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 a large majority of efficient memory 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 to 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.A recent class of models espouse time-parallel execution while being either linear or 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 et al., 2023; Martin & Cundy, 2018; Morad et al., 2023b) are sometimes called Linear Recurrent Models because they usually (but not always) employ a Linear Time-Invariant (LTI) recurrent state update, which can be computed in parallel over the time axis (Gu & Dao, 2023) .Monoids Prior work on efficient sequence modeling primarily updates the recurrent state using linear functions (Schlag et al., 2021; Gu et al., 2021; Smith et al., 2022; Orvieto 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 are linear, as 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, •, e I ) is a monoid if:(a • b) = c a, b, c ∈ H The binary operator • is closed on H (4) (a • b) • c = a • (b • c) a, b, c ∈ H The binary operator • is associative (5) (e I • a) = (a • e I ) = aa, e I ∈ H There exists an identity element e Iwhere • for a single input a is defined as( • a) = (e I • 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 a sequence of length n, a work-efficient parallel scan known as the Blelloch Scan executes O(n) calls to • in O(n) space to produce n outputs (Blelloch, 1990) . With p parallel processors, the parallel time complexity of the scan is O(n/p + log p). For large GPUs where n = p, the parallel time complexity becomes O(log n). 
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 is an image, the recurrent state H is a latent map representation, and the Markov state S is a set of x, 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. EQUATIONGiven n inputs, functions f and g can each be split into n concurrent threads. Recall that monoids have O(log n) parallel time and O(n) space complexity. Consequently, all memoroids have O(log n) parallel time complexity and O(n) space complexity on the length of the sequence1 .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 et al., 2024; Orvieto et al., 2023; Morad et al., 2023b) as memoroids in 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 Blelloch scan.Accelerated Discounted Returns Memoroids can model other recurrences as 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 as memoroids enables us to compute them in a GPU-efficient fashion, using a high-level framework like JAX (Bradbury et al., 2018) . We find that we can compute such quantities orders of magnitude more quickly than the standard approach. We provide the definitions and proofs of these formulations in Theorem D.1 and Theorem E.1.Inline Recurrent State Resets So far, we have assumed that we operate memoroids over a single episode using the Blelloch Scan. To scan over a batch of variable-length episodes, we could truncate and zero pad sequences such that each is a fixed length (i.e., SBB). However, this introduces the issues explained in 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 a new monoid that respects episode boundaries. Theorem 4.2. All monoids (H, •, e I ) can be transformed into a resettable monoid (G, •, g I ) defined asEQUATIONg I = (e I , 0) (11) (A, b) • (A ′ , b ′ ) = ((A • (1 -b ′ ) + e I • b ′ ) • A ′ , b ∨ b ′ ) (12)For a single episode, the A term output by the operator • is equivalent to the output of •. Over multiple contiguous episodes, • prevents information flow across episode boundaries.Proof. See Appendix F.0 0 0 0 1 1 2 3 3 0 4 6 7 I E 1 E 2 Episode Index D Figure 2:A visualization of sampling in TBB, with a batch size of B = 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 I. For example, we might sample 4 from I, corresponding to E 1 in red. Next, we sample 7 from I, corresponding to E 2 in red. We concatenate B = concat(E 1 , E 2 ) and return the result as a train batch.By transforming the monoid (H, •, e I ) within a memoroid, we no longer require separate time and batch dimensions during training. Now, memoroids can process a long sequence comprised of many distinct episodes. Unlike Blelloch (1990); Lu et al. (2024) which provide a reset operator for a specific model, our resettable transformation works for any monoid.Training over the concatenation of episodes would be intractable using Transformers or RNNs due to poor sequence length scaling, while Linear Recurrent Models leak information between episodes. By combining the efficiency of memoroids with our resettable transform, we resolve these issues, enabling us to fold the batch and time dimensions into a single dimension. We call this approach Tape-Based Batching (TBB), which consists of two operations: insertion and sampling. We provide operations for both on-policy and off-policy RL algorithms. Furthermore, we design TBB in a manner that greatly simplifies the implementation of recurrent loss functions.Insertion In RL, we must store the training data (transitions) that we collect during a 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, we discard old transitions by popping the episode begin index from the left of I, and discarding the resulting episode in D. This is guaranteed to discard the oldest episode in D.This method works both for rollouts that contain complete episodes (d n-1 = 1), and those that contain incomplete episodes (d n-1 ̸ = 1), where a rollout might stop before finishing an episode. When combining incomplete episodes with multiple rollout workers, we can experience race conditions. In this scenario, it is easiest to keep one D, I per worker to prevent race conditions.Sampling Once we have constructed D, I, we are ready to train a policy or compute returns. We sample transitions from D, I following Algorithm 2. If we are training on-policy, we can simply train on D. If we are training off-policy, we randomly sample a training batch B from our dataset by slicing D using randomly-sampled sequential pairs of episode indices from I. One could extend our sampling approach to implement Prioritized Experience Replay (Schaul et al., 2015) by assigning each episode or index in I a priority.With TBB, we utilize unmodified, non-recurrent loss functions to train recurrent policies, reducing the implementation complexity of recurrent RL algorithms. Unlike SBB, there is no need to mask outputs or handle additional time dimensions like in SBB. With TBB, the only difference between a recurrent and nonrecurrent update is precomputing the Markov states s, 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 in Algorithm 4. Note that the SBB update equation has an additional time dimension k and requires a padding mask m i,j . . . .I[i + 1]]) ▷ Append episode to batch end while B = B[: B] ▷ Make batch exactly B transitions Algorithm 3 TBB deep Q update Input: params θ, target params ϕ, Q function Q, sequence model M , train batch B, discount γ,M θ ((P B,1 ) ... (P B,L ))   ▷ Estimate Markov state   s ′ 1,1 , ... s ′ 1,L . . . s ′ B,1 , ... s ′ B,L   ←   M ϕ ((P 1,1 ) ... (P 1,L )). . .M ϕ ((P B,1 ) ... (P B,L ))   ▷ Estimate next Markov state ŷj,k = (r j,k + maxa∈A γQ ϕ (s ′ j,k , a)), ∀B[j, k] ▷ Compute target with extra time dimension k θ ← min θ m j,k • ∥Q θ (s j,k , a j,k ) -ŷj,k ∥, ∀B[j, k] ▷ Compute loss and update params ϕ ← ϕβ + (1 -β)θ ▷ Update target network params 
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 et 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 our policies using Double Dueling Figure 3 : We demonstrate that SBB can hurt Q learning through truncated BPTT. We examine the Repeat Previous task, with RML = 10, comparing SBB (left) to TBB (right). For SBB, we set L = RML = 10 to capture all necessary information. After training, we plot the cumulative partial derivative with respect to the observations on 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, a majority of the Q value is not learnable when using SBB. Even when we set L = ∞ 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 et al., 2016; Wang et 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 a 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 k 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.EQUATIONWe 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.e., Q(s n , a n ) = R(s n , a n , s n+1 )). 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 asEQUATIONand 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 to span the entire episode. This means that the recurrent models are 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 on our Q values, yet we are unable to forget or 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 efficiency: (1) As previously discovered, the truncated gradient used by SBB is often a 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 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. We report some combinations of model and task in Figure 4 and present the full results in Appendix A. Figure 4 : We compare TBB (ours) to 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 a noticeable improvement in sample efficiency over SBB, across nearly all configurations of memory model and environment. 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. In our experiments, SBB with larger L always outperforms shorter L, suggesting that the gradient estimation error is a larger contributor to 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 to finish. For SBB, the parallel time complexity is O(log L) while TBB has O(log B) complexity where B > L, but in practice there is no perceivable difference in wall-clock time. One possible reason for this discrepancy is that SBB applies expensive split-and-pad operations to the trajectories, while TBB does not. Each fragment 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. Figure 5 : (Left) We compare how long it takes to compute the discounted return using our memoroid, compared to the standard way of iterating through a batch. Computing the discounted return is orders of magnitude faster when using our 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 to our sensitivity analysis, old observations unexpectedly impacted the Q value across models and tasks. Moving forward, we suggest testing newly-designed memory models to see whether VML = RML, to 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 of 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 L cost. There was no perceptible difference in our experiments, but very long sequences such as those used for in-context RL could incur more noticeable training costs. TBB does not strictly require memoroids, but would likely be intractable for RNN or Transformer-based memory models. 
conclusion: 
 We introduced memoroids as a unifying framework for efficient sequence modeling. We found that memoroids can represent a large number of efficient recurrent models, as 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.A Return Comparison Between TBB and SBB Figure 6 : We run four memoroids on nine different POPGym environments over 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 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 Cartpole has a 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, SBB-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 a deterministic policy, and this environment benefits from a stochastic policy. 2013), plotting the mean and 95% confidence interval over three random seeds. In all environments, we see that TBB outperforms SBB. Figure 8 : We follow a 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 at each timestep. Here, we would expect to see virtually all of the probability mass on 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. 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: Yes, we have adhered to ethics guidelines.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
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 on improving efficiency in RL, as demonstrated on toy problems. It could potentially decrease energy usage in this manner, but we do not want to jump to any conclusions. Further research is required.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The 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 a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: Yes, we 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 a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)? 
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? 
 
 Appendix: After collecting episodes E during a rollout, we split E into fragments F such that each F has a maximum length of L. Fragments are zero padded from the right until they are precisely length L, turning them into segments σ and padding masks m. The segments are stacked into a dataset D, enabling easy batching, storage, and training (Figure 1 ). We define this approach more accurately in the following paragraphs. We define a segment σ as a length L sequence of transitions. During collection, episodes E longer than L transitions are split into fragments F . Fragments are then zero-padded to be length L, resulting in fixed-size segments σ and associated masks m. The resulting segments and masks are stacked into a dataset D, enabling easy batching, storage, and training (Figure 1 ). The split function splits a single episode E into one or more fragments F , each of size L except for the final fragment. . . .EQUATIONThe pad function zero pads a fragment F into a fixed size segment σ and associated mask m denoting the padding elementsEQUATIONUsing our split and pad operators, we split and pad each incoming episode, producing one or more segments and associated masks for each episodeEQUATIONWe represent our training dataset D as the concatenation of segments and masksD = concat         σ 0 ,m 0 . . . σ k ,m k σ k+1 ,m k+1 . . . σ j ,m j . . .         =         σ 0 ,m 0 . . . σ k ,m k σ k+1 ,m k+1 . . . σ j ,m j . . .         (20)During training, we randomly sample rows from D for minibatching (on-policy) or experience replay (off-policy).Theorem D.1. The discounted cumulative return given byG = ∞ t=0 γ t r t (21)is equivalent to computing the following memoroid over r 0 , r 1 , . . .EQUATION)Proof. We prove the correctness of our discounted return memoroid by showing the expansion is equivalent to the discounted return.(1, 0)• (γ, r 0 ) = (γ, r 0 + 0) = (γ, r 0 ) (27) (1, 0) • (γ, r 0 ) • (γ, r 1 ) = (1 • γ • γ, 0 + 1 • r 0 + γr 1 ) = (γ 2 , r 0 + γr 1 ) (28) (1, 0) • (γ, r 0 ) • . . . (γ, r n ) = (1 • γ • γ • • • • γ, 1 • r 0 + γr 1 + . . . γ n r n ) (29) = γ n , n i=0 γ i r i (30)If we let n → ∞, we see that the second element in the monoid tuple approaches the discounted returnEQUATIONLet us define Generalized Advantage Estimation (GAE) in memoroid form: Theorem E.1. The GAE target given byA t = ∞ l=0 (λγ) l δ t+l ; δ t = r t + γV (s t+1 ) -V (s t ) (32)is equivalent to computing the following memoroid over δ t , δ t+1 , . . .EQUATION)Proof. We prove the correctness of our GAE memoroid by showing the expansion is equivalent to the GAE target. This proof is very similar to the proof of the discounted return.(1, 0)EQUATIONIf we let n → ∞, we see that the second element in the monoid tuple approaches the GAE targetlim n→∞ n l=0 (γλ) l δ t+l = ∞ l=0 (λγ) l δ t+l (42)In this section, we reformulate existing models used in our experiments as memoroids. This reformulation is necessary to use inline resets for these models.The Linear Transformer from Katharopoulos et al. (2020) written asEQUATION)can be reformulated as the following memoroidEQUATIONwhere ϕ(x) = 1 + ELU(x).Prior work (Lu et al., 2024) defines an 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 m×m , x ∈ R m×1 } (63) H I = (I m , 0) (64) (X, x) • (X ′ , x ′ ) = (X ′ X, X ′ x + x ′ ) (65) f (o, b) = (W X , W x o) (66) g((X, x), (o, b)) = (W 1 GeLU(W c x) + b 1 ) ⊙ sigmoid(W 2 GeLU(W c x) + b 2 ) (67)where W X , W x , W c are learnable weights, b 1 , b 2 are learnable biases, and I m is the square identity matrix of size m.The Linear Recurrent Unit (Orvieto et al., 2023 ) could be roughly described as a theoretical simplification 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 differentEQUATIONFinally, we can rewrite Fast and Forgetful Memory (FFM) as a memoroid, with the parallel scans simplifying its implementation and fixing numerical instabilities caused by large positive exponentials over long sequences, as discussed in Morad et al. (2023b) . The original formulation is written as an aggregator and cell. First, let us write down the Γ term used in the aggregator.EQUATIONWe then write the aggregator asEQUATION)where ⊙ is the Hadamard product (or power), m is the trace size, c is the context size, and α ∈ R m + , ω ∈ R c are trainable parameters representing decay and context respectively. Multiplying column a vector by 1 ⊤ c "broadcasts" or repeats the column vector c times. The cell is defined asEQUATION)Agg represents the aggregator (Equation 75) and ℓ represents linear layers with mappings ℓ 1 , ℓ 2 :R d → R m , ℓ 3 : R m×2c → R d , ℓ 4 , ℓ 5 : R d → R d .ℜ, ℑ extract the real and imaginary components of a complex number as reals, Flatten reshapes a matrix (m × c → mc) and || is the concatenation operator. LN is nonparametric layer norm, and σ is sigmoid activation. We reformulate and simplify Γ, the FFM aggregator, and the cell as a single memoroidEQUATIONwhere W, b are learnable weights and biases, ℜ, ℑ extract the real and imaginary part of a complex number, ⊙ is the elementwise product, ⊕ is an outer sum, and α ∈ R n , ω ∈ R m are learnable parameters. Note that the original FFM formulation requires distributing Γ(-t) = exp t(|α| + iω) into the sum. Since α is learned, the real component can grow very large and cause numerical instabilities as it overflows even a double precision float. This is discussed in the limitations section of the original paper. Since our formulation utilizes a 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 a time-varying recurrence because the recurrent updates depends on t.The code necessary to reproduce all of our experiments is available at https://github.com/ proroklab/memory-monoids. 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 in prior literature. Once we arrived at a good set of hyperparameters, we simply reused them for our TBB method.We ran out of GPU credits early in the paper. We estimate roughly 70% of experiments were run on CPU only, across a number of hardware configurations. Thus, it is not straightforward to arrive at a single number. Users should be able to run at least one seed for each experiment we did, on a reasonable laptop, over approximately one week.We construct our model using blocks. A block contains a linear layer with nonparametric layer normalization and leaky ReLU activation. Observations feed into a block, followed by a memory model, followed by two more blocks. The hidden size of all blocks is 256 dimensions. For the S5 and LRU models, stacked two S5 and LRU layers, resulting in a sum of 512 dimensions of recurrent state (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.For each task, we selected a 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 τ determines the target network update rate. Batch Size measures the batch size in transitions for each model update. LR is the learning rate with a linear warmup over a specified number of model updates. The ratio describes the number of episodes collected at each epoch, compared to the number of model updates per epoch. 1 : 2 means we would perform 2 gradient updates for each 1 episode collected. ∇ Clip corresponds to gradient clipping, where the gradient magnitude is rescaled to at most ∇ Clip. γ is the decay term used in MDPs. We use a linear learning rate warmup of 200 updates for all tasks. In the Figure 5 plot, we test the wall-clock efficiency of our discounted return monoid against the standard approach of iterating over episodes in a batch. Both the monoid and standard approach are just-in-time compiled on a GPU, however the standard approach requires a for loop when the episode lengths are not fixed. We sample a batch of episodes, where each episode length is sampled from a discrete uniform distribution between one and a maximum episode length. We find that our memoroid computes the discounted return between three orders of magnitude faster. Next, we compare TBB and SBB scaling. TBB scales worse than SBB (O(log B) and O(log L) respectively, where B is the batch size and L is segment length). We question how this overhead translates to wall-clock training time. In the Figure 5 Justification: We detail this in the appendix. We do not have a hard number, but we explain roughly how much compute would be required.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.Justification: There were no human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.38th Conference on Neural Information Processing Systems (NeurIPS 2024).Assuming (1) The binary operator • and f, g are constant-time and constant-space, which is the case for all Linear Recurrent Models listed thus far. (2) Our processor has n parallel threads of execution. 


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

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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
- 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
questions: 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?

- "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: limitations: I agree with the limitations listed by the authors.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.


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

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

- 2. **Tone**

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




[(Williams and Peng, 1990)]( https://direct.mit.edu/neco/article-abstract/2/4/490/5561/An-Efficient-Gradient-Based-Algorithm-for-On-Line)


- 3. **Some missing Related Works**

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:
- 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.
- 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
questions: 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: limitations: Yes.
paper_summary: paper_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.

The 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.
strength_weakness: strength_weakness: 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. 
2. 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?
3. 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.
questions: 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?
2. 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? 
3. Could the authors please elaborate on the ‘What are the Consequences of Truncating BPTT’ experiment?
limitations: limitations: Yes, the authors mention the limitations of their method.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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?
questions: questions: - The tradeoff between L and B for SBB is nowhere given formally (did i miss it)?

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

- 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: limitations: The authors are upfront about limitations concerning an increased time cost of the approach compared to segment-based batching.

Comments: title: Thank you for the response
comment: Thank you for the response.

In the final paper, we will have:
- Two Atari experiments
- Rewritten the `What are the Consequences of Truncating BPTT` paragraphs to be more intuitive, roughly following the above explanation
- Moved some of the more technical bits of this experiment to the Appendix
- A sentence explaining why SBB is not faster than TBB
- A sentence explaining that SBB with fixed length episodes, and $L$ set specifically to the episode length, produces identical outputs to TBB

If we have adequately addressed your concerns, please consider updating your score.
I would like to thank the authors for their detailed response and clarifications and especially for the explanation in the Questions section! 

It would add to your paper if you could include something similar to this explanation in the final manuscript.

I also appreciate running more experiments and for the clarification on the `CountRecall` task and the PopGym benchmark suite.
Thank you for the prompt reply and insightful feedback.

While TBB is certainly useful, it is only one of our contributions. We argue that the unifying memoroid formalism, inline resets, and $O(\log n)$ advantage and return operations are probably more useful to a broad range of readers than TBB. Yes, comparing TBB to the random update would be interesting. But we argue that the current paper is sufficiently useful for the community at large.

Let us discuss the tone next. The papers by (Hausknecht & Stone, 2015; Kapturowski et al., 2019; Igl et al., 2018; Ni et al., 2024) are foundational works and we did not mean to disparage them. We have deleted

> Prior work (Hausknecht & Stone, 2015; Kapturowski et al., 2019; Igl et al., 2018; Ni et al., 2024) $\nabla \approx \nabla_\sigma$, although we are not aware of any theoretical justification for this assumption, as the error between the true and approximated gradient is unbounded. In fact, our experiments show that $\nabla_\sigma$ is often a poor approximation of $\nabla$.

as these sentences are not necessary, given that:
1. We now have a few sentences at the start of Section 3 discussing (Williams and Peng, 1990) as a necessity for tractability in deep learning (not just RL)
2. The following paragraph already explains the (Hausknecht & Stone, 2015; Kapturowski et al., 2019) approaches more thoroughly and in a better light -- as _solutions_ to the gradient truncation problem.

We hope we that have addressed at least some of your concerns. If we did, we hope that you reconsider your reject rating.
Thank you for the prompt response. Your feedback has already helped us improve the readability of the paper.

> Is the $P$ here also a tuple $(o, b)$?

Yes, our apologies. It should be $\overline{P}$ and explicitly listed under the algorithm `Input:`. We have updated the text.

>  how is the Markov-state mapping $M$ in Algorithm 3 different from Algorithm 4?

We partially address this in lines 215-222, but these lines are on another page. We will see if we can fit the algorithm and paragraph on the same page. The function $M$ itself is identical between Algorithms 3 and 4. The difference is how we structure the inputs to $M$ (either as one long TBB sequence, or a $B \times T$ SBB batch). Note that this is only possible with the reset operator.

> This is what I was referring to. I think that Figure 7 does not paint such a clear picture. In general, I feel like this paper would benefit from investigating more thoroughly the reason for SBB's inferiority. For example, comparing magnitudes and principal components of gradients for different $L$ or $B$; Or showing that SBB with  equal to the episode length produces the same updates as TBB.

We agree. Originally, we were focused on simplifying recurrent loss functions -- we were not expecting TBB to outperform SBB by so much. The sensitivity study was a late addition. 

Surprisingly, we show that the unnecessary dependence on prior inputs is still an issue with TBB, though to a much lesser extent. If we can backpropagate over the whole sequence, why are we still sensitive to prior inputs? To what extent does this occur, in say, Atari games? Is this a property of the objective or target network? The RNN weight initialization? The optimizer? It is not clear.

It was already a challenge to fit memoroids, parallel return computations, inline-resets, and TBB into the same paper while maintaining readability. We suspect that figuring out and fixing the root cause could be a paper on its own.

Thank you for updating your score.
Thank you for the comments! How to use the replay buffer is a deeply understudied part of using recurrent networks in Deep RL, and I want this paper to be successful. While I laid some heavy criticism initially (some of which I think is still valid), I do believe the ideas presented here are good ideas (that is not in dispute). This conversation is around how to make them as impactful as they can be.

The authors are correct. I was mistaken and had forgotten that statement about Kapturowski et al, and missed it when writing this review. While Kapturowski et al do do this, other papers such as Hausknecht and Stone (and several following papers linked below) follow the random update strategy I initially described. While there is room to compare to SBB, I believe there is a need to also compare with the random update strategy (Kapturowski also doesn't really do this comparison in a straightforward manner i.e. afaict there is no ablation showing random starting points vs SBB). This part of the architecture does tend to be under-reported, so maybe it is worth going through and trying to uncover what the choices are in the background sections of this work. I don't think relying on open-source package's implementations as a way to decide what algorithms to compare against is a valid decision. 

Some other papers using random updates, but others exist:
- https://arxiv.org/pdf/1704.07978
- https://arxiv.org/pdf/2110.05038: You link the codebase of this paper. On inspecting the sequence replay buffer code, I do believe they use the random strategy. You are correct they don't do burn-in though.
- I was about to link to Impala (https://arxiv.org/pdf/1802.01561), but their explanation of their batching strategy is sufficiently obscured that it could be either. I don't have time to dig in to the code (https://github.com/google-deepmind/scalable_agent/), but might be another example.

I think there was a misunderstanding of my critic of the language surrounding prior work. While I agree we need to be critical of previous approaches, there is a difference between stating weaknesses of the work in a constructive manner and implying previous work was misleading/dishonest. Language like "Prior work <> assumes $\nabla \approx \nabla_\sigma$ but we are not  aware of any theoretical justification..." is not generous when understanding prior work, and to me implies the previous work mislead their readers. While I agree, it would be better to clearly state the underlying assumptions we are making in our algorithms, the language of this paper suggests [Hausknecht and Stone] and following papers were the first to make such an assumption. There are clear reasons why the assumption was made in context of the literature, and is a well known assumption made throughout the use of recurrent models in all of machine learning. A better way to approach this would be to lay out the history of the assumption, and note that several of the RL works use this assumption as well.

I have also updated my scores/review to reflect my misremembering of Kapturowski. I think this paper could be accepted now, but would be more impactful taking into account some of the other weaknesses.
I would like to thank the authors for answering my questions. I’m particularly pleased with their experiments on Atari asteroid, which reassures my rating of Accept. I understand that experiments on Frostbite and memory maze would require more forethought and run times. I’m looking forward to seeing them in the final manuscript.
> Would it be helpful to move Algorithms 1, 2, 3 before their descriptions in the text?

Algorithms 3 and 4 are the part that confuses me. Is the $P$ here also a tuple $(o,b)$? If so, how is the Markov-state mapping $M$ in Algorithm 3 different from Algorithm 4?

> We have added a primer on parallel scans to the Appendix.

This addresses my concern here adequately.

> If you are referring to the VML/RML plots, we also investigate another environment in Figure 7 in the Appendix. 

This is what I was referring to. I think that Figure 7 does not paint such a clear picture. 
In general, I feel like this paper would benefit from investigating more thoroughly the reason for SBB's inferiority. For example, comparing magnitudes and principal components of gradients for different $L$ or $B$; Or showing that SBB with $L$ equal to the episode length produces the same updates as TBB.

> Exceptions include CountRecall and Autoencode.

Honestly, I was expecting that all these environments are supposed to be "solvable" in the sense of being published alongside a model that reaches a reward of 1. After going back to the POPGym paper I realized that this was not the case.

> We feed TBB and SBB the same number of transitions for each update. 

This very sentence is missing in the manuscript. Thank you for clarifying!

After reading the rebuttal, I have to admit that my initial judgment was indeed a bit harsh due to some misunderstandings on my part. Therefore, I increase the rating to Weak Accept.