Can Learned Optimization Make Reinforcement Learning Less Difficult?

While reinforcement learning (RL) holds great potential for decision making in the real world, it suffers from a number of unique difficulties which often need specific consideration. In particular: it is highly non-stationary; suffers from high degrees of plasticity loss; and requires exploration to prevent premature convergence to local optima and maximize return. In this paper, we consider whether learned optimization can help overcome these problems. Our method, Learned Optimization for Plasticity, Exploration and Non-stationarity (OPEN), meta-learns an update rule whose input features and output structure are informed by previously proposed solutions to these difficulties. We show that our parameterization is flexible enough to enable meta-learning in diverse learning contexts, including the ability to use stochasticity for exploration. Our experiments demonstrate that when meta-trained on single and small sets of environments, OPEN outperforms or equals traditionally used optimizers. Furthermore, OPEN shows strong generalization characteristics across a range of environments and agent architectures.

Paper

References (90)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer qjkS7/10 · confidence 4/52024-07-06

Summary

This paper presents OPEN, a meta-learning approach with a learned optimiser. The approach is based on addressing three challenges for RL: 1) the non-stationary nature of input and output distributions; 2) the loss of plasticity; and 3) exploration. The approach performs better than a traditional Adam optimiser and other learned optimisation techniques (Optim4RL, L2O, VeLO, and Lion), and works for both single- and multi-task environments.

Strengths

The paper is clearly written and clearly highlights (and makes a good case for mitigating) general problems inherent in the RL framework, namely non-stationarity, plasticity loss, and exploration. The background and definitions are clear. The technique uses per-parameter learned parameter noise to both increase exploration and reduce loss of plasticity. The optimiser update rules are learned based on evolutionary strategies taking dormancy, layer proportion, training proportion, and batch proportion as inputs, which is a novel approach as far as I am aware. The approach works well compared to state-of-the art and provides a scalable approach. Using the approach, the exploration parameter is increased automatically for large environments, which is intuitively what you would expect. The study takes care of ablations to verify the impact of each of the parameters in the meta-optimisation.

Weaknesses

The techniques used for the empirical comparisons are not described in enough detail. The review of directly related techniques is limited, even though I am aware that this is a niche domain. One of the key components is parameter noise. It is not clear how the paper compares to approaches for parameter noise, e.g. citation [11] in your paper. It may for instance be possible that parameter noise alone is beneficial. This is important to verify but I currently do not see any such analysis in the paper. The paper tries to address different points at the same time and there are many components. In a sense, the contribution of each individual part is limited and the technique is a "combination" of techniques rather than a new technique altogether.

Questions

l.75 is use --> is used The automated exploration rate tuning based on the environment and meta-learning of algorithmic steps is similar to Active Adaptive Perception [1]. It is relevant to review it, e.g. in the “Discovering RL algorithms” subsection. Appendix E. Why is Ant being evaluated for more time steps than the other conditions? It would be good to see an ablation comparing your system with citation [11] or other techniques for parameter exploration. It is not clear whether there is an algorithmic improvement in terms of the parameter exploration rule itself. ------ [1] D.M. Bossens, N.C. Townsend, A.J. Sobey (2019). Learning to learn with active adaptive perception, Neural Networks, 115, 2019, Pages 30-49.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The paper describes limitations in terms of bias towards environments where Adam underperforms, the limits of the scope of the study in terms of meta-optimisation (although I think the scope of meta-optimisation is already large), and the base-learner used. In addition, I think there is an additional limitation that should be mentioned. The study focused on gridworlds. It would be interesting to see how the technique scales to high-dimensional problems.

Reviewer qjkS2024-08-09

