Title: Efficient Recurrent Off-Policy RL Requires a Context-Encoder-Specific Learning Rate
Venue: NeurIPS 2024
Accepted: yes
Keywords: Reinforcement Learning

Recurrent Off-Policy Reinforcement Learning

Partially Observable Markov Decision Process
Abstract: Real-world decision-making tasks are usually partially observable Markov decision processes (POMDPs), where the state is not fully observable. Recent progress has demonstrated that recurrent reinforcement learning (RL), which consists of a context encoder based on recurrent neural networks (RNNs) for unobservable state prediction and a multilayer perceptron (MLP) policy for decision making, can mitigate partial observability and serve as a robust baseline for POMDP tasks. However, prior recurrent RL algorithms have faced issues with training instability. In this paper, we find that this instability stems from the autoregressive nature of RNNs, which causes even small changes in RNN parameters to produce large output variations over long trajectories. Therefore, we propose **R**ecurrent Off-policy RL with Context-**E**ncoder-**S**p**e**cific **L**earning Rate (RESeL) to tackle this issue. Specifically, RESeL uses a lower learning rate for context encoder than other MLP layers to ensure the stability of the former while maintaining the training efficiency of the latter. We integrate this technique into existing off-policy RL methods, resulting in the RESeL algorithm. We evaluated RESeL in 18 POMDP tasks, including classic, meta-RL, and credit assignment scenarios, as well as five MDP locomotion tasks. The experiments demonstrate significant improvements in training stability with RESeL. Comparative results show that RESeL achieves notable performance improvements over previous recurrent RL baselines in POMDP tasks, and is competitive with or even surpasses state-of-the-art methods in MDP tasks. Further ablation studies highlight the necessity of applying a distinct learning rate for the context encoder. Code is available at https://github.com/FanmingL/Recurrent-Offpolicy-RL.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors investigate the stability of training deep recurrent policies for POMDP tasks. They hypothesize that the recurrent part of the encoder faces less stable training compared to the fixed-length parts of the encoder (e.g. the MLP input layer) and propose that the former should use a lower learning rate. This is shown to be effective mostly through an empirical evaluation.\n\nSTRENGTHS:\nThe paper's ideas are communicated clearly and the authors present excellently summarize the state of the art for several partially observable settings, including the general setting, meta RL, and credit assignment. I appreciated the wide range of experimental domains and the sensitivity analysis demonstrating that the recurrent layers and the rest of the network should have different learning rates for optimal performance. The suggestion to introduce a new hyperparameter for the recurrent encoder could be broadly adopted.\n\nWEAKNESSES:\nUnfortunately, I believe the work is not sufficiently novel for acceptance at NeurIPS. The problem identified by the authors --- that a recurrent encoder is numerically less stable than a fixed-length encoder --- has been identified many times over the years with various techniques proposed to mitigate it. This includes classical works such as the original LSTM paper [1] aimed at tackling exponentially scaling gradients. More recently, techniques such as gradient clipping and truncating the number of recurrent backpropagation steps (which is not the same as shortening the context length) have been suggested to tackle this problem as well.\n\nExperimentally, I would have liked to see that these last two tricks do not already solve the stability issues, as there was no mention of them in the paper. It also seems that the hyperparameters were tuned by hand (since I could not find any mention of a hyperparameter selection scheme) which I think is insufficient for a work for which the sole contribution is an empirical result that we are meant broadly apply. I would have expected to see a very rigorous and objective hyperparameter selection method such as a grid search over reasonable hyperparameters. While I appreciated the sensitivity analysis in Fig 9, it was only conducted on a single domain.\n\nUltimately however, my opinion is that the idea is not substantial enough for publication in a top venue.\n \n[1]: Hochreiter, Sepp, and Jürgen Schmidhuber. \"Long short-term memory.\" Neural computation 9.8 (1997): 1735-1780.\n\nQUESTIONS:\n1. How were hyperparameters tuned in the experiments?\n2. Have you considered the approaches mentioned in the weaknesses? (gradient clipping and truncated gradients)\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes RESeL which improves recurrent off-policy RL in POMDPs mainly by applying a lower learning rate to the context encoder. This is justified by their theoretical analysis that the recurrence amplifies the output difference in the long run. In practice, they also incorporate several techniques (Mamba architecture, critic ensembles, efficient data sampling, slow policy update). The main contribution is on strong empirical performance. Across several benchmarks (classic POMDP, meta-RL, credit assignment), RESeL attains SOTA performance.\n\nSTRENGTHS:\nThe problem setting of partial observability that this paper tackles is crucial and challenging. As indicated by the experiments, the partial observability is severe in some tasks and may require full context lengths (around 1000). \n\nThe paper is well written and easy to follow.\n\nThe performance gain averaged on a wide range of POMDP tasks is large enough to make RESeL a very strong empirical work. RESeL uses Mamba, which also accelerates training time a lot. \n\nThe ablation experiments on section 5.1 and 5.3 clearly shows that a smaller learning rate helps stability (through reduced action variation).\n\nWEAKNESSES:\nNo major weaknesses. Please see the questions below.\n\nOne missing point is about the theoretical understanding -- why a smaller learning rate on context encoder can stabilize training? Reducing the output variation between consecutive updates, in my opinion, is a starting point, but not enough to explain the stability as a whole. This is connected with two-timescale update of feature vs value learning, e.g., https://openreview.net/forum?id=rJleN20qK7 Perhaps this related work is worth a discussion.\n\nQUESTIONS:\nAbout architectural design. Is there a reason for using respective context encoders in actor-critic? \n\nAbout proposition 1. The connection between a smaller learning rate and the bound (1) is not very clear. Is the logic that a smaller learning rate leads to a smaller \\|\\theta - \\theta’\\| and thus a smaller \\epsilon, then a tighter upper bound in (1)? \n\nIs this a typo in the unusually large batch size described in Table 2? \n\nAs the main finding is on the learning rates, could the sensitivity analysis (Figure 9) be applied to more tasks? \n\nIndicated by Figure 10, is RESeL-GRU more sample-efficient than RESeL-Mamba? Is Mamba a crucial component of RESeL?\n\nLIMITATIONS:\nYes, no major limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper mitigates the training instability issue of the recurrent off-policy RL by using a smaller learning rate for the RNN block.\n\nSTRENGTHS:\n- The paper is well-written and easy to follow.\n- The paper proposes a simple solution with analysis.\n- The proposed solution is thoroughly verified in different experimental settings.\n\nWEAKNESSES:\n- My main concern is the novelty of the proposed solution. Although the paper gives a reason for using different learning rate for the RNN block and the MLPs, the proposed method is to choose a lower learning rate for the RNN block. In practice, it's somehow common to choose different learning rate for different components. For example, Dreamer [Hafner, Danijar, et al.], which also uses an RNN block in their model, chooses a smaller learning rate for the RNN-based world model. Also, when facing the training stability issue, tuning the learning rate is always on the checklist. \n\n    However, I would consider accepting the paper if it proposes an approach to automatically decide the learning rate of the RNN block by leveraging the analysis in section 4.2.\n- Some details need further explanation:\n    - In equation (1), it could be better to distinguish two $\\epsilon$.\n    - Will the amplification factor $\\beta = \\frac{K_y}{1-K_h}$ always larger than 1 (L160)?\n    - The average variation in network output $\\frac{1}{t}\\sum^{t-1}_{i=0} || y_i - y'_i||$ converges to $\\beta \\epsilon + \\epsilon$, which is not involved in $t$, but when this indicates with **longer** sequence lengths, the average variations in the RNN output induced by gradient deselect are amplified (L163-164)?\n    - In the caption of Fig 3, what does it mean by \"after one-step gradient-update\"? Does it mean one-gradient update per environment step?\n    - In L218, why does the learning rate of MLP must be increased **twentyfold**?\n    - In L221, \"The right panel shows the green and blue curve remain at similar levels until the final time step\". Where are the green and blue curves?\n\nQUESTIONS:\nSee weakness above\n\nLIMITATIONS:\nThe limitation is discussed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper contributes to the important and long-standing issue of representing latent state information for successfully finding optimal RL control policies in the context of partially observable Markov decision processes (POMDP). The major contribution is a newly composed RL algorithm (RESeL), which, by design, enables a thorough study of the impact and importance of the actual learning rate (LR) used for the recurrent (latent) context encoder part of the solution architecture. The author highlight the fact that current state-of-the-art (SOTA) methods commonly use an equal LR for training both the latent representation and the policy/critic networks. \n\nTheir idea of separating the LR between recurrent CE and policy networks, as well as the overall performance benefits w.r.t current SOTA algorithms, is demonstrated by an extensive survey of selected POMDP tasks and benchmarks. These experiments together with theoretical arguments let the authors conclude that using a single  (and often too large) LR parameter approach of previous methods leads to sub-optimal performance results in solving POMDP tasks.\n\nSTRENGTHS:\nThe paper presents a sound, well-written and focused approach of studying the impact of using dedicated LR parameters for both the context encoding and policy/critic related NN representation as part a newly proposed RL algorithm (RESeL) to solve POMDP tasks with high performance. It is implicitly stated (by not mentioning previous work) that the authors' work represents the first time that such a clear separation and impact analysis of different LR between the major architecture components is made, which - if true - points to a relatively high degree of originality.\n\nThe authors underline their central proposition (of better using lower LR for the CE related network parts) both by theoretical arguments/proof and an extensive experimental study on a relatively large set of well-known POMDP tasks and benchmarks, including direct comparisons to a large variety of SOTA algorithms. This presentation, together with additional material in the appendix, is well-suited to allow reproducibility and comparability of their results.\n\nWEAKNESSES:\nEven though mentioned in the Limitiation section, it would have been very helpful to add some benchmark comparisons of the RESeL algorithm with its RNN-CE block entirely bypassed/skipped to clearly separate the impact of the latent space presentation from the rest of the RL algorithm. In particular for the classic POMDP tasks, where only position and velocity features are masked, respectively, it is important to know a baseline performance when simply feeding the current and last-step observation vector into the MLP networks (given that RESeL uses 2 x 256 neurons layers in their MLPs, adding a few more inputs show be more than tractable). Including such extension (not necessarily for all benchmarks but maybe a smaller subset) could significantly improve the scientific insight.\n\nHaving clear focus on the impact of varying LR the remainder of the RESeL algorithm is entirely built from existing RL building blocks, like the overarching SAC RL architecture or convential RNN/GRU components and, hence, provides no original contribution in itself (as correctly mentioned by the authors, using RNN for latent space representation is known for a long time, as is the SAC approach). Re-using existing building blocks per-se is not a bad idea (and I don't recommend addressing this issue in the rebuttal phase) but one has to accept that the algorithmic novelty in this work is hence limited. \n\nAs a minor remark: the authors spend a relatively large portion of their introduction on the mathematical formulation of Proposition 1, whose central statement (variation of RNN outputs grow as function of rollout-step and LR) appears not entirely original (if it really is a first-time, the authors should highlight this fact more strongly) since the result is quite intuitive for people working in the RNN/RL community. Besides that, the connection between lowering the LR as part of the update rule when training RNNs and the resulting formula of Proposition 1 is not obvious and could be highlighted more strongly.\n\nQUESTIONS:\nl. 127: Why is the last-step observation also needed as input to the MLP if both current obs and context-encoded obs are already fed into the  NN? According to Fig. 2, the last-step obs is fed into the CE only (but not into the MLP)\n\nl. 132: Is there any justification/reasoning behind the choice of the MLP architecture, in particular why a relatively large capacity (2 x 256 neurons) was chosen? Is there an estimate of how much this choice impacts the overall performance results of the policy training?\n\nl. 163-164: It seems to me that the claim of Proposition 1 and the subsequential statements (\"average variations in the RNN output ... are amplified\") are not a novel discovery but have been studied and found by others in similar or different contexts of RNN before. In that case, it would be good to provide corresponding citations of pior art or - if the others believe this is the first time such a claim was made - highlight the significance of their finding. \n\nl. 165-166: Please provide further arguments/proof, why the effect of Proposition 1 can be mitigated by smaller learning rate (even if such claim sounds plausible, it would be good to refer to previous work or more detailed reasoning). In other words: how does the learning rate affect the result of ||y - y'|| as a result of Propsotion 1 (it is likely related to epsilon in the formula but if so this relation could be mentioned more explicitly)?\n\nl. 167-168: Similar remark as before: this sentence remarks a general claim about MLPs which should either be backed up by reference citations or by stronger reasoning or explanations.\n\nl. 195: Given the typical level of stochasticity or dependency on random intial states, a trial number of 6 seems relatively low to incorporate statistical fluctuations in the return evaluation, in particular for the following ablation studies. Have the authors made sure that their reported results are not prone to larger statistical errors? What in detail is included in the choice of \"different random seeds\"? Does it only refer to the initial weight setup of the RNNs/MLPs or does it reflect varying intial conditions of the environemnts? Do all environments provide deterministic rollouts?\n\nFigure 3: Why are error bars only visible for the brown curves/datapoints?\n\nl. 218: How is \"action variation after one-step update\" as a function of rollout steps defined in the case of MLPs as they don't have an intrinsic time-convolution? I.e., what is the meaning of the i-th rollout time step in the case of MLP \"only\" (LR_CE = 0)? And why is that variation not close to zero for the 0-th rollout step, as indicated in Figure 3 (gray dashed curve) but starting from a value around 0.6?\n\nl. 223-229: At this point of studying the performance in various benchmarks as a function of LR, and also for some of the subsequent ablation studies of Sec. 5.3, it would be great to show a benchmark reference where the CE module is skipped or bypassed altogether, feeding the current and last obserivation only into the MLP networks but not using a CE representation of the latent state at all (while there is one case where LR_CE = 0 in Fig. 9a, this is still different from skipping the entire module and following the suggestion above). This would generally help to underline the significance of the CE for the chosen set of benchmarks and tasks.\n\nFig. 6: It appears that SAC does hardly profit from using a recurrent CE (SAC-GRU) if compared to the MLP variant (SAC-MLP); is there any explanation why? It could stress the importance of showing the RESeL performance with and without CE, as mentioned in my comment above\n\nMinor remarks:\n\nl. 126: Better don't use \" ' \" behind name of algorithm (\"RESeL's\") but speak of RESeL policy as a joint expression\n\nl. 171: Printing style (formatting) of scientific number notation \"3e - 4\" looks somewhat odd; check if correct Math formula style was chosen in the LaTex document or even better use decimal notation (\"$3 10^-4$\") or capital \"E\" instead of \"e\". This applies to other instances of number formatting of the main text as well.\n\nl. 240: SOTA (state-of-the-art) should be introduced as an abbreviation\n\nFig. 5 (and others): Quality of plots could generally be improved by avoiding overlapping of images and axes captions in some cases\n\nFig. 6: Colors/symbols of \"EPI\" and \"SAC-MLP\" are hardly distinguishable\n\nLIMITATIONS:\nThe authors have addressed the main limitation well. As they state, what is missing is a study on the impact of using the RNN as a context encoder in general, which refers to my request for providing a reference baseline where the CE is bypassed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-25"}