Efficient Recurrent Off-Policy RL Requires a Context-Encoder-Specific Learning Rate

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, previous recurrent RL methods face training stability issues due to the gradient instability of RNNs. In this paper, we propose Recurrent Off-policy RL with Context-Encoder-Specific Learning 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.

Paper

References (51)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer B2hv3/10 · confidence 4/52024-07-01

Summary

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

Strengths

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

Weaknesses

Unfortunately, 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. Experimentally, 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. Ultimately however, my opinion is that the idea is not substantial enough for publication in a top venue. [1]: Hochreiter, Sepp, and Jürgen Schmidhuber. "Long short-term memory." Neural computation 9.8 (1997): 1735-1780.

Questions

1. How were hyperparameters tuned in the experiments? 2. Have you considered the approaches mentioned in the weaknesses? (gradient clipping and truncated gradients)

Rating

3

Confidence

4

Soundness

2

Presentation

4

Contribution

1

Limitations

Yes.

Reviewer B2hv2024-08-12

Re: Rebuttal

Thank you to the authors for the effort put into the rebuttal. I appreciate the inclusion of an investigation on gradient clipping and truncating the number of gradient steps, although I would have liked to see truncations with a smaller number of steps as well. Many implementations consider 4 or 8 steps (e.g. [1,2]) which may do well given that the output variation is exponential in the number of steps -- truncating to 32 steps doesn't provide much insight into whether it can tackle the exponential variation. I'm still not convinced about the novelty of the work after reading the authors' rebuttal and the other reviews. The authors claim that the novelty is not so much in their implementation trick but in the theoretical and intuitive understanding gained from this work. However, the idea that RNNs tend to be more unstable than MLPs is not a particularly new or surprising finding. Or is the novel understanding something different from what I just described? In reply to Reviewer 89BQ's argument for the novelty of this work, I can agree that the trick of using separate learning rates is not considered common practice. With a more rigorous investigation to ensure that current methods don't already solve this, such as gradient truncation with a smaller number of steps, I can agree that this knowledge could be useful to many practitioners (though perhaps it doesn't require a 9-page NeurIPS paper to disseminate). However, this paper doesn't seem to offer much beyond that. Most papers that propose a new algorithm also come with new and surprising insights that future works can build off of. Unfortunately, with this work, I don't think the insights are sufficiently novel. I don't see them pushing a research area forward, inspiring future works, or significantly improving our scientific understanding. I've increased by score since the experimentation was improved, but unfortunately my main complaint regarding lack of novelty still stands. [1] https://github.com/MarcoMeter/recurrent-ppo-truncated-bptt?tab=readme-ov-file [2] https://github.com/lcswillems/torch-ac?tab=readme-ov-file

Authorsrebuttal2024-08-12