Thanks for the response, I still have a few comments. The new results look promising at first glance but we would need to see the spread across runs to indicate the confidence in the result. The authors mention the standard error but either there is a mistake or it is too small. Perhaps the standard deviation makes more sense if it is too small to see. > Directly swapping out learnable stochasticity for parameter space noise is, unfortunately, infeasible. Parameter space noise for exploration introduces two new hyperparameters: how much noise to add; and the threshold when noise scales. To tune these while learning to optimize would be excessively costly, as each hyperparameter run would take on the order of days. This was not the question. The question was how your approach compares to RL with parameter noise, not meta-RL with parameter noise. > However, our analysis (Appendix F.4, F.5) shows that OPEN learns stochastic behaviour that could not be represented with vanilla parameter space noise, like producing higher levels of stochasticity in larger Deep Sea environments without awareness of the environment size. It is important to note that our learnable stochasticity conditions on all of OPEN's inputs, meaning there are complex interactions between all of our features and the stochasticity. Therefore, it is unlikely that vanilla parameter space noise could replicate our performance. Actually, the claim does not necessarily follow from these plots. It could simply be the case that (due to dividing by the parameter value) the noise remains constant but the parameter value is changing all the time. The larger environment size might be correlated with smaller parameter values. There is also no comparable plot for the vanilla parameter space noise case. > We believe our work offers significant contribution as the first to show learned optimization consistently outperforming baselines in RL. While we refer you to the expanded list in our global response, our key contributions include: • Learning to optimize in RL using evolution strategies, which has only recently made feasible thanks to JAX acceleration. • Producing a small number of input features, verified via ablation, that improve the performance of the learned optimizer without requiring excessive hand-designed structure. These features are grounded in theory, but are derived novelly. • Introducing meta-learned exploration directly in the optimizer, using learnable stochasticity. • Stabilizing training with zero-meaning. This does not excessively limit the expressiveness of the optimizer but enables meta-training in environments with continuous actions. While OPEN has several components, it is a mischaracterization to reduce it to only a 'combination of techniques'. Our method significantly advances the SOTA in learned optimization for RL. We also note that many successful papers integrate multiple findings into a single method, such as [1]. I think that these points of novelty are not clear from the text, for instance, these points are not highlighted in the intro. Also the novelty and mechanism of each of these points is difficult to expand upon (and therefore evaluate) seriously due to having so many components. Whether or not it advances SOTA performance is not directly related to whether it is a combination of techniques or not. That said, this is not a point of rejection for me and I believe that many methods can be characterised like this – some more easily than others.

Authorsrebuttal2024-08-12

