Can Learned Optimization Make Reinforcement Learning Less Difficult?

Authors: alexander d goldie
chris  lu
matthew t jackson
shimon  whiteson
jakob n foerster
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 11:37:02

Abstract: 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 **O**ptimization for **P**lasticity, **E**xploration and **N**on-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.

Main paper: 
 introduction: 
 Reinforcement learning [1, RL] has undergone significant advances in recent years, scaling from solving complex games [2, 3] towards approaching real world applications [4] [5] [6] [7] . However, RL is limited by a number of difficulties which are not present in other machine learning domains, requiring the development of numerous hand-crafted workarounds to maximize its performance.Here, we take inspiration from three difficulties of RL: non-stationarity due to continuously changing input and output distributions [8] ; high degrees of plasticity loss limiting model capacities [9, 10] ; and exploration, which is needed to ensure an agent does not converge to local optima prematurely [1, 11] . Overcoming these challenges could enable drastic improvements in the performance of RL, potentially reducing the barriers to applications of RL in the real-world. Thus far, approaches to tackle these problems have relied on human intuition to find hand-crafted solutions. However, this is fundamentally limited by human understanding. Meta-RL [12] offers an alternative in which RL algorithms themselves are learned from data rather than designed by hand. Meta-learned RL algorithms have previously demonstrated improved performance over hand-crafted ones [13] [14] [15] .On a related note, learned optimization has proven successful in supervised and unsupervised learning (e.g., VeLO [16] ). Learned optimizers are generally parameterized update rules trained to outperform handcrafted algorithms like gradient descent. However, current learned optimizers perform poorly in RL [16, 17] . While some may argue that RL is simply an out-of-distribution task [16] , the lack of consideration for, and inability to target, specific difficulties of RL in naïve learned optimizers leads us to believe that they would still underperform even if RL were in the training distribution.We propose an algorithm, shown in Figure 1 , for meta-learning optimizers specifically for RL called Learned Optimization for Plasticity, Exploration and Nonstationarity (OPEN). OPEN meta-learns update rules whose inputs are rooted in previously proposed solutions to the difficulties above, and whose outputs use learnable stochasticity to boost exploration. OPEN is trained to maximize final return only, meaning it does not use regularization to mitigate the difficulties it is designed around.In our results (Section 6), we benchmark OPEN against handcrafted optimizers (Adam [18] , RMSProp [19] ), open-source discovered/learned optimizers (Lion [20] , VeLO [16] ) and baseline learned optimizers trained for RL (Optim4RL [17] , 'No Features'). We show that OPEN can fit to single and small sets of environments and also generalize in-and out-of-support of its training distribution. We further find that OPEN generalizes better than Adam [18] to gridworlds and Craftax-Classic [21] . Therefore, our novel approach of increasing optimizer flexibility via extra inputs and an explorationdriven output, rather than enforcing more structured updates, enables significant performance gains.Figure 1 : A visualization of OPEN. We train N agents, replacing the handcrafted optimizer of the RL loop with ones sampled from the meta-learner (i.e., evolution). Each optimizer conditions on gradient, momentum and additional inputs, detailed in Section 5.3, to calculate updates. The final returns from each loop are output to the meta learner, which improves the optimizer before repeating the process. A single inner loop step is described algorithmically in Appendix B.1. 
background: 
 Reinforcement Learning The RL problem is formulated as a Markov Decision Process [1, MDP] described by the tuple ⟨A, S, P, R, ρ, γ⟩. At discrete timestep t, the agent takes action a t ∈ A sampled from its (possibly stochastic) policy, π(•|s t ) ∈ Π, which conditions on the current state s t ∈ S (where s 0 ∼ ρ). After each action, the agent receives reward R(s t , a t ) and the state transitions to s t+1 , based on the transition dynamics P (s t+1 |s t , a t ). An agent's objective is to maximize its discounted expected return, J π , corresponding to a discount factor γ ∈ [0, 1), which prevents the agent making myopic decisions. This is defined asEQUATIONProximal policy optimization [22, PPO] is an algorithm designed to maximize J π . It uses advantage, A π (s, a), which is calculated from the state value function,V π = E π [ ∞ t=0 γ t R t |S t = s],and stateaction value function,Q π (s, a) = E π [ ∞ t=0 γ t R t |S t = s, A t = a].It measures the improvement of a specific action over the current policy and takes the form A π (s, a) = Q π (s, a) -V π (s).(2) PPO introduces a loss for optimizing the policy, parameterized by θ, that prevents extreme policy updates in gradient ascent. This uses clipping, which ensures there is no benefit to updating θ beyond where the policy probability ratio, r t (θ) = π θ (at|st) π θ old (at|st) , exceeds the range [1 ± ϵ]. The clipped loss is L CLIP (θ) = E [min(r t (θ)A π (s t , a t ), clip(r t (θ), 1 ± ϵ)A π (s t , a t ))] .(3) PPO is an actor-critic [1] method, where the policy and value functions are modeled with different neural networks, or separate heads of the same neural network, conditioning on state. The PPO objective to maximize combines the clipped loss, a value function error, and an entropy bonus intoEQUATIONMeta-Learning Optimizers in RL Algorithms like Adam [18] or RMSprop [19] are designed to maximize an objective by updating the parameters of a neural network in the direction of positive gradient with respect to the function. They are often applied with augmentations such as momentum or learning rate schedules to better converge to optima. Learned optimizers offer an alternative: parameterized update rules, conditioning on more than just gradient, which are trained to maximize an objective [23, 16, 24] . For any parameterized optimizer opt ϕ , which conditions on a set of inputs x, the update rule to produce update u can be described as a function, opt ϕ (x) → u.We treat learning to optimize as a meta-RL problem [12] . In meta-RL, the goal is to maximize J π over a distribution of MDPs P (M). For our task, an optimizer trained to maximize J (ϕ) = E M J π |opt ϕ yields the optimal meta-parameterization opt ϕ * .Evolution Strategies Evolution algorithms (EA) are a backpropagation-free, black-box method for optimization [25] which uses a population of perturbed parameters sampled from a distribution (here, θ ∼ N (θ, σ 2 I)). This population is used to maximize a fitness F (•). EA encompasses a range of techniques (e.g., evolution strategies (ES) [26, 27] , genetic algorithms [28] or CMA-ES [29] ).Natural evolution strategies (NES) [30] are a class of ES methods that use the population fitness to estimate a natural gradient for the mean parameters, θ. This can then be optimized with typical gradient ascent algorithms like Adam [18] . Salimans et al. [26] introduce OpenAI ES for optimizing θ using the estimatorEQUATIONwhich is approximated using a population average. In practice, we use antithetic sampling (i.e., for each sampled ϵ, evaluating +ϵ and -ϵ) [31] . Antithetic task sampling enables learning on a task distribution, by evaluating and ranking each antithetic pair on different tasks [32] .Historically, RL was too slow for ES to be practical for meta-training. However, PureJaxRL [15] recently demonstrated the feasibility of ES for meta-RL, owing to the speedup enabled by vectorization in Jax [33] . We use the implementation of OpenAI ES [26] from evosax [34] .3 Related WorkOur full list of features is concatenated and fed into the optimizer, as in algorithm 1 below. We use some notation from Table 3 . These features are:• Gradient g, processed with the transformation g → {sign(g), log(g + ϵ)}.• Momentum m calculated with the following β coefficients: [0.1, 0.5, 0.9, 0.99, 0.999, 0.9999].This is processed with the transformation m i → {sign(m i ), log(m i + ϵ)} for each β i .• Current value of the parameter p.• Training proportion, or how far the current batch is through the total number of batches.Since the total number of batches is calculated as N batch = (T //N steps )//N envs , this is defined as (t//(L * N minibatches )/N batch , where t is the current update iteration. We use a floor division operator as training proportion should be constant through a batch.• Batch proportion, or how far the current update is through training with the same batch, calculated as ((t//N minibatches ) mod L)/L.• Dormancy, D, calculated for the neuron downstream of the parameter using equation 6. This is repeated over parameters with the same downstream neuron. • Proportion of the current layer through the network, h/H. The gradient, momentum and parameter value features are normalized to have a second moment of 1 across the tensor, following Metz et al. [90, 45] , as a normalization strategy which preserves direction. Like Lan et al. [17] , we process gradients and momentum as x → {sign(x), log(x + ϵ)} to magnify the difference between small values. The details of the layer sizes of OPEN in both the single-task and gridworld settings are in Table 1 . In our update rule, we let α 1 = α 2 = α 3 = 0.001. We mask out the third output in the critic by setting the noise (ϵ) to zero; otherwise, updates to our critic would be non-deterministic.Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: We discuss limitations of this work thus far with respect to curricula used in multi-task settings, and discuss potential avenues for future work in this regard or in terms of considering other areas of RL which this method could be applied to. This is included in Section 8. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
optimization in rl: 
 Wilson et al. [35] and Reddi et al. [36] show that adaptive optimizers struggle in highly stochastic processes. Henderson et al. [37] indicate that, unlike other learning regimes, RL is sufficiently stochastic for these findings to apply, which suggests RL-specific optimizers could be beneficial.The idea that optimizers designed for supervised learning may not perfectly transfer to RL is reflected by Bengio et al. [38] , who propose an amended momentum suitable for temporal difference learning [39] . This is related to work by Sarigül and Avci [40] , who explore the impact of different types of momentum on RL. While these works motivate designing optimization techniques specifically for RL, we take a more expressive approach by replacing the whole optimizer instead of just its momentum calculation, and using meta-learning to fit our optimizer to data rather than relying on potentially suboptimal human intuition.Discovering RL Algorithms Rather than using handcrafted algorithms, a recent objective in meta-RL is discovering RL algorithms from data. While there are many successes in this area (e.g., Learned Policy Gradient [13, LPG] , MetaGenRL [14] , Active Adaptive Perception [41] and Learned Policy Optimisation [15, LPO] ), we focus on meta-learning a replacement to the optimizer due to the outsized impact a learned update rule can have on learning. We also use specific difficulties of RL to guide the design of our method, rather than simply applying end-to-end learning.Jackson et al. [32] learn temporally-aware versions of LPO and LPG. While their approach offers inspiration for dealing with non-stationarity in RL, they also rely on Adam [18] , an optimizer designed for stationarity that is suboptimal for RL [9] . Instead, we propose replacing the optimizer itself with an expressive and dynamic update rule that is not subject to these problems.Learned Optimization Learning to optimize [23, 42, 43, L2O] strives to learn replacements to handcrafted gradient-based optimizers like Adam [18] , generally using neural networks (e.g., [24, [44] [45] [46] ). While they show strong performance in supervised and unsupervised learning [16, 47] , previous learned optimizers do not consider the innate difficulties of RL, limiting transferability. Furthermore, while VeLO used 4000 TPU-months of compute [16] , OPEN requires on the order of one GPU-month.Lion [20] is a symbolic optimizer discovered by evolution that outperforms AdamW [48] using a similar update expression. While the simplistic, symbolic form of Lion lets it generalize to RL better than most learned optimizers, it cannot condition on features additional to gradient and parameter value. This limits its expressibility, ability to target the difficulties of RL and, therefore, performance.Learned Optimization in RL Learned optimization in RL is significantly more difficult than in other learning domains due to the problems outlined in section 4. For this reason, SOTA learned optimizers fail in transfer to RL [16] . Optim4RL [17] attempts to solve this issue by learning to optimize directly in RL. However, in their tests and ours, Optim4RL fails to consistently beat handcrafted benchmarks. Also, it relies on a heavily constrained update expression based on Adam [18] , and it needs expensive learning rate tuning. Instead, we achieve much stronger results with a completely black-box setup inspired by preexisting methods for mitigating specific difficulties in RL.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [NA] Justification: We do not make theoretical claims as this is a purely empirical paper. This is made clear in the paper. Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
difficulties in rl: 
 We believe that fundamental differences exist between RL and other learning paradigms which make RL particularly difficult. Here, we briefly cover a specific set of prominent difficulties in RL, which are detailed with additional references in appendix A. Our method takes inspiration from handcrafted heuristics targeting these challenges (Section 5.3). We show via thorough ablation (Section 7) that explicitly formulating our method around these difficulties leads to significant performance gains.(Problem 1) Non-stationarity RL is subject to non-stationarity over the training process [1] as the updating agent causes changes to the training distribution. We denote this training non-stationarity. Lyle et al. [9] suggest optimizers designed for stationary settings struggle under nonstationarity.(Problem 2) Plasticity loss Plasticity loss, or the inability to fit new objectives during training, has been a theme in recent deep RL literature [9, 49, 50, 10] . Here, we focus on dormancy [49] , a measurement tracking inactive neurons used as a metric for plasticity loss [10, [51] [52] [53] . It is defined asEQUATIONwhere h l i (x) is the activation of neuron i in layer l with input x ∈ D for distribution D. H l is the total number of neurons in layer l. The denominator normalizes average dormancy to 1 in each layer.A neuron is τ -dormant if s l i ≤ τ , meaning the neuron's output makes up less than τ of its layer's output. For ReLU activation functions, τ = 0 means a neuron is in the saturated part of the ReLU. Sokar et al. [49] find that dormant neurons generally stay dormant throughout training, motivating approaches which try to reactivate dormant neurons to boost plasticity.(Problem 3) Exploration Exploration is a key problem in RL [1] . To prevent premature convergence to local optima, and thus maximize final return, an agent must explore uncertain states and actions. Here, we focus on parameter space noise for exploration [11] , where noise is applied to the parameters of the agent rather than to its output actions, like ϵ-greedy [1] .Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes] Justification: We give full details of our method, including mathematical expressions for all inputs to the optimizer (Appendix B), the architecture used (Appendices B.4, B.5) and how it was trained (5.2). This was all detailed in the main body of the paper, in Section 5.We provide a full algorithm in Appendix B.1 which can be used to completely replicate the results of the paper. Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. Justification: For all experiments in the main body of the paper, we report 95% stratified bootstrap confidence intervals, using the implementation of rliable by Agarwal et al. [70] .For return curves in Appendices E and G, in addition to our analysis F, we make it clear that we plot standard error. We discuss that all error bars are for a single learned optimizer, due to the extreme cost of learned optimization, and cover multiple runs with the same optimizer. We discuss that other works use the same evaluation protocol [16, 17, 47] .Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
method: 
 There are three key considerations when doing learned optimization for RL: what architecture to use; how to train the optimizer; and what inputs the optimizer should condition on. In this section, we systematically consider each of these questions to construct OPEN, with justification for each of our decisions grounded in our core difficulties of RL.To enable conditioning on history, which is required to express behavior like momentum, for example, OPEN uses a gated recurrent unit [54, GRU] . This is followed by two fully connected layers with LayerNorm [55] , which we include for stability. All layers are small; this is important for limiting memory usage, since the GRU stores separate states for every agent parameter, and for maintaining computational efficiency [45] . We visualize and detail the architecture of OPEN in Appendix B.We split the calculation of the update in OPEN into three stages, each of which serves a different purpose. The first stage, which follows Metz et al. [24] , isEQUATIONwhere m i and e i are optimizer outputs for parameter i, and α (1,2) are small scaling factors used for stability. This update can cover many orders of magnitude without requiring large network outputs.(P3) Secondly, we augment the update rule for the actor only to increase exploration; there is no need for the critic to explore. We take inspiration from parameter space noise for exploration [11] , since it can easily be applied via the optimizer. To be precise, we augment the actor's update asEQUATIONHere, α 3 is a small, stabilizing scaling factor, δ actor i is a third output from the optimizer, and ϵ ∼ N (0, 1) is sampled Gaussian noise. By multiplying δ actor i and ϵ, we introduce a random walk of learned, per-parameter variance to the update which can be used for exploration. Since δ actor i depends on the optimizer's inputs, this can potentially learn complex interactions between the noise schedule and the features outlined in Section 5.3. Unlike Plappert et al. [11] , who remove and resample noise between rollouts, OPEN adds permanent noise to the parameters. This benefits plasticity loss (i.e., (P2)), in principle enabling the optimizer to reactivate dormant neurons in the absence of gradient.Unfortunately, naïve application of the above updates can cause errors as, even without stochasticity, agent parameters can grow to a point of numerical instability. This causes difficulty in domains with continuous action spaces, where action selection can involve exponentiation to get a non-negative standard deviation. Therefore, the final update stage stabilizes meta-optimization by zero-meaning, asEQUATIONWhile this limits the update's expressiveness, we find that even traditional optimizers tend to produce nearly zero-mean updates. In practice, this enables learning in environments with continuous actions without harming performance for discrete actions. Parameter i is updated as p(t+1) i = p (t)i -u i .We train our optimizers with OpenAI ES [26] , using final return as the fitness. We apply the commonly-used rank transform, which involves mapping rankings over the population to the range [-0.5, 0.5], to the fitnesses before estimating the ES gradient; this is a form of fitness shaping [30, 26] , and makes learning both easier and invariant to reward scale.For training on multiple environments simultaneously (i.e., multi-task training, Section 6), we evaluate every member of the population on all environments. After evaluation, we: 1) Divide by the return Adam achieves in each environment; 2) Average the scores over environments; 3) Do a rank transform.Normalizing by Adam maps returns to a roughly common scale, enabling comparison between diverse environments. However, this biases learning to environments where Adam underperforms OPEN. We believe finding better curricula for multi-task training would be highly impactful future work.Carefully selecting which inputs to condition OPEN on is crucial; they should be sufficiently expressive without significantly increasing the computational cost or meta-learning sample requirements of the optimizer, and should allow the trained optimizers to surgically target specific problems in RL. To satisfy these requirements, we take inspiration from prior work addressing our focal difficulties of RL.In spirit, we distill current methods to a 'lowest common denominator' which is cheap to calculate. We provide additional details of how these features are calculated in Appendix B.2.(P1) Two training timescales Many learned optimizers for stationary problems incorporate some version of progress through training as an input [16, 45, 46] . Since PPO learns from successively collected, stationary batches of data [56] , we condition OPEN on how far it is through updating with the current batch (i.e., batch proportion). This enables behavior like learning rate scheduling, which has proved effective in stationary problems (e.g., [57, 58] ), and bias correction from Adam to account for inaccurate momentum estimates [18] .Inspired by Jackson et al. [32] , who demonstrate the efficacy of learning dynamic versions of LPG [13] and LPO [15] , we also condition OPEN on how far the current batch is through the total number of batches to be collected (i.e., training proportion). This directly targets training non-stationarity.(P2) Layer Proportion Nikishin et al. [59, 60] operate on higher (i.e., closer to the output) network layers in their attempts to address plasticity loss in RL. Furthermore, they treat intervention depth as a hyperparameter. To replicate this, and enable varied behavior between the different layers of an agent's network, we condition OPEN on the relative position of the parameter's layer in the network.(P2) Dormancy As Sokar et al. [49] reinitialize τ -dormant neurons, we condition OPEN directly on dormancy; this enables similar behavior by allowing OPEN to react as neurons become more dormant.In fact, in tandem with learnable stochasticity, this enables OPEN to reinitialize dormant neurons, just like Sokar et al. [49] . Since dormancy is calculated for neurons, rather than parameters, we use the value of dormancy for the neuron downstream of each parameter. 
results: 
 In this section, we benchmark OPEN against a plethora of baselines on large-scale training domains.Due to computational constraints, we meta-train an optimizer on a single random seed without ES hyperparameter tuning. This follows standard evaluation protocols in learned optimization (e.g., [16, 17, 47] ), which are also constrained by the high computational cost of meta-learned optimization.We provide the cost of experiments in Appendix J, including a comparison of runtimes with other optimizers. We detail hyperparameters in Appendix C.5. We define four evaluation domains, based on Kirk et al. [61] , in which an effective learned optimization framework should prove competent:Single-Task Training A learned optimizer must be capable of fitting to a single environment, and being evaluated in the same environment, to demonstrate it is expressive enough to learn an update rule. We test this in five environments: Breakout, Asterix, Space Invaders and Freeway from MinAtar [62, 63] ; and Ant from Brax [64, 65] . This is referred to as 'singleton' training in Kirk et al. [61] .To show an optimizer is able to perform under a wide input distribution, it must be able to learn in a number of environments simultaneously. Therefore, we evaluate performance from training in all four environments from MinAtar [62, 63] 2 . We evaluate the average normalized score across environments with respect to Adam.In-Distribution Task Generalization An optimizer should generalize to unseen tasks within its training distribution. To this end, we train OPEN on a distribution of gridworlds from Jackson et al. [66] with antithetic task sampling [32] , and evaluate performance by sampling tasks from the same distribution. We include details in Appendix D.Out-Of-Support Task Generalization Crucially, an optimizer unable to generalize to new settings has limited real-world usefulness. Therefore, we explore out-of-support (OOS) generalization by testing OPEN on specific gridworld distributions defined by Oh et al. [13] , and a set of mazes from minigrid [67] which are not in the training distribution, with unseen agent parameters. Furthermore, we test how OPEN performs 0-shot in Craftax-Classic [21] , a Jax reimplementation of Crafter [68] .We compare against open-source implementations [69] of Adam [18] , RMSProp [19] , Lion [20] and VeLO [16, 45] . We also learn two optimizers for the singleand multi-task settings: 'No Features', which only conditions on gradient and momentum; and Optim4RL [17] (using ES instead of meta-gradients, as in Lan et al. [17] ). Since Optim4RL is initialized close to sgn(Adam), and tuning its learning rate is too practically expensive, we set a learning rate of 0.1 × LR Adam based on Lion [20] (which moves from AdamW to sgn(AdamW)). The optimizer's weights can be scaled to compensate if this is suboptimal. We primarily consider interquartile mean (IQM) of final return with 95% stratified bootstrap confidence intervals [70] . All hyperparameters can be found in Appendix C.Figure 2 shows the performance of OPEN after single-task training. In three MinAtar environments, OPEN significantly outperform all baselines, far exceeding previous attempts at learned optimization for RL. Additionally, OPEN beat both learned optimizers in every environment. Overall, these experiments show the capability of OPEN to learn highly performant update rules for RL.Return curves (Appendix E) show that OPEN does not always achieve high returns from the start of training. Instead, OPEN can sacrifice greedy, short-term gains in favor of long-term final return, possibly due to its dynamic update rule. We further analyze the optimizers' behavior in Appendix F. We plot 95% stratified bootstrap confidence intervals for each environment.Figure 3 shows each optimizer's ability to fit to multiple MinAtar environments [63, 62] , where handcrafted optimizers are tuned per-environment. We normalize returns with respect to Adam, and aggregate scores over environments to give a single performance metric. Here, we increase the size of the learned optimizers, with details in Appendix B.5. In addition to IQM, we consider the mean normalized return to explore the existence of outliers (which often correspond to asterix, where OPEN strongly outperforms Adam), and optimality gap, a metric from Agarwal et al. [70] measuring how close to optimal algorithms are. Unlike single-task training, where optimizers are trained until convergence, we run multi-task experiments for a fixed number of generations (300) to limit compute.OPEN drastically outperforms every optimizer for multi-task training. In particular, OPEN produces the only optimizer with an aggregate score higher than Adam, demonstrating its ability to learn highly expressive update rules which can fit to a range of contexts; no other learned optimizers get close to OPEN in fitting to multiple environments simultaneously. As expected, OPEN prioritizes optimization in asterix and breakout, according to the return curves (Appendix G); we believe better curricula would help to overcome this issue. Interestingly, Optim4RL seems to perform better in the multi-task setting than the single-task setting. This may be due to the increased number of meta-training samples. Figure 3 : Mean, IQM and optimality gap (smaller = better), evaluated over 16 random seeds per environment for the aggregated, Adam-normalized final returns after multi-task training on MinAtar [62, 63] . We plot 95% stratified bootstrap confidence intervals for each metric.Figure 4 shows OPEN's ability to generalize in a gridworld setting. We explore in-distribution generalization on the left, and OOS generalization on the right, where the top row (rand_dense, rand_sparse and rand_long) are from LPG [13, 66] and the bottom row are 3 mazes from Minigrid [67, 66] . We explore an additional dimension of OOS generalization in the agent's network hidden size; OPEN only saw agents with network hidden sizes of 16 in training, but is tested on larger and smaller agents. We include similar tests for OOS training lengths in Appendix H. We normalize OPEN against Adam, which was tuned for the same distribution and agent that OPEN learned for. a) On the left, we show performance in the distribution OPEN and Adam were trained and tuned in. On the right, we show OOS performance: the top row shows gridworlds from Oh et al. [13] , and the bottom row shows mazes from Chevalier-Boisvert et al. [67] .We mark Hidden Size = 16 as the in-distribution agent size for OPEN and Adam.OPEN learns update rules that consistently outperform Adam in-distribution and out-of-support with regards to both the agent and environment; in every OOS environment and agent size, OPEN outperformed Adam. In fact, in all but rand_dense, its generalization improves compared to Adam for some different agent widths, increasing its normalized return. In combination with our findings from Section 6.3, which demonstrate our approach can learn expressive update rules for hard, diverse environments, these results demonstrate the effectiveness of OPEN: if trained on a wide enough distribution, OPEN has the potential to generalize across a wide array of RL problems. Finally, in Figure 5 , we test how the optimizer trained above (i.e., for Figure 4 ) transfers to Craftax-Classic, an environment with completely different dynamics, 0-shot. We compare against two baselines: Adam (0-shot), where we tune the hyperparameters of Adam on the distribution of gridworlds that OPEN was trained on; and Adam (Tuned), where the Adam and PPO hyperparameters are tuned directly in Craftax-Classic. The latter gives an idea of the maximum achievable return.We find that OPEN transfers significantly better than Adam when both are tuned in the same distribution of environments, and performs only marginally worse than Adam tuned directly indistribution. These results exemplify the generalization capabilities of OPEN and lay the foundation for learning a large scale 'foundation' optimizer which could be used for a wide range of RL problems. 
ablation study: 
 In our ablation study, we explore how each constituent of OPEN contributes to improved performance. Firstly, we focus on two specific, measurable challenges: plasticity loss and exploration. Subsequently, we verify that OPEN can be meta-trained for an RL algorithm besides PPO.We ablate each individual design decision of OPEN in Figure 6 . For each ablation, we train 17 optimizers in a shortened version of Breakout [62, 63] and evaluate performance after 64 PPO training runs per optimizer. Further details of the ablation methodology are in Appendix I.While measuring plasticity loss is important, dormancy alone is not an appropriate performance metric; a newly initialized network has near-zero dormancy but poor performance. Instead, we include dormancy here as one possible justification for why some elements of OPEN are useful.(P1) Ablating training proportion directly disables the optimizer's ability to tackle training nonstationarity. Similarly, removing batch proportion prevents dynamic behavior within a (stationary) batch. The corresponding decreases in performance show that both timescales of non-stationarity are beneficial, potentially overcoming the impact of non-stationarity in RL.(P2) Removing dormancy as an input has a drastic impact on the agent's return, corresponding to a large increase in plasticity loss. While dormancy plays no direct role in the optimizer's meta-objective, including it as an input gives the optimizer the capability to react as neurons grow dormant, as desired.(P2) Not conditioning on layer proportion has a negative effect on final return. The increase in dormancy from its removal suggests that allowing the optimizer to behave differently for each layer in the network has some positive impact on plasticity loss.(P2)/(P3) Stochasticity benefits are two-fold: besides enhancing exploration (Section 7.2), it notably lowers dormancy. Though limited to the actor, this reduction likely also contributes to improve return. (P3) We verify the benefits of learnable stochasticity in Deep Sea, an environment from bsuite [71, 63] designed to analyze exploration. This environment returns a small penalty for each right action but gives a large positive reward if right is selected for every action. Therefore, agents need to explore beyond the local optimum of 'left at each timestep' to maximize return. Deep Sea's size can be varied, such that an agent needs to take more consecutive right actions to receive reward.Naïvely applying OPEN to Deep Sea struggles; we posit that optimizing towards the gradient can be detrimental to the actor since the gradient leads to a local optimum, whereas in the critic the gradient always points to (beneficially) more accurate value functions. Therefore, we augment OPEN to learn different updates for the actor and critic ('separated'). In Figure 7 , we show the disparity between a separated optimizer trained with and without learnable stochasticity after training for sizes between [4, 26] in a small number of generations (48) , noting that larger sizes are OOS for the optimizer. We include full results in Appendix I.2.The optimizer with learnable stochasticity consistently achieves higher return in large environments compared to without, suggesting significant exploration benefits. In fact, our analysis (Appendix F.5) finds that, despite being unaware of size, stochasticity increases in larger environments. In other words, the optimizer promotes more exploration when size increases. However, stochasticity does marginally reduce performance in smaller environments; this may be due to the optimizer promoting exploration even after achieving maximum reward, rather than converging to the optimal policy. To ensure that OPEN can learn for algorithms beyond PPO, and thus is applicable to a wider range of scenarios, we explore learning an optimizer for PQN [72] in Figure 8 . PQN is a vectorized and simplified alternative to DQN [73] . We include analysis and the return curve in Appendix I, and hyperparameters in Appendix C.We find that, in our single meta-training run, OPEN is able to outperform Adam by a significant margin. Therefore, while PPO was the principle RL algorithm for our experiments, our design decisions for OPEN are applicable to a wider class of RL algorithms. While exploring the ability of OPEN to learn for other common RL algorithms would be useful, Figure 8 provides a confirmation of the versatility of our method. 
limitations and future work: 
 While OPEN demonstrates strong success in learning a multi-task objective, our current approach of normalizing returns by Adam biases updates towards environments where Adam underperforms learned optimizers. We believe developing better curricula for learning in this settings, akin to unsupervised environment design [66, 74] , would be highly impactful future work. This may unlock the potential for larger scale experiments with wider distributions of meta-training environments. Additionally, the OPEN framework need not be limited to the specific difficulties we focus on here. Exploring ways to include other difficulties of RL which can be measured and incorporated into the framework (e.g., sample efficiency or generalization capabilities of a policy) could potentially elevate the usefulness of OPEN even further. Finally, while we show that OPEN can learn for different RL algorithms, further testing with other popular algorithms (e.g., SAC [75] , A2C [76] ) would be insightful. Furthermore, training OPEN on many RL algorithms simultaneously could unlock a truly generalist learned optimization algorithm for RL. Successfully synthesizing these proposals with additional scale will potentially produce a universal and performant learned optimizer for RL.Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: In Appendix J, we estimate how much compute was required per learned optimizer, how much compute was used for hyperparameter tuning and roughly how much compute the project used in total. We also detail that additional compute was used in the preliminary design of our method. We give a full breakdown of what hardware was used for experiments.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
conclusion: 
 In this paper, we set out to address some of the major difficulties in RL by meta-learning update rules directly for RL. In particular, we focused on three main challenges: non-stationarity, plasticity loss, and exploration. To do so, we proposed OPEN, a method to train parameterized optimizers that conditions on a set of inputs and uses learnable stochasticity in its output, to specifically target these difficulties. We showed that our method outperforms a range of baselines in four problem settings designed to show the expressibility and generalizability of learned optimizers in RL. Finally, we demonstrated that each design decision of OPEN improved the performance of the optimizer in an ablation study, that the stochasticity in our update expression significantly benefited exploration, and that OPEN was sufficiently versatile to learn with different RL algorithms.A Additional Details: Difficulties in RL (Problem 1) Non-stationarity Igl et al. [8] highlights different sources of non-stationarity across a range of RL algorithms. For PPO, non-stationary arises due to a changing state visitation distribution, the target value function V π (s) depending on the updating policy and bootstrapping from the generalized advantage estimate [77] . Since PPO batches rollouts, non-stationarity occurs between each batch; within the same batch, the agent solves a stationary problem. As the non-stationarity occurs over the course of training, we denote it training non-stationarity.Optimizers designed for stationarity, like Adam [18] , have been shown to struggle in non-stationary settings [9] and techniques to deal with non-stationarity have been proposed. Asadi et al. [78] prevent contamination between batches by resetting their optimizer state, including resetting their momentum. However, this handcrafted approach is potentially too severe, as it fails to take advantage of potentially useful commonality between batches.(Problem 2) Plasticity loss Plasticity loss, which refers to an inability to fit to new objective over the course of training, is an important problem in RL. As an agent learns, both its input and target distributions change (i.e., nonstationarity, (P1)). This means the agent needs to fit new objectives during training, emphasizing the importance of maintaining plasticity throughout training. Therefore, there have been many handcrafted attempts to reduce plasticity loss in RL. Abbas et al. [53] find that different activation functions can help prevent plasticity loss, Obando-Ceron et al. [79] suggest using smaller batches to increase plasticity and Nikishin et al. [60] introduces new output heads throughout training. Many have demonstrated the effectiveness of resetting parts of the network [80, 59, 81, 49] , though this runs the risk of losing some previously learned, useful information. Additionally, these approaches are all hyperparameter-dependent and unlikely to eliminate plasticity loss robustly.(Problem 3) Exploration Exploration in RL has a rich history of methodologies: [82] [83] [84] , to name a few. While there are simple, heuristic approaches, like ϵ-greedy [1] , recent methods have been designed to deal with complex, high dimensional domains. These include count based methods [85, 86] , learned exploration models [87, 88] or using variational techniques [89] . Parameter space noise [11] involves noising each parameter in an agent to enable exploration across different rollouts while behaving consistently within a given rollout. This inspired our approach to exploration since it is algorithm-agnostic and can be implemented directly in the optimizer.Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: This paper does not violate any part of the NeurIPS Code of Ethics; there are no potential harms caused by our research process, nor are there are any potential societal consequences.• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: We detail an example update step when using OPEN. We use notation from Appendix B.2, and use z GRU t to refer to the GRU hidden state at update t. In this algorithm, we assume the dormancy for each neuron has been tiled over the relevant parameters. In practice, we build our optimizer around the library fromMetz et al. [45] . Algorithm 1: Example update step from OPEN. Given: opt, N batch , N minibatches , L, H, β, α 1 , α 2 , α 3 Input: g t , m t-1 , p t , t, dormancies t , h, z GRU t-1 Output: (p t+1 , m t , z GRU t ) /* Compute new momentums */ for each β do m β t = β × m β t-1 + (1 -β) × g t end /*, h i ] GRU out , z GRU (t,i) = opt GRU (input, z GRU (t-1,i) ) m i , e i , δ i = opt M LP (GRU out ) ûi = α 1 m i exp α 2 e i if i in actor then ûi = ûi + α 3 δ i ϵ end end/* Zero-mean updates and apply them */ for each parameter i in agent dou i = ûi -E i [û i ] p (t+1,i) = p (t,i) -u i end return (p t+1 , m t , z GRU t ) B.We use layers of the same size in 'No Features', but condition on less inputs (14) and do not have the third output. We use the same architecture as Lan et al. [17] for Optim4RL, which includes two networks, each comprising a size 8 GRU, and two size 16 fully connected layers.When training optimizers for the multi-task setting (Section 6.3), we find increasing the network sizes consistently benefits performance. We show the architecture for OPEN in this experiment in Table 2 , masking out the noise for the critic by setting ϵ = 0. We also enlarge No Features and Optim4RL by doubling the GRU sizes to 16 and the hidden layers to size 32. C.1 PPO Hyperparameters PPO hyperparameters for the environments included in our experiments are shown in Table 3 . For our gridworld experiments, we learned OPEN and tuned Adam [18] for a PPO agent with W = 16, but evaluated on widths W ∈ [8, 16, 32, 64, 128] . Hyperparameters for PPO are taken from [32] where possible. For Adam, RMSprop and Lion, we tune hyperparameters with fixed PPO hyperparameters. For each environment we run a sweep and evaluate performance over four seeds (eight in the gridworld), picking the combination with the highest final return. In many cases, we found the optimizers to be fairly robust to reasonable hyperparameters. For Lion [20] , we search over a learning rate range from 3-10× smaller than Adam and RMSprop as suggested by Chen et al. [20] . For Optim4RL [17] , hyperparameter tuning is too expensive; instead, we set learning rate to be 0.1 × LR Adam following the heuristic from [20] .For Adam, RMSprop and Lion, we also test whether annealing learning rate from its initial value to 0 over the course of training would improve performance.A full breakdown of which hyperparameters are tuned, and their value, is shown in the following tables. We use notation and implementations for each optimizer from optax [69] . For Craftax-Classic (Table C .3), we test 0-shot optimizers using the gridworld hyperparameters to ensure transferring between the settings is as fair as possible. For the finetuned case, we use PPO hyperparameters from [21] which were tuned in domain. Also, for the finetuned case we set β 1 and β 2 values to defaults in optax [69] and tune the learning rate in the range {0.00005, 0.0005}. All optimizers are set to W = 64 so as to keep the network size close to, but still out of, distribution. For PQN, we use algorithm hyperparameters from [72] . For experiments with Adam, we tune learning rate, β 1 and β 2 and Anneal LR. Due to the length of meta-optimization, it is not practical to tune hyperparameters for meta-training. We, however, find the following hyperparameters effective and robust for our experiments. We use common hyperparameters when learning OPEN, Optim4RL and No Features.For the single-task and gridworld settings, we train all optimizers for a number of generations after their performance stabilizes, which took different times between optimizers, to ensure each learned optimizer has reached convergence; this occasionally means optimizers were trained for different number of generations, but enables efficient use of computational resources by not continuing training far past any performance gains can be realized. For the multi-task setting, due to the extreme cost of training (i.e., having to evaluate sequentially on four MinAtar environments), we train each optimizer We visualize the three mazes in Figure 10 . In Figure 11 , which shows return curves over training for the five 'single-task' environments (section 6), OPEN significantly beats 3 of the 5 baselines, performs similarly to Lion [20] in space invaders and performs comparably to hand-crafted optimizers and 'No Features' in ant. OPEN is clearly able to learn strongly performing update rules in a range of single-task settings. Throughout this section, we include analysis and figures exploring the behavior of optimizers learned by OPEN. In particular, we consider the behavior of OPEN with regards to dormancy, momentum, update size and stochasticity in different environments. Where relevant, we also make comparisons to the 'No Features' optimizer, introduced as a baseline in Section 6, to explore possible differences introduced by the additional elements of OPEN.Due to the number of moving parts in OPEN, it is difficult to make strong claims regarding the behaviour of the learned optimizers it produces. Instead, we attempt to draw some conclusions based on what we believe the data plots included below suggest, while recognizing that the black-box nature of its update rules, which leads to a lack of interpretability, is a potential drawback of the method in the context of analysis.All plots included below pertain to the single task regime, besides section F.5 which focuses specifically on deep sea.Figure 12 shows the τ = 0 dormancy curves during training for each of the MinAtar [62, 63] environments. These environments were selected as the only ones where the agent uses ReLU activation functions; ant [65, 64] used a tanh activation function for which τ = 0 dormancy is not applicable. In all environments excluding freeway, dormancy of all agents quickly converged to 0. Additionally, in every environment, OPEN produced more plastic (i.e., less dormant) agents than the featureless optimizer. We draw conclusions about this behavior below:• OPEN optimizes towards zero dormancy in a subset of environments, demonstrating it is capable of doing so. The fact that this only held in three out of four environments, despite OPEN outperforming or matching baselines in every environment, suggests it optimizes towards zero dormancy only when the agent plasticity is limiting performance (i.e., when plasticity loss prevents the agent from improving). Therefore, though a dormancy regularization term in the fitness function for meta-training may have been beneficial for asterix, breakout and spaceinvaders, it is possible that this could harm the performance in freeway and would require expensive hyperparameter tuning of the regularization coefficient to maximize performance. Our meta-learning approach completely sidesteps this problem. • Despite having the same optimization objective, No Features always had higher dormancy than OPEN. In tandem with our ablation from Figure 6 , it is likely that the inclusion of each additional feature has given OPEN the capability to minimize plasticity loss.We evaluate the cosine similarity between the update from OPEN and 'No Features' with momentum over different timescales, and the gradient, in Figure 13 . Momentum is calculated as m t+1 = βm t + (1 -β)g t , where g t is the gradient at time t. In OPEN, inspired by [24] , we input momentum using βs in [0.1, 0.5, 0.9, 0.99, 0.999, 0.9999]. For the ith timescale, we denote the momentum as m i in Figure 13 . For the gradient, we include an additional comparison against Adam. with the gradient or momentum at different β timescales. Each column corresponds to a different environment, and each row to a different timescale. In order, the rows show cosine similarity with gradient, followed by momentum at β = [0.1, 0.5, 0.9, 0.99, 0.999, 0.9999]. We plot mean cosine similarities with standard error over 16 runs.From Figure 13 , we can conclude that:• For both OPEN and 'No Features', alignment with momentum seems to be maximized at similar (though not always the same) timescales as the β 1 values tuned for Adam in Section C.2.• OPEN and 'No Features' rarely had similar cosine similarities. While it is difficult to discern whether this arises as a result of randomness (optimizers trained from a single seed) or something more fundamental, it may demonstrate how the additional elements of OPEN have changed its optimization trajectory.• For both learned optimizers, cosine similarity with gradient and momentum generally peaked and decreased over the training horizon. While OPEN includes lifetime conditioning, which may partially influence this behavior, it is likely that both optimizers rely on their own hidden states, which can be thought of as a learned momentum, more as training progresses.In the following experiments (i.e., Appendices F.3, F.4, F.5), we find it useful to divide updates by their respective parameter value. Unlike most handcrafted optimizers which produce shrinking updates from an annealed learned rate and generally output parameters with magnitudes close to zero, OPEN generally increases the magnitude of parameters over training with roughly constant update magnitudes. As such, normalizing updates with respect to the parameter value shows the relative effect (i.e., change) when an update is applied. In all cases, we plot the average absolute normalized value over time.In Figure 14 , we explore how the magnitude of the non-stochastic update changes over time (i.e., the update before any learned stochasticity is applied, defined as ûi = α 1 m i exp α 2 e i ). This is normalized with respect to the parameter value, and so is calculated We find that optimizers learned with OPEN all anneal their step sizes. In particular, the relative size of the update with respect to the parameter size decreases over time, with final updates generally being much smaller than at the beginning of training. This holds in all environments, though was weaker in ant which had significantly smaller updates from the start. Interestingly, for all handcrafted optimizers, hyperparameter tuning also consistently produces smaller learning rates for ant. As expected, this decreases significantly over time, suggesting the optimizers promotes less exploration over time. This makes sense; it is beneficial to explore early in training, when there is time to continue searching for better local or global optima, but later in training it is important to converge to the nearest optimum.|û/p| = E i | ûi pi | .We also note the significantly different scales of noise which are applied to different environments. In particular, ant uses very small noise, suggesting it is not an environment which needs a significant amount of exploration. This may be one reason why OPEN performed similarly to handcrafted optimizers, and 'No Features'.Finally, we consider how the randomness used by OPEN changes during training in environments of different sizes from Deep sea [71, 63] . This randomness was shown to be crucial for learning in larger environments in Section 7.2. To analyze this, we consider 5 different sizes (10, 20, 30, 40, 50) , with the size corresponding to how many 'rights' the agent has to take in a row to receive a large reward. We look at randomness/p = E i | We focus on two behaviors of the optimizer. Firstly, the applied randomness progressively decreases over training, reducing the exploration as the agent converges. This follows similar behavior to figure 15 , where the level of randomness decreases during the training horizon. Secondly, the amount of noise seems to grow with the environment size; the optimizer promotes more exploration in larger environments, despite lacking any awareness about the environment size. The combination of these two elements allows the agent to explore sufficiently in larger environments, while converging to good, close to optimal policies towards the end of training.Figure 17 shows return curves the four MinAtar [62] environments included in 'multi-task' training (Section 6. As expected, the method of normalization used in training (i.e., dividing scores in each environment by scores achieved by Adam and averaging across environments) leads OPEN, and the other learned optimizers, to principally focus on those which they can strongly outperform Adam in. In particular, OPEN marginally underperforms Adam and other handcrafted optimizers in freeway and spaceinvaders. However, its score in asterix, which is approximately double Adam's, will outweigh these other environments. We leave finding better methods for multi-task training, such as using more principled curricula to normalize between environments, to future work. However, we still note that this demonstrates that OPEN is capable of learning highly expressive update rules which can fit to many different contexts; it performs comparatively to handcrafted optimizers in two environments and significantly outperforms them in two others, in addition to consistently matching or beating the other learned baselines. Similar to Figure 4 , which looked at how OPEN generalized to different agent sizes, we consider how OPEN adapts to different training lengths in Figure 18 . In particular, we note that OPEN was only ever trained for gridworlds ('all', Appendix D) which ran for 3e7 training steps, so all lengths outside of this are OOS tasks. We also consider that in [13] and the bottom row inspired by mazes from [67] , all implemented by [66] . We mark 3e7 timesteps as the in-distribution training length for both OPEN and Adam.We train 17 optimizers for each feature ablation, and evaluate performance on 64 seeds per optimizer, using a shortened version of Breakout. We use the same PPO hyperparameters for Breakout as in Table 3 besides the total timesteps T , which we shorten to 2 × 10 5 . Each optimizer takes ∼ 60 minutes to train on one GPU, and we train a total of 119 optimizers.We train each optimizer for 250 generations, keeping all other MinAtar ES hyperparameters from Table 9 .We show the final return against size for the Deep Sea environment [71, 63] in Figure 16 for all optimizers; both shared and separated, with stochasticity and without. The one without can be thought of as a deterministic optimizer, and is denoted 'ablated' in the figure. For this experiment, we trained each optimizer for only 48 iterations on sizes in the range [4, 26] ; any sizes outside of this range at test time are out of support of the training distribution.Clearly the separated optimizer with learned stochasticity is the only one which is able to generalize across many different sizes of Deep Sea. This occurs marginally at the behest of optimization in smaller environments, where the stochastic optimizer still explores after obtaining reward and so does not consistently reach maximum performance. In Deep Sea, the final return scale is between 0.99 and -0.01. Figure 20 shows the performance of OPEN and Adam when training a PQN agent [72] in Asterix [62, 63] . While OPEN outperforms Adam in this setting, we note that it exhibits similar behavior to the other return curves (Appendix E, G) in having much lower return at the beginning of training before overtaking Adam at the end of training. As discussed in section 6.2, this suggests that OPEN sacrifices greedy optimization in favor of long-term performance. For all of our hyperparameter tuning experiments, we used a total of ∼ 16 GPU days of compute, using Nvidia A40 GPUs.For each learned optimizer in our single-task experiments, we used approximately ∼ [16, 32] GPU days of compute. This equates to around 1 GPU year of compute for 3 learned optimizers on 5 environments (=15 optimizers total). This was run principally on Nvidia GeForce GTX 1080Tis, Nvidia GeForce RTX 2080Tis and Nvidia GeForce RTX 3080s.For each multi-task optimizer, we used around 60 GPU days of compute, on Nvidia GeForce RTX 2080Tis. This totals 180 GPU days of compute. We also ran experiments for each optimizer with the smaller architecture, which took a similar length of time. These are not included in the paper; in total multi-task training used another 1 GPU year of compute, taking into account the omitted results.To train on a distribution of gridworlds, we used 7 GPU days of compute on Nvidia GeForce RTX 2080 Tis.Our ablation study used ∼ 8 GPU days of compute, running on Nvidia A40 GPUs.Each optimizer took 2 days to train in deep sea on an Nvidia A40 GPU. In total, this section of the ablation study took 8 GPU days.As stated in Table 13 , inference was a negligible cost (on the order of seconds) in this paper.Totaling the above, all experiments in this paper used approximately 2.1 GPU years to run, though it is difficult to aggregate results properly due to the variety of hardware used.We ran a number of preliminary experiments to guide the design of OPEN, though these are not included in the final manuscript. We are unable to quantify how much compute was used to this end. Due to the excessive cost of learned optimization, we ran only one run for each setting and optimizer (besides the multi-task setting, where we tested both small and large architectures and found that the larger architectures performed best for every optimizer on average). After converging upon our method, and finishing the implementations of Optim4RL and No Features, we did not have any failed runs and so included all experiments in the paper.Below we include a full list of assets used in this paper, in addition to the license under which it was made available.Table 14 : Main libraries used in this paper, links to the github repository at which they can be found and the license under which they are released.Name Link License evosax [34] evosax Apache-2.0 gymnax [63] gymnax Apache-2.0 brax [64] brax Apache-2.0 learned_optimization [45] learned_optimization Apache-2.0 groove (gridworlds) [66] groove Apache-2.0 purejaxrl [15] purejaxrl Apache-2.0 optax [69] optax Apache-2.0 rliable [70] rliable Apache-2.0 craftax [21] craftax MIT purejaxql (PQN) [72] purejaxql Apache-2.0NeurIPS Paper Checklist 1. Claims Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: We make claims regarding the performance of our learned optimizers for RL, and the means by which these performance gains are realized. These are justified by empirical results in a range of environments (Section 6, Appendices E, G, H), including three extensive different ablations (Section 7). Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.Seaquest, which is a part of MinAtar[62], does not have a working implementation in gymnax[63]. 


Decision: The paper proposes an interesting research direction to tackle the specific issues that arise in deep RL optimization. The approach suggests to make use of a "learned optimization" to overcome these problems in a meta-learning setting.

Here are the strong points of the paper provided in the reviews:
- The contributions has the potential to lead to follow-up research with a significant impact (all reviewers). 
- The strong evaluation and extensive ablation studies are convincing and support well the goal and claims of the paper (reviewers Fr37 and KVk8).
- Discussion of acceleration via parallelization in JAX is appreciated (reviewer sKXo).
- The quality of the paper is appreciated in terms of writing concerning introduction, related work, analysis of results, etc. (all reviewers).

A few (minor) weaknesses or suggestions:
- The idea about meta-learning an optimizer is not fully original (see reviewer KVk8).
- Even more extensive experiments could strengthen the results even further (see reviewer Fr37) and some additional experiments were suggested (e.g. reviewers sKXo and qjkS).

**Summary**: Given the strong quality of the paper, we recommend the paper for acceptance and even a spotlight. The originality of the paper is good but also not highly significant and we don't recommend the paper for an oral presentation.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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?
limitations: limitations: None
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: questions: Can the authors say something about the computational cost of their method versus the other methods used in their evaluation?
limitations: limitations: Related to my question, maybe a small section about computational differences should be added to the main text.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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.
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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.
limitations: 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.

Comments: title: Response to Rebuttal
comment: Thanks for your thorough response. I will keep my positive score.
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
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.
Dear KVk8,

Thank you for your additional response, and we are very grateful that you appreciate our work! Please do let us know if you have any other queries about our work that we might be able to answer.
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.
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.
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.
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.