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

Thank you for taking the time to review our paper and providing us with your valuable feedback.

• [LSTM] Our paper discovered a totally different issue from gradient exploding (numerical instability) resolved in LSTM: even if the RNN gradients do not explode and are similar in magnitude to those of MLPs, the gradient updates can still cause significant variations in the RNN outputs. These large output variations lead to instability in RL training, requiring a slow down in RNN updates. This is evident in Figure 3, particularly when the rollout step is 0. The action variation of the brown curve is lower than that of the orange curve, indicating that the RNN parameter changes are small and its gradient is not large. However, as the rollout steps increase, the brown curve grows by several orders of magnitude, which can destabilize reinforcement learning training.

• [The gradient clipping & truncated gradient tricks]

  • We have tried gradient clipping (0.5 maximum gradient norm as did in Dreamer) and truncated gradient (gradient maximum horizon 32) tricks, but we found that these tricks are irrelevant to the problem of RNNs in RL. Their purpose is to prevent gradient explosion in RNNs and cannot slow down the RNN updates. In Fig. R1, we show the effects of these tricks. It can be seen that gradient truncation may still lead to instability and infinite values in tasks like Ant and HalfCheetah. Gradient clipping can suppress extreme values but does not improve performance. We conducted an in-depth investigation and compared the policy gradient norms of RESeL and the gradient clipping approach. As shown in Fig. R3, the policy gradient norm of the gradient clipping method has many spikes, reflecting training instability. This instability does not stem from gradient value instability, as we have already employed gradient clipping.

  • We find the gradient instability originates from the instability of value function learning. In Fig. R4, we observe that in the baseline with gradient clipping but unchanged learning rates, the value loss frequently encounters outliers. This is due to large variation in the outputs of the policy and value networks, which make value function learning, based on bootstrapping, more unstable. This instability results in large value loss, which in turn causes spikes in the policy gradients. This issue is distinct from the instability and explosion of RNN gradients.

• [Sensitivity analysis was only conducted on a single domain] We have conducted similar experiments in WalkerBLT-P-v0 and AntBLT-V-v0. The results are shown in Fig. R2. Limited by computational resources, we only included the setting of Figs. 9(a) and 9(c) to illustrate that the optimal $LR_{CE}$ is around $10^{-5}$ in these tasks, and to show that $LR_{CE} = LR_{MLP}$ does not achieve optimal performance. The conclusions drawn from Fig. R2 are consistent with those presented in Figs. 9(a) and 9(c).

• [Hyper-parameters] Some hyper-parameters (discount factor, reward as input) were chosen based on the nature of the environments, while others (Context encoder LR, policy and value function LR, and batch size) are determined via grid search:

  • Context Encoder LR: We fixed the learning rate of the context encoder to be 1/30 of the policy learning rate in all tasks. We aim to have the action variation caused by RNN updates to be comparable to that caused by MLP updates, thus fixing the learning rate ratio between the context encoder and MLP. After conducting parameter searches in several environments, as shown in Figure 9, we find the 1/30 ratio performed well (this conclusion is also evident in Figure R2), so we fixed this ratio across all environments.
  • Policy and Value Function LR: We mainly followed the default learning rate settings in CleanRL [1], using a policy learning rate of $3 \times 10^{-4}$ and an action-value function learning rate of $10^{-3}$. We conducted a grid search between the aforementioned learning rates and those reduced by a factor of $5$ (aiming to further stabilize training). We found that a smaller learning rate was more stable in tasks where the partial observability is not very significant (e.g., MDPs like classic MuJoCo).
  • Discount Factor ($\gamma$): For general tasks, we used a common $\gamma$ value of 0.99. For environments requiring long-term memory (e.g., Key-to-Door), we used a larger $\gamma$ of 0.9999.
  • Reward as Input: The Classic Meta-RL tasks require inferring the true reward function based on real-time rewards. Thus, we also included real-time reward feedback as an input to the policy and value in these tasks.
  • Batch Size: We conducted a grid search for batch size, mainly exploring sizes that were 1x or 2x the maximum trajectory length. We found that in classic POMDP tasks, a larger batch size performed better. In other tasks, the advantage of a large batch size was not significant and could even reduce training speed.

