Recurrent Reinforcement Learning with Memoroids

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.

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?

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.

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

The authors are upfront about limitations concerning an increased time cost of the approach compared to segment-based batching.

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.

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. 2. The formalism is clearly explained and the first sections of the paper are clear to follow. Segment-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. 3. 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. 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

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?

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

Yes, the authors mention the limitations of their method.

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.

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

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

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

Yes.

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.

Strengths

- The paper is very well-written. It was a pleasure reading this work! - Clear experiments that demonstrate the superiority of tape-based batching to segment-based batching - Fig.3 discusses the consequences of truncating BPTT, which is quite interesting - 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 - 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

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

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

I agree with the limitations listed by the authors.

Thank you for the response

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.