Thank you for your reply. It appears that there may be a misunderstanding. The reviewer seems to believe that the instability in Recurrent off-policy RL is due to numerical explosion or instability of RNN gradients, and that addressing the gradient issues alone would resolve the overall training stability. This understanding aligns with the traditional issues of RNNs and corresponding solutions. However, there will be different issues in off-policy RL. Our experimental results (as shown by the red line in Fig. R3) indicate that the gradient explosion problem is not significant, as advanced RNNs (such as GRU, Mamba, etc.) have already addressed gradient instability quite effectively. If the gradients were indeed growing exponentially with the number of steps, merely reducing the RNN learning rate by 30 times would not be sufficient to counteract this exponential divergence, whereas gradient clipping would effectively mitigate such issues. Contrary to this, the experimental results show that our method is significantly more stable than gradient clipping. This is because the fundamental cause of instability in Recurrent off-policy RL is not gradient instability, but rather excessively large output variations between consecutive updates of the RNN. Even if the gradients are not large (on the same order of magnitude as those in MLPs), the outputs of the recurrent policy/value function can vary significantly between updates. These large output changes introduce instability into RL. For instance, in value function training, the (simplified) optimization target is $||Q(s,a)-r(s,a)-\gamma\hat{Q}(s',a')||_2^2$, where $\hat{Q}$ and $a'$ are the target value function network and the outputs of the current policy on $s'$, respectively. If the outputs of the value function network and policy network vary significantly, $a'$ will change greatly, and $\hat{Q}(s',a')$ will also fluctuate substantially. This results in large shifts in the optimization target of the value function after each update, leading to training instability. Similar instabilities occur in policy training, contributing to overall RL training instability. This is why, as shown in Fig. R4, even after clipping the gradients, the value loss remains unstable. **To the best of our knowledge, this finding has not been reported in previous work.** We will include these discussions in the revised version of our paper. Our theoretical results demonstrate that output variation does not increase exponentially with the number of steps but converges to a certain value. Therefore, we can balance the output variations, amplified by time steps, with a small RNN-specific learning rate. In contrast, if gradient clipping or truncation is used to keep the gradients within a normal range (similar to MLPs), the RNN output variations will still be large, and RL training will remain unstable. The results of Gradient truncation with 4 or 8 truncation steps are shown in the table below. In both variants, the algorithm produced infinite outputs before 250 iterations, leading to early stopping. | | $LR_{CE}$=$10^{-5}$ | $LR_{CE}$=$3\times10^{-4}$, grad-step-truncation-4 | $LR_{CE}$=$3\times10^{-4}$, grad-step-truncation-8 | $LR_{CE}$=$3\times10^{-4}$, grad-step-truncation-32 | | :-----------------: | :-------------------------------------: | :------------------------------------------------: | :------------------------------------------------: | :-------------------------------------------------: | | AntBLT-V-v0 | $ \mathbf{1986} \pm\mathbf{73} ^\star$ | $408\pm 74$ | $269\pm 88$ | $499±55$ | | HalfCheetahBLT-V-v0 | $ \mathbf{2679} \pm\mathbf{176} ^\star$ | $-107\pm 320$ | $-703\pm 330$ | $−458±244$ |

Reviewer B2hv2024-08-12

Re: Additional Comment

Thank you for clarifying the difference between gradient instability vs large output variations and also for gathering further results for truncated BPTT. This is much appreciated. I now realize that I misinterpreted the equation in Proposition 1 to mean the output variation exponentially increases with $t$ when in fact it is bounded by a constant factor. I will reassess the work with this in mind. It would help if you could answer the following questions: - Could you give insight on the scope of the results with respect to off-policy RL vs RL in general vs supervised learning? The theoretical foundations don't seem particularly tied to off-policy RL, and the other paradigms also use neural architectures with recurrent and MLP layers. Is there any reason this work specifically targets off-policy RL? > To the best of our knowledge, this finding has not been reported in previous work. - When you mentioned the above statement, which finding are you specifically referring to? Is it that "the fundamental cause of instability in Recurrent off-policy RL is not gradient instability, but rather excessively large output variations between consecutive updates of the RNN"?

Authorsrebuttal2024-08-12

Thank you very much for your prompt response. We will address each of your concerns in detail. 1. **[The scope of the results]** We observed that significant RNN output variations can introduce instability in value learning when using bootstrap-based methods. This is because the target for value learning, $r(s,a) + \gamma \hat{Q}(s',a')$, experiences substantial changes. Off-policy RL algorithms such as TD3, SAC, TD7, and DQN, which heavily rely on value estimation, are especially prone to this instability. Additionally, in actor-critic off-policy RL methods, the optimization objective of the policy, $\max_{\pi} Q(s,a)$, also experiences large fluctuations, leading to even greater instability. On the other hand, on-policy algorithms, especially PPO and TRPO, have trust region constraints that prevent significant changes in policy outputs, which inherently mitigates the impact of RNN output variations. However the value function in these methods can still exhibit instability. We believe that incorporating the RESeL framework into on-policy algorithms could further enhance training stability. In supervised learning, the impact of large RNN output variations is less significant. This is mainly due to the use of fixed labels, which contributes to more stable training. Additionally, certain loss functions in supervised learning, such as MSE loss in the case of continuous labels, naturally suppress the issue of large RNN output variations as the gradient decreases rapidly with the loss. We will include this discussion in the revised version of the paper. 2. **[This finding]** Yes, we are the first to identify this fundamental cause of instability in recurrent off-policy RL. Please feel free to reach out if you need further clarification on any point. We are more than happy to answer any additional questions you may have.

Reviewer vPkL7/10 · confidence 4/52024-07-07

Summary

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

Strengths