[1] Huang et al. "CleanRL: High-quality single-file implementations of deep reinforcement learning algorithms." JMLR 2022.

We sincerely appreciate your time and expertise in reviewing our paper. We would greatly appreciate it if you could re-evaluate our paper based on the above responses.

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

We appreciate your time to review and provide positive feedback for our work.

• [Theoretical understanding] Thank you very much for pointing out this relevant work. It has been very enlightening for us. RESeL and TTN are closely related, with the MLP value and the small learning rate RNN encoder corresponding to the value function and representation network, respectively. The convergence guarantee of TTN also explains why a small output variation between consecutive updates can enhance training stability.

• [Architectural design] As noted in [1], sharing a context encoder could lead to a large gradient norm, causing instability. Therefore, we chose to use separate context encoders.

• [The connection between a smaller learning rate and the bound (1)] You are right, reducing the learning rate can lower $\epsilon$, thereby reducing the upper bound. We will include this discussion in the revised version of our paper.

• [Large batch size described in Table 2] The batch size here refers to the number of transitions. For each update, we collect at least one complete trajectory to calculate the loss. In environments like HalfCheetah, one trajectory consists of 1000 steps. So, the 1000 and 2000 mentioned correspond to one or two complete trajectories in these environments. To maintain consistency in the number of valid transitions across different environments, especially those with varying trajectory lengths, we use the same batch size for them. For these environment, we will dynamically sample uncertain number of complete trajectories to ensure that the number of valid transitions is no less than this batch size.

• [Sensitivity analysis in more tasks] Yes, we have conducted similar experiments in WalkerBLT-P-v0 and AntBLT-V-v0. The results are shown in Fig. R2. Limited by computational resources, we only included the setting of Figs. 9(a) and 9(c) to illustrate that the optimal $LR_{CE}$ is around $10^{-5}$ in these tasks, and to show that $LR_{CE} = LR_{MLP}$ does not achieve optimal performance.

• [Comparison between RESeL-Mamba and RESeL-GRU] It's true that GRU is more sample-efficient than Mamba in some environments. Mamba loses some expressive power due to its parallelzability (the hidden state at time $t$ is not used as input to the nonlinear network at time $t+1$ but is accumulated through a linear transition matrix) and also is not crucial in RESeL. However, Mamba significantly reduces computation time during training and requires much less memory in environments with variable trajectory lengths.

[1] Ni et al. "Recurrent Model-Free RL Can Be a Strong Baseline for Many POMDPs." ICML 2022.

Thank you for taking the time to provide us with your feedback. We appreciate your valuable comments and suggestions, which will help us improve our work. We look forward to receiving your further feedback.

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.

Thank you for taking the time to review our paper, and for your insightful comments.

• [Dreamer] It is not a common practice in reinforcement learning to use different learning rates for specific layers in a neural network. Typically, in reinforcement learning, the number of learning rates corresponds to the number of loss functions, with each loss function's associated network parameters sharing the same learning rate. For instance, in Dreamer, the world model is updated based on a prediction loss. Despite comprising various modules such as the non-RNN encoder, decoder, predictor, and RNN sequence model, these modules share a single learning rate in Dreamer families [1-3]. However, Figure 9(c) demonstrates that using the same learning rate for all layers during RL policy training could be suboptimal. Conversely, RESeL assigns a separate smaller learning rate specifically for the sequence model, while using a different learning rate for others.

• [Tuning the learning rate is always on the checklist] We are not naively tuning the learning rate here. We have shown that uniformly tuning the overall learning rate does not achieve optimal performance. Instead, we propose RESeL to specifically set a smaller learning rate for the RNN, based on our theoretical analysis and intuitive understanding. Despite being straightforward, it is not trivial.

• [An approach to automatically decide the learning rate of the RNN block]

  • In all experiments, the learning rate for the RNN was set to 1/30 of the policy learning rate. We found this setting to be highly generalizable. From a practical standpoint, we recommend setting it to 1/30 of the policy learning rate or conducting a grid search centered around this ratio.

  • On the other hand, for long trajectories, the upper bound of the average change in RNN outputs will eventually converge to $(1+\beta)\epsilon_{RNN}$, where $\epsilon_{RNN}$ is the variation of the RNN’s first step output. Our goal is to find a scaling factor $\alpha$ for the RNN learning rate so that the variation in RNN outputs matches $\epsilon_{MLP}$: $\alpha(1+\beta)\epsilon_{RNN}\approx\epsilon_{MLP}$. However, the values of $\beta$, $\epsilon_{RNN}$, and $\epsilon_{MLP}$ are highly dependent on the neural network weights. Thus, we can periodically test $(1+\beta)\epsilon_{RNN}$ and $\epsilon_{MLP}$ and use their ratio to determine $\alpha$. Specifically, for every 500 gradient updates, we perform a pseudo-update to observe how much the outputs of the policy change after individually updating the RNN block or the MLP. We then use the ratio of these changes to scale the RNN learning rate. To avoid instability in training caused by frequent changes in the learning rate, we only adjust the learning rate during the initial 50 iterations (a warmup phase, corresponding to 50000 gradient updates) and keep it fixed thereafter. We set the initial RNN learning rate to $3\times10^{-4}$. Our experimental results, shown in Fig. R6, indicate that the automated tuning method performs similarly to RESeL. We will add this discussion in our revised version. The final RNN learning rates are listed in the following table, which is close to our setting.

    Env. Name	Auto-Tuned RNN Learning Rate
    AntBLT-V-v0	$6.0\times 10^{-6}\pm 3\times 10^{-7}$
    HalfCheetahBLT-V-v0	$1.1 \times 10^{-5}\pm 1\times 10^{-7}$
    HopperBLT-V-v0	$1.1 \times 10^{-5}\pm 4\times 10^{-7}$
    WalkerBLT-V-v0	$1.2 \times 10^{-5}\pm 1\times 10^{-7}$

• [Distinguish two $\epsilon$] Thank you for the suggestion. We will distinguish the two $\epsilon$ in the revised version.

• [Is $\beta$ always larger than 1] It depends on the weights and gradient magnitude of the neural network. However, $\beta$ is always larger than $0$, so the upper bound of $||y_i-y_i'||(i>0)$ is always greater than $\epsilon$.

• [Average variation is not involved in $t$] This needs to be understood in conjunction with Eq. (11) in Appendix B. Since the right-hand side of Eq. (11) is an increasing function of $t$, the mean value increases over time. We will include this explanation in the revised paper.

• [After one-step gradient-update in Fig. 3] The action variation quantifies the difference in policy output after the gradient update compared to its output before the update, as considered in proposition 1. "After a one-step gradient update" indicates that we updated the policy only once, but we compare the differences in action outputs at various rollout steps before and after the policy update.

• [Learning rate of MLP must be increased twentyfold] We found that increasing the MLP learning rate twentyfold to 0.006 resulted in action variation comparable to a CE learning rate of 0.0003. We will revise this statement in the updated paper to avoid confusion.

• [The green and blue curve] This was a typo, it should actually refer to the orange and purple curves.

[1] Hafner et al. "Dream to control: Learning behaviors by latent imagination." arXiv preprint 2019.

[2] Hafner et al. "Mastering atari with discrete world models." arXiv preprint 2020.

[3] Hafner et al. "Mastering diverse domains through world models." arXiv preprint 2023.

We hope that the above response can address your concerns adequately. We would greatly appreciate it if you could re-evaluate our paper based on the above responses.

We sincerely appreciate the time and effort you have invested in reviewing our paper, as well as your insightful and constructive feedback. Your comments have greatly assisted us in improving the quality of our work.

• [Skipping CE blocks] Thank you for pointing out the missing baseline. We implemented this baseline by setting the CE outputs to zero while keeping everything else unchanged to isolate the effect of the latent space. The results are shown as the NO_CE curve in Fig. R1. The results indicate that historical information is crucial in our algorithm. Without it, achieving high returns on these POMDP tasks is hard.

• [Algorithmic novelty] Our work primarily focuses on the impact of learning rates in recurrent RL. We believe the importance and ubiquity of this issue merit a dedicated discussion.

• [Originality of Proposition 1] We referenced the proof from model imitation learning [1]. We consider the RNN as the transition function of an MDP and then, based on the derivation in [1], combined with the characteristics of RNNs, we derived the compounding error of RNNs. To the best of our knowledge, we are the first to derive the compounding error bound in RNNs. Thank you for your suggestion, we will highlight this point and include the relevant citations in the revised paper.

• [Relationship between Proposition 1 and reducing learning rate] Lowering the learning rate reduces the single-step network output variation $\epsilon$, thereby lowering the upper bound of Eq. (1).

• [Last-step observation] Here we use an MLP-based pre-encoder, which is not the MLP policy, to map the last-step observation to a latent space. The MLP policy does not take the last-step observation as input.

• [MLP architecture] We adopted the MLP network design from algorithms like SAC [2] and TD7 [3], which use a two-layer MLP for the policy network, with each layer consisting of 256 neurons. We previously tried smaller network structures, we find the impact was not significant.

• [Claims on MLP] A smaller learning rate makes MLP training very slow, requiring many gradient updates to train the MLP adequately, leading to inefficiency. We will add this explanation to the revised paper.

• [Random seed]

  • We set a random seed at the start of each experiment to randomly initialize everything, including neural network weights and the environment. Every 1000 gradient updates, we perform a policy evaluation without resetting the environment's random state (forking the current random state). Using the current policy, we deterministically collect five trajectories for evaluation, where the policy outputs mean actions, but the initial state of the environment is re-sampled at the start of each trajectory. We repeated the experiment six times, each with a different random seed.
  • To verify that our choice of seed count does not introduce larger statistical errors, we conducted an additional experiment with six more seeds on WalkerBLT-V-v0 and HopperBLT-V-v0. The results, shown in Fig. R5, demonstrate that the new curve (seed_7-12) closely matches the mean and shadow of the old curve (seed_1-6), indicating no significant statistical deviation. This shows that six seeds are sufficient for our experimental setup.

• [Figure 3] Regarding Figure 3, we trained the policy using RESeL and collected trajectories with non-deterministic actions in the environment. We obtained the mean actions $\{a_0,a_1,a_2,...\}$ for each trajectory, where $1,2,3,...$ are rollout steps. We then performed a single gradient update on the policy and obtained the updated policy's mean action outputs at each state in the trajectory $\{a_0', a_1', a_2', \ldots\}$. Figure 3 demonstrates $\{||a_0-a_0'||_2,||a_1-a_1'||_2,||a_2-a_2'||_2,...\}$ for different learning rates, with each experiment repeated 20 times. Although each curve has certain variance, the high repetition count results in a small standard error. The brown curves have a larger variance, hence the noticeable shadow (standard error), while other curves have smaller variance, making the shadow less apparent but present. In the gray dashed curve, the excessively high learning rate caused unpredictable behaviors, leading to significant initial variations in the policy. The gray dashed curve starts high and then decreases, possibly due to the larger action amplitude in the latter half of the trajectory, reaching the boundary values. As a result, many actions reach the boundary values and get compressed after substantial policy updates, leading to less variation compared to the first half.

• [SAC-GRU and SAC-MLP] We believe the too large output variations of RNNs cause the value and policy learning process being instable, therefore preventing the policy from efficiently utilizing the historical information.

• [Minors] Thank you for pointing these out. We will correct all these issues in the revised paper.

[1] Xu et al. "Error bounds of imitating policies and environments." NeurIPS 2021.

[2] Haarnoja et al. "Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor." ICML 2018.

[3] Fujimoto et al. "For sale: State-action representation learning for deep reinforcement learning." NeurIPS 2023.

Thank you for guiding us towards making it a better work. We hope our responses have addressed your concerns effectively and enhanced the clarity of our main contributions. We look forward to your further comments.

Thank you very much for your prompt response and kind comments. We are delighted that our reply has addressed your concerns. Your valuable suggestions have greatly improved our work.