Dear qjkS, Thank you for your response and continued engagement in the rebuttal process. We apologize for the delay in our response; we have spent some time generating additional results. # Craftax Results If you zoom in to the plot, there are standard error intervals, but as in the original Craftax paper these are very small. For your benefit, we provide the specific values for final return below: | Optimizer | Mean ± Std Err | |---------------------|-------------------| | Adam (Finetuned) | 13.842 ± 0.037 | | OPEN (0-shot) | 12.066 ± 0.036 | | Adam (0-shot) | 7.490 ± 0.053 | # Additional Baselines We apologize for misunderstanding your original review point. We have now implemented Noisy Networks for Exploration [1] with PPO on the MinAtar environments. This allows for a *learned per-parameter variance* which is comparable to the noise applied by OPEN, rather than parameter space noise (where noise variance is shared across the whole agent). We try to follow the implementation of noisy networks for A3C, though highlight a few important implementation details and differences below: - We maintain a constant noise *between updates* rather than over episodes, as would be standard; in JAX, with vectorized environment rollouts, it is far from trivial to have constant episodic noise unless episodes are of a constant length (which is not the case in MinAtar). OPEN also receives new noise every update rather than episodically. - We tune Adam (*LR*, $\beta_1$, $\beta_2$) for each MinAtar environment. - Unlike OPEN, where noise is *sticky* (i.e., maintained across updates), in noisy networks the noise gets resampled regularly. - In OPEN, the learnable stochasticity is applied *in addition* to the entropy bonus, whereas in noisy networks the noise is applied *instead of* entropy. In our noisy networks implementation, we remove the entropy bonus from PPO to align to literature. - We use the standard initialization of variance in noisy networks (0.017). Below, we provide results for PPO + noisy networks alongside the scores for Adam and OPEN. We report the IQM and 95% stratified bootstrap confidence intervals (min, max) for 16 seeds, following the reporting procedure from the paper. | Optimizer | Freeway | Breakout | Asterix | Space Invaders | |-------------|--------------------------|--------------------------|--------------------------|--------------------------| | OPEN | **64.42 (64.06, 64.74)** | **66.94 (60.60, 72.02)** | **36.84 (32.11, 39.96)** | **167.22 (157.47, 175.76)** | | Adam | 62.26 (61.93, 62.58) | 47.82 (37.44, 55.81) | 14.97 (13.30, 16.42) | **165.42 (162.71, 168.52)** | | Noisy Networks (+Adam) | 54.02 (53.47, 54.59) | 38.09 (30.36, 48.47) | 9.75 ( 5.46, 14.38) | 124.86 (121.58, 129.58) | We note that applying noisy networks, or any form of parameter noise, is uncommon with PPO; as a result, the fact that the more canonical implementation of PPO outperforms the noisy networks should not prove surprising. It is, however, interesting that learnable stochasticity provides such a benefit to OPEN; as we highlight in our paper, it may be that improvements arise from **both** improved exploration *and* reduced dormancy. It is also possible that it is significantly easier to *meta-learn* a noise profile than learn how much noise to apply *online* (as in noisy networks, where variance is an optimization parameter learned by RL). # Novelty We appreciate your comments around this, and have included some sentences describing the contributions of our method into the introduction. To evaluate each contribution, we provide thorough ablation of each of these components and find that each contributes to the overall performance of OPEN. The novelty of these contributions should not be difficult to interpret; no one has trained optimizers in RL with ES, our stabilization method has not been applied before, and while our inspiration for the various features and randomness is included in the paper, their incorporation into a meta-learned system is novel besides temporal conditioning, which was inspired by [2] (though we implement it slightly differently). In our original paper, we attempted to make this clear, but our new sentences in the introduction make this more concrete. We hope we have answered your concerns, but please do send a comment if you have additional questions that we can answer before the deadline. --- [1] Meire Fortunato et al., Noisy Networks for Exploration, 2017 [2] Matthew Thomas Jackson et al., Discovering Temporally-Aware Reinforcement Learning Algorithms, 2024

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

Summary

This paper presents OPEN, a meta-learned optimizer for reinforcement learning gradient updates. The update rule for OPEN is split into 3 stages, one for each difficulty in RL they define: non-statiarity, plasticity loss and exploration. They train their optimizer with an evolutionary strategy. To analyze their results they run experiments on different setups with varying training and deployment domains. They also perform an ablation study to verify their choises. They find that their method significantly outperforms their baselines (mainly Adam), and is able to generalize well to different environments.

Strengths

Quality: Major. The paper is sound tehcnically and their ablation study/experimental results are very well done. They are very thorough in their testing. Clarity: Major. The paper usees bolding and italics exellently to provide roadmaps and make it easy to read. The ablations study plots are extremely easy to interpret. Information from the text and appendix should be enough to reproduce all results. Significance: Above average. The results seem to be promising to improving RL algorithm performance only by improving the optimizer. Other than that, on the tasks they do experiment with their results show a large improvement over all baselines and seem to provide a step toward better RL performance. Originality: Slightly below average. The high level idea here (learning an optimizer for RL) is not a novel idea, being done in the Optim4RL work. I don't think that this high level novelty is the main contribution of the paper and they don't claim it to be. They do take a different approach to learning the optimizer which is their novelty. Works are accurately cited.

Weaknesses

I wish it was a little more clear how this work is novel (from the authors perspective). From my reading I gathered that it is the use of ES and specific identification/use of RL-specific challenges but a sentence or two in the intro would be fantastic. They only apply their method to gridworld tasks, while this is completely reasonable due to possibly limited compute I would like to know if this generalizes to more complex tasks, or at least if the authors believe it will.

Questions

If you feel your work is more novel than I am giving credit for this is something I am very willing to change my mind on if I'm wrong. Your optimizer seems to generalize well, is it possible to train it once on a wide variety of environments and then freeze it and use it for basically any environment? You address this in the appendix, but it would be nice to put in the main text a small sentence or two about how your method isn't much slower than Adam. That was a concern of mine as far as applying this if it is magnitudes slower than Adam that cancels out some benefit but it doesn't seem like it is.