The 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). The paper is well written and easy to follow. The 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. The ablation experiments on section 5.1 and 5.3 clearly shows that a smaller learning rate helps stability (through reduced action variation).

Weaknesses

No major weaknesses. Please see the questions below. One 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.

Questions

About architectural design. Is there a reason for using respective context encoders in actor-critic? About 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)? Is this a typo in the unusually large batch size described in Table 2? As the main finding is on the learning rates, could the sensitivity analysis (Figure 9) be applied to more tasks? Indicated by Figure 10, is RESeL-GRU more sample-efficient than RESeL-Mamba? Is Mamba a crucial component of RESeL?

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

4

Limitations

Yes, no major limitations.

Reviewer vPkL2024-08-13

Thank you for the rebuttal and the additional experiments, although what I expected was a thorough sensitive analysis across all the tasks (currently only 3 tasks are shown). Overall the paper and the rebuttal looks good to me. I keep my rating.

Authorsrebuttal2024-08-13

Thank you for your response. Due to time and computational resource limitation, we were only able to add two additional environments during the rebuttal phase. We will continue with sensitivity studies and include the results and discussion of the remaining tasks in the next version of our paper. We appreciate your constructive suggestions, which have helped improve our paper.

Reviewer dCQa5/10 · confidence 3/52024-07-12

Summary

The paper mitigates the training instability issue of the recurrent off-policy RL by using a smaller learning rate for the RNN block.

Strengths

- The paper is well-written and easy to follow. - The paper proposes a simple solution with analysis. - The proposed solution is thoroughly verified in different experimental settings.

Weaknesses

- 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. 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. - Some details need further explanation: - In equation (1), it could be better to distinguish two $\epsilon$. - Will the amplification factor $\beta = \frac{K_y}{1-K_h}$ always larger than 1 (L160)? - 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)? - 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? - In L218, why does the learning rate of MLP must be increased **twentyfold**? - 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?

Questions

See weakness above

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

The limitation is discussed.

Reviewer 89BQ7/10 · confidence 4/52024-07-25

Summary

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

Strengths

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

Weaknesses

Even 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. Having 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. As 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.

Questions

l. 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) l. 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? l. 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. l. 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)? l. 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. l. 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? Figure 3: Why are error bars only visible for the brown curves/datapoints? l. 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? l. 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. Fig. 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 Minor remarks: l. 126: Better don't use " ' " behind name of algorithm ("RESeL's") but speak of RESeL policy as a joint expression l. 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. l. 240: SOTA (state-of-the-art) should be introduced as an abbreviation Fig. 5 (and others): Quality of plots could generally be improved by avoiding overlapping of images and axes captions in some cases Fig. 6: Colors/symbols of "EPI" and "SAC-MLP" are hardly distinguishable

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

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

Reviewer 89BQ2024-08-08

I'd like to thank the authors for their detailed response to my remarks and concerns, and in particular for their additional study of bypassing the CE to provide another baseline for the performance evaluation. Together with other explanations and clarifications done in the revised version my remaining questions and suggestions also have been sufficiently addressed. Hence, I continue to vote for "Accept (7)". Towards the main concern regarding lack of novelty, as raised in my own review and reflected by some of the other reviewers' remarks: It is true that tuning the LR in any context of machine learning is on the checklist of every good practioner, so I was tempted to come to the same conclusion in the beginning that this work doesn't add enough novelty regarding this aspect. What changed my final perspective is related to the answer given by the authors in their rebuttal to some of the other reviews: In the particular context of multi-component POMDP-RL solution architectures it is not common practice to use separate LR for modules which otherwise contribute to the same overall loss function. And even if that idea was previously applied in this particular context, it would highly likely not entail a comprehensive study as provided in the present work.

Reviewer dCQa2024-08-12

Thanks for the detailed reply and the efforts in preparing the rebuttal. After reading your rebuttal, I agree that although only a separate learning rate is used for RNN, the paper explains the motivation behind this and supports the conclusion with good experiments. I would improve my score to borderline accept.

Authorsrebuttal2024-08-13

Thank you

We are glad that our rebuttal was able to address your concerns. Your suggestions have been very helpful to us. On the other hand, should you have any additional questions or lingering concerns, please do not hesitate to contact us.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC