Reinforcement Learning with Fast and Forgetful Memory

Summary

The paper proposes a new memory model for reinforcement learning agents, targeting POMDP setting. The memory (FFM) aggerates past data via an exponentially weighted summation operation where the weights are learnable parameters. The memory is used to compute the Markov state through sophisticated gating mechanisms. The paper evaluates FFM on 2 POMDP benchmarks and shows improvement over baselines such as GRU.

Strengths

- The memory design is motivated by computational psychology - The proposed memory is light-weight and fast

Weaknesses

- The novelty is questionable. The main contribution is Eq. 11, and it is very similar to Eq. 1 in [1]. The difference is that the decay term here is learnable and consists of real and imaginary components. Table 2 reveals that non-learned decay also shows decent performance, which questions the need of learning the decay. - The literature review lacks a line of research using MANN on RL [2,3,4]. They are memory models specifically designed for RL. It is fairer to compare FFM with them rather than comparing with general memory models such as DNC or GRU. Please note that general memory models may require heavy hyperparameter tuning and adjustment to make them work on RL, thus comparing with well-established memory-based RL methods is better. For example, Transformer needs some tricks to work on RL [5]. Please consider using this baseline instead of general Transformer models. - The experimental results are not convincing. In POPGym, the gap between FFM and GRU is not clear. The ablation study in Table 2 indicates that the full model does not outperform other configurations significantly, especially when the std is high and only 3 runs are executed. It would suggest that the complicated design might not add much to the performance gain. In Fig. 5, there are many tasks FFM shows similar or worse results than RNN. It questions the benefit of using FFM. [1] Ba, Jimmy, Geoffrey E. Hinton, Volodymyr Mnih, Joel Z. Leibo, and Catalin Ionescu. "Using fast weights to attend to the recent past." Advances in neural information processing systems 29 (2016). [2] Wayne, Greg, Chia-Chun Hung, David Amos, Mehdi Mirza, Arun Ahuja, Agnieszka Grabska-Barwinska, Jack Rae et al. "Unsupervised predictive memory in a goal-directed agent." arXiv preprint arXiv:1803.10760 (2018). [3] Fortunato, Meire, Melissa Tan, Ryan Faulkner, Steven Hansen, Adrià Puigdomènech Badia, Gavin Buttimore, Charles Deck, Joel Z. Leibo, and Charles Blundell. "Generalization of reinforcement learners with working and episodic memory." Advances in neural information processing systems 32 (2019). [4] Lampinen, Andrew, Stephanie Chan, Andrea Banino, and Felix Hill. "Towards mental time travel: a hierarchical memory for reinforcement learning agents." Advances in Neural Information Processing Systems 34 (2021): 28182-28195. [5] Parisotto, Emilio, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg et al. "Stabilizing transformers for reinforcement learning." In International conference on machine learning, pp. 7487-7498. PMLR, 2020.

Questions

- It is unclear why the inductive bias (composite memory and contextual drift) is designed specifically for RL setting. The bias is very general and can be used for both SL and RL. This raises the questions such as whether FFM can show good performance under SL? Why only test it on RL? - It is hard to link Eq. (11) to the contextual drift motivation where $\omega$ should be recurrently computed. - It is also unclear why the decay should be complex-valued. The authors should present the explanation for the changes compared to the original motivation Eq. 8-10. - The choice of gating for the Cell seems arbitrary. Could you explain the reasons behind these designs, or do you follow any prior papers/ideas?

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

More investigation is needed on the initialization of the decay. Now, it seems to be heuristic and unclear if it still works for other tasks or settings.

Summary

This work introduces a new recurrent neural network layer called Fast and Forgetful Memory (FFM), which, differently from traditional architectures such as GRUs and LSTMs, FFM is explicitly designed for deep reinforcement learning

Strengths

The authors present a novel and promising memory layer for reinforcement learning, the core ideas of the layer are intuitive. The theoretical foundation is well explained and easy to follow, I didn't spot any errors. Also the empirical analysis is extensive. Authors cover time and space complexity of the layer which is very beneficial to get a general idea of its advantages. In general I think is a good work and a strong accept

Weaknesses

No major weaknesses that I can tell, just a few minor points: - In the abstract I would not say so broadly that all model-free approaches are using architectures from SL I think that applies to RNNs but mainstream transformers used in RL were specifically redesigned for this framework - Line 104, I would explain what traces are - Line 117, remove brackets from citations that are part of the sentence - Page 4, I would do an effor to inegrate the footnote in the main text, it would make the section much easier to follow

Questions

- Do you foresee any possible issues of instability? - How long were the temporal dependencies tested in the experiments?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

I think authors did fine work addressing the limitations of their contribution at the end of the paper

Summary

The paper introduces Fast and Forgetful Memory (FFM), a memory model designed for reinforcement learning. FFM is based on theories of composite memory and contextual drift from computational psychology. It consists of two main components: an aggregator and a cell. The aggregator computes complex-valued summaries of inputs, while the cell performs input gating and extracts real-valued outputs from the summaries. Conceptually, the aggregator implements composite memory by blending individual traces and incorporating an exponential forgetting mechanism. It further enables relative-time reasoning by applying a temporal context. On a suite of POMDP tasks from POPGym and POMDP-Baselines, FFM-enabled recurrent RL algorithms outperform those running on standard RNN. Moreover, since FFM belongs to the class of hybrid memory models, which provides equivalent recurrent and closed-form formulae, it achieves linear space complexity and logarithm time complexity with parallelism. This makes FFM an efficient and performant memory model that can be used as a drop-in replacement for RNN in recurrent RL algorithms.

Strengths

- The paper proposes a novel memory mechanism for reinforcement learning by drawing inspiration from the field of computational psychology. It represents its cell state as a collection of traces and features numerous inductive biases for composition, forgetting, and relative-time reasoning. These inductive biases collectively lead to improved performance on RL tasks. - The experiments are substantial and thorough, providing ample evidence that FFM can be used as a drop-in replacement for RNN in recurrent RL algorithms. The ablation experiments and analysis provide further insight into the significance of each design choice. - The contributions of this paper are significant, especially considering the importance of memory models in partially-observable RL.

Weaknesses

- The intuition behind some inductive biases is elusive. Please refer to the Questions section.

Questions

Besides drawing inspiration from computational psychology, are there intuitive or formal reasons that the proposed inductive biases lead to better performance on reinforcement learning tasks compared to the various instantiations of RNN? In particular: - Why is storing memory as individual traces more desirable than having a single context vector? - Is there a benefit to using imaginary temporal context over standard position encoding? - Is there evidence that temporal context enables relative-time reasoning?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors address the following limitations in their paper: - Performance decreases with larger context sizes, presumably because of a larger number of periodic functions resulting in more local extrema in optimization. - Limitation in the number of experiments. - Repeated multiplication of exponentials can lead to a loss of numerical precision.

Summary

The paper aims to propose an algorithm-agnostic memory model specifically designed for reinforcement learning (RL). The proposed neural network (NN) model is claimed to offer advantages over existing sequence-based memory models like Recurrent Neural Networks (RNNs) and can be seamlessly substituted for RNNs. The key advantage over RNNs is that the model can process a sequence in parallel, making it more efficient compared to traditional RNNs. The main design components of the model include: A context drift formula that computes the embedding of sensory information with context. A decaying mechanism incorporated into the recurrent formula to compute a summary of the embedding. Gating mechanisms used to extract the final output from the learned summary of the embedding.

Strengths

1. Appropriate architecture may be suitable for a particular type of tasks. I agree that there might be some interesting special NN architecture that is particularly beneficial for RL tasks. Studying RL-specific NN structure is an important and practically useful research direction. 2. The proposed NN architecture seems to be new.

Weaknesses

The main weaknesses are: clarity of presentation, lack of justifications of some design choices, experiments are not persuasive enough. Clarification questions: 1. What is the difference between subscript t and n? I found it to be confusing in eq(9). What should be the relation between n and t there? 2. eq(4) already defined x_n as the embedding of observation-action pair, why eq(9) redefined this notation? And what is hat{x}? It is said that it is the raw sensory information, isn’t it part of observation variables? What is the context variable in RL? Do you simply define omega as a constant one vector? I strongly suggest the authors to use a particular RL environment to explain all these notations. I found it to be confusing when reading notations in the paper. Justify some designs: 1. As a new NN architecture, it is critical to study which component of the proposed methods matters. At least, the paper should study the necessity of exponential decay mechanism, outer product design (eq 9), etc. One possible baseline in the experiments could be simply adding a decaying mechanism in the hidden state of a regular RNN so that it can somehow show the importance of the decay mechanism. I have doubts on why decaying is good. Why it makes sense to set up an expiration for each feature unit? In existing literature, it is critical to always maintain a so-called coverage in the state space to learn an optimal policy (please refer to references [1][2][3]). If setting expiration for some feature units, doesn’t it mean manually discard some of the states? 2. The paper should also include discussions with successor features, as the method of computing the hidden state (i.e. eq(7)) is highly similar to how a successor representation is learned. An in-depth study of successor representation may result in better understanding of the proposed architecture. Experiments. I found the experiments a bit disappointing. As a general NN architecture, why not also test it on those commonly used RL domains like mujoco, Atari, etc. as mentioned in the abstract, “it is a drop-in replacement for RNN.” There are already existing work using RNN in a variety of RL benchmark tasks. It is better to test the proposed architecture on those well-tested RL domains. I would like to emphasize that i don’t think it is necessary to restrict the test on POMDP. In function approximation setting, a MDP can be always viewed as a POMDP. [1] An equivalence between loss functions and non-uniform sampling in experience replay by Scott Fujimoto et al. [2] Understanding and mitigating the limitations of prioritized experience replay by Pan et al. [3] regret minimization ER in off-policy RL et al.

Questions

see weaknesses.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

na

Rebuttal Response

Thank you for addressing my questions and for the extra effort for the changes, while the paper was quite clear I think they will make the paper even easier to follow for future readers. I keep my strong acceptance recommendation