Rating

8

Confidence

3

Soundness

4

Presentation

4

Contribution

4

Limitations

Limitations: Maybe this is implied somehwere or I missed it, but applying this to more complex domains would be an interesting future work in my opinion. If you feel that your environments generalize to complex domains then this is not a limitation. None other than that.

Reviewer Fr377/10 · confidence 4/52024-07-18

Summary

This paper explores whether learned optimization can address specific challenges in reinforcement learning (RL), such as non-stationarity, plasticity loss, and the need for exploration. The authors propose an optimizer, OPEN, which meta-learns update rules related to these issues. OPEN demonstrates strong performance and generalization across various environments compared to traditional and other learned optimizers. The study confirms that learned optimizers, when well-parameterized and trained, can improve the PPO algorithms' effectiveness.

Strengths

- Very clear and well-written paper. - Comprehensive Introduction and Related Work. - Strong evaluation using Adam-Normalized scores. - Extensive ablation studies.

Weaknesses

- Line 135: Plasticity loss. The use of ReLU dormancy is, in my opinion, not a viable metric for plasticity in RL. Although interesting work, [1] has shown that combatting ReLU dormancy does not improve 'vanilla' nature DQN, unless you increase the replay ratio. Further work [2] has also shown that combatting dormant neurons as in [1] significantly reduces Impala DQN performance, while extensive network pruning increases performance. This raises the question of whether neuron dormancy is actually related to plasticity in RL. - An evaluation on another RL algorithm or a more complex environment (Atari?) would strengthen the results. (The PureJaxRL Github Repo [3] also has DQN) Small remarks ##: - Line 73: "Evolution algorithms" -> 'Evolutionary algorithms' - Line 75: "This population is use" -> 'This population is used' - Line 250: "significantly outperform" -> 'significantly outperforms' [1] Sokar et al: The Dormant Neuron Phenomenon in Deep Reinforcement Learning [2] Obando-Ceron et al: In value-based deep reinforcement learning, a pruned network is a good network [3] Lu et al: Discovered policy optimisation

Questions

Can the authors say something about the computational cost of their method versus the other methods used in their evaluation?

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

Related to my question, maybe a small section about computational differences should be added to the main text.

Reviewer Fr372024-08-08

Response to Rebuttal

Dear authors, Thank you for your extensive rebuttal. # Dormancy I agree with your conclusion that dormancy is very likely related to performance. However, I tend to see dead neurons more as a symptom of poor optimization, rather than having a causal relationship to plasticity. I think your use of dormancy as an OPEN feature allows it to treat the disease rather than treat the symptoms. But I might be wrong :). # Harder Environments Thanks for adding the Craftax environment. Even though the consensus already seems to lead to acceptance, this will certainly strengthen the paper. # Computational cost A more visible reference to the computational cost appendix will be useful, thanks. Given that the authors have done an extensive rebuttal, and are committed to adding an additional environment in their evaluation, I have increased the score to a 7.

Authorsrebuttal2024-08-12

Dear Fr37, Thank you for your detailed review and engagement in our rebuttal; we are very grateful that you have increased your score. Since our experiment has concluded, we now share our final results comparing OPEN against Adam for PQN [1], a parallelised and simplified implementation of DQN designed for JAX. Here, we slotted OPEN *directly* into the PQN code available on github, with no algorithmic changes made to OPEN. Due to compute limitations, we were only able to train OPEN from a single seed in one environment, Asterix, though we see no reason that these results would not transfer to other environments. As was standard in our work, we used many hyperparameters directly from the PQN paper but extensively tuned Adam (*LR*, $\beta_1$, $\beta_2$, *Anneal_LR*) with 8 seeds per hyperparameter. We are happy to give full hyperparameters if desired. Our results are as follows, assessed over 64 seeds with standard error. OPEN significantly outperforms Adam. | Optimizer | Final Return ± Std Err | |---------------|-------------------| | Adam | 36.86 ± 0.68 | | OPEN | **40.06 ± 0.42** | While I am unable to show you our return curve due to the limit of a 1 page PDF submitted during the rebuttal period only, we also observe some interesting behaviour similar to that of OPEN with PPO; OPEN only overtakes Adam later in training. This suggests that many of our findings about the behaviour of OPEN with PPO also hold with other algorithms. However, further analysis would be required to verify this, which we have not had the time to carry out during the rebuttal period. We hope that this satisfies your concerns regarding other algorithms. We will ensure to update our paper with these additional results. [1] Matteo Gallici et al., Simplifying Deep Temporal Difference Learning. 2024.

Reviewer sKXo7/10 · confidence 4/52024-07-25

Summary

This paper presents OPEN, which meta-learns an update rule (a learned optimizer) for reinforcement learning. The authors test in single-task and multi-task settings, and show good generalization to tasks outside of the training set.

Strengths

This is a good paper on the application of learned optimization to RL. Of particular note: - I like the focus on learning capacity via single-task learning. This is a good sanity check for learned optimizers. - The out-of-distribution experiments are valuable. - Discussion of acceleration via parallelization in JAX is an extremely important engineering detail. - The explicit conditioning on dormancy is an interesting contribution. There is interesting potential future work within the learned optimization community on useful theoretically-driven inputs to learned optimizers.

Weaknesses

- The zero mean last step (eq (9)) seems surprising for an update. Indeed, it seems quite different compared to eg momentum or other normalization strategies (such as adam-style normalization). Better ablations on this term would be helpful. It makes sense that adding random noise at each timestep could destabilize the model (the variance as t -> \infty will grow without limit), but could you instead use an additive noise process with a contraction term that has bounded or shrinking variance instead? - The notation in eq (8) is a bit hard to parse---it reads more like an assignment in CS than an equality. Even though a def eq is used here, I think the notation could be improved for clarity.

Questions

- How does the OpenAI ES strategy compare to the ES strategies used in the previous work on learned optimization, like standard ES, PES, NRES, etc? - Are there alternate methods for preventing parameter growth from the exploration term?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

None

Reviewer KVk82024-08-08

Thank you for a response to my questions. The idea to possibly train on a variety of environments and universally apply the optimizer is one that I believe has far reaching implications. Great job.

Authorsrebuttal2024-08-12

Dear reviewers, As promised in our general response, our experiment running OPEN with PQN [1] has now concluded and we can share the results. PQN is a recent version of DQN built around using vectorized sampling in JAX. Due to both compute and time constraints, we train OPEN for one seed on Asterix and evaluate against Adam. We use most hyperparameters directly from PQN, but completed an extensive finetuning on Adam covering *LR*, $\beta_1$, $\beta_2$ and *Anneal_LR*. We evaluate each hyperparameter configuration on 8 seeds. We run our evaluation of OPEN and Adam for 64 seeds, and provide standard error. Our results are as follows: | Optimizer | Final Return ± Std Err | |---------------|-------------------| | Adam | 36.86 ± 0.68 | | OPEN | **40.06 ± 0.42** | OPEN significantly outperforms Adam in this setting. While we were unable to explore different environments to Asterix, we also see no reason why these results would not transfer to the other environments explored in our paper. We will update our manuscript to include these new results. Though we can not share any additional figures due the 1 page PDF constraint, we also want to highlight that there seems to be interesting behavior of OPEN with PQN which parallels that of OPEN with PPO. In particular, performance is not uniformly better for OPEN over the training horizon, and OPEN only begins to outperform Adam from about halfway through training. We will include the return curve into our appendix, like with all our other experiments, to demonstrate the occurence of this effect. While additional analysis would be interesting in this setting, to see if the behavior of OPEN is consistent across algorithms, we have not had the time within the rebuttal period to carry this out. We hope that the reviewers find this additional result interesting, and are happy to answer any extra queries about the performance of OPEN with other algorithms. [1] Matteo Gallici et al., Simplifying Deep Temporal Difference Learning. 2024.

Program Chairsdecision2024-09-25

Decision

Accept (spotlight)

© 2026 NYSGPT2525 LLC