Can Learned Optimization Make Reinforcement Learning Less Difficult?
Authors: Alexander D. Goldie
Chris Lu
Matthew Thomas Jackson
Shimon Whiteson
Jakob Nicolaus Foerster
Venue: NeurIPS 2024 spotlight (2024)
Keywords: Meta-Learning
Reinforcement Learning
Learned Optimization
Deep Learning
Primary area: reinforcement_learning
TL;DR: We propose OPEN, a method for learning optimizers designed to improve final return in reinforcement learning by tackling the difficulties of plasticity loss, exploration and non-stationarity.

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.

Can Learned Optimization Make Reinforcement Learning Less Difficult?
Alexander D. Goldie
Chris Lu
Matthew T. J.Jcksson'
Shimon Whiteson 12
Jakob Foerster
FLAIR, University of Oxford
WhiRL, University of Oxford
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 opti- mization 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 these difficulties. We show that our parameterization flexible enough to enable meta-learning in diverse learning contexts, including the ability to use stochasticity exploration. Our experiments demonstrate that when meta-trained single and small sets of environments, OPEN outperforms or equals traditionally used optimizers. Furthermore, OPEN shows strong generalization characteristics across range of environments and agent architectures.
Introduction
Reinforcement learning RLJ has undergone significant advances in recent years, scaling from solving complex games 2 towards approaching real world applications (477.. 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 maximize its performance.
Here, we take inspiration from three difficulties RL: -ootsnationaritty due to continuously changing input and output distributions [8; high degrees of plasticity loss limiting model capacities ( [0];
and exploration, which is needed to ensure an agent does not converge to local optima prematurely [[] 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 tackle these problems have relied human intuition find hand-crafted solutions. However, this is fundamentally limited by human understanding. Meta-RL. [12) offers an alternative which RI. algorithms themselves are learned from data rather than designed by hand. Meta-learned RI. algorithms have previously demonstrated improved performance over hand-crafted ones [3]15] 15
On related note, learned optimization has proven successful in supervised and unsupervised learning (e.g., VeLO [6)). Learned optimizers are generally parameterized update rules trained to outperform handcrafted algorithms like gradient descent. However, current learned optimizers perform poorly in RL [6 17). While some may argue that RL is simply an dfistt-iiittibuttio task [16], the lack of
"Correspondence to goldie @ www.os.oo....oooooo 'Equal supervision. Open-source code available here
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

consideration for, and inability to targer, specific difficulties of RL in naïve learned optimizers leads us believe that they would still underperform even if RL were in the training distribution.
We propose an algorithm, shown in Figure I for meta-learning optimizers specifically for called Learned Optimization for Plasticity, Exploration and Nonstationarity (OPEN). OPEN meta-leams update rules whose inputs are rooted in previously proposed solutions to the difficulties above, and whose outputs use learnable stochasticity to boost exploration. OPEN 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 [8]. R.SPrrp [9).. open-source discovered/lcarnet optimizers (Lion (20), VeLO [T6] and baseline learned optimizers trained for RL (Optim4RL [77. 'No Features'). We show that OPEN
fit to single and small sets of environments and also generalize in- and out-of-suppor of its training distribution. We further find that OPEN generalizes better than Adam [18 gridworlds and Traftax-Classtf (21). Therefore, our novel approach of increasing optimizer flexibility via extra inputs and an exploration- driven output, rather than enforcing more structured updates, enables significant performance gains.
Sample new
grad
update
opt
opt
Mela Learner
OPEN
inputss
Final Returns
Environment
action
agenni
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 [.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
2 Background
Reinforcement I Learning The RL problem is formulated Markov Decision Process 1 MDP] described by the tuple (A, P. R. p. 7)- At discrete timestep f,, the agent takes action a E A sampled from its (possibly stochastie) policy, <--[se) II, which conditions
the current state 8, E S (where 8p ~ p). After each action, the agent receives reward R(s,, a,) and the state transitions 114-1 based on the transition dynamics P(s i (s,,ar),, An agent's is objective is to maximize its discounted expected return, J",, corresponding to a discount factor 7 E [0, 1), which prevents the agent making myopic decisions. This is defined
00
J
E Ban Ma
RR
(1)
1=0
Proximal policy optimization 22 PPO] is algorithm designed to maximize J" It uses advantage, (" a). which is calculated from the state value function, V/8 E, Ei R' Rg [S. s). and state- action value function, Q (s, a) E. E 7 S, 8, a). It measures the improvement of specific action over the current policy and takes the form
A²(s, (-.
(2)
PPO introduces loss for optimizing the policy, parameterized by 8, that prevents extreme policy updates in gradient ascent. This uses clipping, which ensures there is no benefit to updating beyond mo(aels))
where the policy probability ratio, re( = Feaia ( exceeds the range [1 + c). The clipped loss is
(s,, ax), PPO is an actor-critic I 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, value function error, and an entropy bonus into
CLIP E [min(r.(0)A (T),
clip(r:(0))
e) (sc,
(o) ELL (0) C. (0) (=)(
(4)
2

---

Meta-Learning Optimizers in RL Algorithms like Adam [8] or RMSprop [9 are designed to maximize an objective by updating the parameters of 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 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 [6 24]. For any parameterized optimizer opt which conditions on a set inputs æ, the update rule to produce update u. can be described as a function, (ptx(x)) + u.
We treat learning to optimize as a meta-RL problem [1]. In meta-RL, the goal is to maximize JR over a distribution of MDPs P(M). For our task, an optimizer trained maximize J()) EM [J+opp yields the optimal patttaaaaamttniiatooonn opt
Evolution Strategies Evolution algorithms (EA) are a buckpropagaatio--free, black-box method for optimization (25 which uses a population of perturbed parameters sampled from a distribution (here,
N~0000 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] CMA-ES (2)).
Natural evolution strategies (NES) [30 are class of ES methods that use the population fitness
to estimate natural gradient for the mean parameters, 0. This can then be optimized with typical
gradient ascent algorithms like Adam [8] Salimans et 26 introduce OpenAl ES for optimizing f using the estimator
V VB VeE,~ σe) (00 oe)e).
(5)
which is approximated using a population average. practice, we use antithetic sampling (i.e., for each sampled evaluating + and -c) [3]. Antithetic task sampling enables learning on 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 re-
cently demonstrated the feasibility of ES for meta-RL, owing the speedup enabled by vectorization in Jax 33. We use the implementation of OpenAl ES (26) from evosax [34].
3 Related Work
3.1 Optimization in RL
Wilson et al. [35] and Reddi et al. (36 show that adaptive optimizers struggle in highly stochastic processes. Henderson al. 33 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 perfeetly 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 RI, take expressive approach by replacing the whole optimizer instead of just its momentum calculation, and using meta-lcarning to fit our optimizer to data rather than relying on potentially suboptimal human intuition.
3.2 Meta-learning
Discovering RL Algorithms Rather than using handcrafted algorithms, a recent objective in meta-RL is discovering RL algorithms from data. While there many successes in this area (e.g.. Learned Policy Gradient [3 LPG), MetaGenRL 14), Active Adaptive Perception 4T and Learned Policy Optimisation [1]] LPO]), we focus meta-learning 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-awee 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 RI. (9]. Instead, we propose replacing the optimizer itself with an expressive and dynamic update rule that is not subject to these problems.
3

---

Learned Optimization Learning to optimize 223 42 43
L20] strives to learn replacements to
handcrafted gradient- crasee optimizers like Adam 18 generally using neural networks (e.g.. [24
46).. While they show strong performance in supervised and unsupervised learning 16 47 previous
learned optimizers do not consider the innate difficulties of RI., limiting transferability. Furthermore, while VeL.O used 4000 TPU-months of compute [6), OPEN requires on the order of one GPU-month.
Lion 220 is symbolic optimizer discovered by evolution that outperforms AdamW 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.
48 using
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]. Optim4RI. [17] attempts to solve this issue by learning to optimize directly in RL. However, in their tests and ours, Optim4RL fails to consistently beat handerafted benchmarks. Also, it relies heavily constrained update expression based on Adam 18], and it needs expensive learning rate tuning. Instead, we achieve much stronger results with completely black-box setup inspired by preexisting methods for mitigating specific difficulties in RL.
4 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-stationall non-stationarity over the training process [ as the updating agent causes changes to the training distribution. We denote this training non- 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 theme in recent deep RL literature [ 49 50 [0]. Here, we focus dormancy [49] measurement tracking inactive neurons used as a metric for plasticity loss [10] 5 It is defined as (x)D)
I EnER (x)
HT
where (a:) is the activation of neuron i in layer / with input E D for distribution D. H' is the total number of neurons in layer 1. The denominator normalizes average dormancy to in each layer.
A neuron is 7-dormnnt if sf < T, meaning the neuron's 's output makes up less than 7 of its layer's output. For ReLU activation functions, т 0 means 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 Exploration Exploration is a key problem in RL [.] To prevent premature conver- gence 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 [II], where noise is applied the parameters of the agent rather than to its output actions, like c-greedy [1.
5 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.
5.1 Architecture and Parameterization
To enable conditioning on history, which required to express behavior like momentum, for example, OPEN uses gated recurrent unit [54] GRUJ. 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
Update Expression 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 is
i O: m exp a2e4
where my and e; are optimizer outputs for parameter i, and (122 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 need for the critic to explore. We take inspiration from parameter space noise for exploration [[], since it can easily be applied via the optimizer. To be precise, we augment the actor's update actor, new
: aatcorr + ox36
Here, 0g is small, stabilizing scaling factor, opctor is third output from the optimizer, and (0,, 1) is sampled Gaussian noise. By multiplying Spctor and 6, we introduce a random walk of learned, per-parameter variance the update which can be used for exploration. Since 6 depends on the optimizer's ss inputs, this can potentially learn complex interactions between the noise schedule and the features outlined in Section 5.3 Unlike Plappert al. [. 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, naive application of the above updates can cause errors as, even without stochasticity, agent parameters can grow to 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, as
u û B[q]]
While 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 is updated as P; P U4
(++1) (5)
5.2 Training
We train our optimizers with OpenAl ES 126]. using final return as the fitness. We apply the commonly-usee 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 33 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 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 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.
5.3 Inputs
Carefully selecting which inputs to condition OPEN on is crucial; they should be sufficiently expres- sive 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
46 Since PPO learns from successively
collected, stationary batches of data [56], we condition OPEN how far it is through updating with
5

---

the current batch (i.e., batch proportion). This enables behavior like learning rate scheduling, which has proved effective stationary problems (e.g.. [57 58) ) and bias correction from Adam account for inaccurate momentum estimates [88]
Inspired by Jackson et al. (32], who demonstrate the efficacy of learning dynamic versions of LPG [33] and LPO [15. we also condition OPEN on how far the current batch is through the total number of batches 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 hyperparameter. To replicate this, and enable varied behavior between the different layers of an agent's 's network, we condition OPEN on the relative position of the parameter's layer in the network.
(P2) Dormancy As Sokar et al. 44 reinitialize dormannt 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. [4]. 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.
6.1 Experimental Setup
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 (c.g.. [16 (7|47) which are also constrained by the high computational cost of meta-learned optimization. We provide the cost of experiments in Appendix J optimizers. We detail hyperparameters in Appendix C We define four evaluation domains, based on Kirk al. [61), in which effective learned optimization framework should prove competent:
including comparison of runtimes with other
Single-Tast Training A learned optimizer must be capable of fitting 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 and Ant from Brax [64, 65]. This is referred to as 'singleton' training in Kirk et al. [6].
62 63]
Multi-Task Training To show an optimizer is able to perform under a wide input distribution, it must be able to learn in number of environments simultaneously. Therefore, evaluate performance from training in all four environments from MinAtar 662 63 normalized score across environments with respect to Adam.
We evaluate the average
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. 666) with antithetic task sampling [32], and evaluate performance by sampling tasks from the same distribation. We include details in Appendix D
Out-Or-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 al. [3] and set of mazes from minigrid 677 which are not in the training distribution, with unseen agent parameters. Furthermore, we test how OPEN performs 0-shot Craftax -Caassic (2 a Jax reimplementation of Crafter [68]
Baselines We compare against open-source implementations [69] of Adam [88, RMSProp 19] Lion [20 and VeLO 16 45). We also learn two optimizers for the single- and multi-task settings: 'No Features', which only conditions instead of meta-gradients, as in Lan et al. [7]). Since Optim4RL is initialized close to sgn (Adam), and tuning its learning rate is too practically expensive, we set learning rate of 0.1 LR Adam based on Lion [20 (which moves from Adam W to sgn( AdamW)). The optimizer's ss 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
gradient and momentum; and Optim4RL [17] (using ES
SSeaquest, which is part of MinAtar 62]. does not have working implementation in gymnax 63

---

6.2 Single-Task. Training Results
Figure Sesure2shooe 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 its dynamic update rule. We further analyze the optimizers' behavior in Appendix F
freeway
spaceinvaders
breakout
asterix
ant
Optim4RL
VeLO
RM
Adam
OPEN
64
40
5.6
163
Figure 2: IQM of final returns for the five single-task training environments, evaluated over 16 random environment seeds. We plot 95% stratified bootstrap confidence intervals for each environment.
6.3 Multi-Task Training Results
Figure 3 shows each optimizer's s ability to fit to multiple MinAtar environments 63 62 where handcrafted optimizers tuned per-environment. We normalize returns with respect Adam, and aggregate scores over environments to give 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 al. [70] measuring how close to optimal algorithms are. Unlike single-task training, where optimizers are trained until convergence, run multi-task experiments for a fixed number of generations (300) limit compute.
OPEN drastically outperforms every optimizer for multi-task training. In particular, OPEN produces the only optimizer with aggregate score higher than Adam, demonstrating its ability learn highly expressive update rules which can fit to a range of contexts; no other learned optimizers get close OPEN in fitting 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, Optim4RI. seems to perform better in the multi-task setting than the single-task setting. This may be due to the increased number of meta-traine aata-training samples.
Mean
IQM
Optimality Gap
No reatume
Optm4RL
VeLO
Lion
RMSpmp
Adam
OPEN
0.90
0.18
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 We plot 95% stratified bootstrap confidence intervals for each metric.
62
63
6.4
Generalization
4 shows OPEN's ability to generalize in 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 B3 (66 and the bottom row are 3 mazes from Minigrid We explore an additional dimension 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 tuned for the same distribution and agent that OPEN learned for.
Figure
67
66
7

---

rand_dent cense
rand_sppares
rand long
in distribution
1.20
OPEN
Normalzze Retum
1.15
1.10
Hokee
Hide
standard maze
sixteen rooms
labyrinth
1.00
0.95
16
84
128
600
Hidden Size
Figure 4: IQM of return, normalized by Adam, in seven gridworlds, with 95% stratified bootstrap confidence intervals for 64 random seeds. 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. 667 We mark Hidden Size = 16 as the in- dn-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; 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 wide enough distribution, OPEN has the potential to generalize across wide array of RL problems.
craftax classic
Finally, in Figure 5 we test how the optimizer trained above (i.e., for Figure 4 transfers to Craftax-Classicn an environment with completely different dynamics, 0-shot. We compare against two baselines: Adam (0-shot), where we tune the hyperparameters of Adam the distribution of gridworlds that OPEN was trained on; and Adam (Tuned), where the Adam and PPO hyperparame- ters are tuned directly gives idea of the maximum achievable return.
12.5
10.0
Return
7.5
5.0
Craftax-Classic The latter
2.5
Adam (Tuned)
OPEN (0-shol)
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 in- distribution. These results exemplify the generaliza- tion capabilities of OPEN and lay the foundation for learning a large scale 'foundation' optimizer which could be used for wide range of RI. problems.
0.0
Adam (0-shot)
3
Frames
1e7
Figure 5: A comparison of OPEN and Adam with and without hyperparameter tuning in Craftax-Classic. We plot mean return over 32 seeds. Standard error negligible 0.06).
7 Ablation Study
In our ablation study, explore how each constituent of OPEN contributes 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.
7.1 Individual Ablations
We ablate each individual design decision of OPEN in Figure For each ablation, we train 17 optimizers a shortened version of Breakout 62 63 and evaluate performance after PPO training runs per optimizer. Further details of the ablation methodology are in Appendix]

---

While measuring plasticity loss is important, dormancy alone is not an appropriate performance metric: newly initialized network has near-zero dormancy but poor performance. Instead, 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 non- stationarity. 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 -on-stationarity RL.
(P2) Removing dormancy as input has drastic impact on the agent's return, corresponding to large increase in plasticity loss. While dormancy plays no direct role in the optimizer's 's meta- meeaa-bbeective, including it as an input gives the optimizer the capability react as neurons grow dormant, as desired.
(P2) Not conditioning layer proportion has negative effect on final return. The increase in dormaney from its removal suggests that allowing the optimizer behave differently for each layer in the network has some positive impact plasticity loss.
(P2)(P3) Stochasticity benefits are two-fold: besides enhancing exploration (Section 7.2) notably lowers dormancy. Though limited to the actor, this reduction likely also contributes improve return.
No Features
No Stochasticity
No Training Proportion No Batch Proportion
No Dormancy
No Layer Proportion OPEN
10.4
11.2
12.0
12.8
16
Final Return
Dormant Neurons [%]
Figure 6: IQM of mean final return for 17 trained optimizers per ablation, evaluated on 64 random seeds each, alongside mean т O dormancy for optimizers in the interquartile range. We show 95% stratified bootstrap confidence intervals.
7.2
Exploration
+0.8
(P3) We verify the benefits of learnable stochas- ticity Deep Sea, an environment from bsuite designed to analyze exploration. This en- vironment returns a small penalty for each right
+0.7
Improvement with Stochasticity
+0.6
[71 63
+0.4
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 take more consecutive right actions to receive reward.
+0.3
+0.2
+0.1
-0.1
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 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 ('sepa- rated"). In Figure 7 we show the disparity between separated optimizer trained with and without learnable stochasticity after training for sizes between [4, 26) a small number generations (48), noting that larger sizes are OOS for the optimizer. We include full results in Appendix 1.2
-0.2
Deep Sea Size
Figure 7: IQM performance improvement of an optimizer with learnable stochasticity over without. We plot 95% stratified bootstrap confidence intervals over 128 random seeds.
The optimizer with learnable stochasticity consistently achieves higher return 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.
7.3 Alternative RL Algorithm
asterix (PQN)
To ensure that OPEN can learn for algorithms be- yond PPO, and thus is applicable to wider range of scenarios, we explore learning an optimizer for PQN [72] in Figure X PQN is vectorized and simplified alternative to DQN [73]. We include analysis and the return curve in Appendix [ and hyperparameters in Appendix C
Adam
OPEN
34
Figure 8: IQM of final return for Adam and OPEN after training PQN in Asterix. We plot 95% stratified bootstrap confidence intervals over 64 seeds.
We find that, in single meta-training run, OPEN able to outperform Adam by a significant mar- gin. 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 a confirmation of the versatility of our method.
Limitations and Future Work
While OPEN demonstrates strong success in learning multi-task objectíve, our curtent 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 veeaatraining environments. Additionally, the OPEN framework need not be limited to the specific difficulties we focus on here. Exploring ways 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 775] 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.
Conclusion
In this paper. we set out 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 so, we proposed OPEN, method to train parameterized optimizers that conditions on 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, 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.
Author Contributions
AG implemented the code, designed and performed all experiments and made the architectural and training design decisions of OPEN. AG also provided all analysis and wrote all of the paper.
CL proposed the initial project, contributed the initial implementation framework and provided assistance throughout the project's timeline.
MJ contributed to experimental and ablation design and helped with editing the manuscript.
SW and JF provided advice throughout the project, and in writing the paper.
10

---

Acknowledgements
We thank Benjamin Ellis for offering guidance throughout this project, and Pablo Samuel Castro for providing advice in the early stages of the research. We are also grateful to Michael Matthews, Scott Roux, Juliusz Ziomek and our reviewers for their comments and feedback.
Funding
AG and MJ are funded by the EPSRC Centre for Doctoral Training Autonomous Intelligent Machines and Systems EP/S024050/1. MJ is also sponsored by Amazon Web Services. JF is partially funded by the UKI grant EP/Y028481/1 (originally selected for funding by the ERC). JF is also supported by the JPMC Research Award and the Amazon Research Award. Our experiments were also made possible by a generous equipment grant from NVIDIA.
References
[1] Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. The MIT Press, second edition, 2018. URL http:://incompleeteeiiieee
nwt.beet//eeeeeee book/ooteeeboo--oeeeeee 2nd. html
[2] David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George den Driessche, Julian Schrittwieser, loannis Antonoglou, Veda Panncershelvam, Mare Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of Go with deep neural networks and tree search. Nature, 529(7587)4484 489, January 2016. ISSN 1476-4687. doi: 10. 1038/natuur 16961.
[3] David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot. Laurent Sifre, Dharshan Kumaran, Thore Graepel. Timothy Lillicrap. Karen Simonyan, and Demis Hassabis. A general reinforcement learning algorithm that masters chess, shogi. Go through self-play. Science (New York, N.Y.), 362(6419)) 19:1 19:1 1440-114, 2018. doi: 10.1 22l1////ence..a.e 126/scccncc.aar6400. aar6404.
[4] Matt Barnes, Matthew Abueg. Oliver F. Lange. Matt Deeds, Jason Trader, Denali Molitor, Markus Wulfmeier, and Shawn O' Banion. Massively Scalable Inverse Reinforcement Learning in Google Maps. arXiv preprint 23Xiv22305.1126 2023. doi: 10.48550/arXiv.2305... 12290.
[5] Daniel J. Mankowitz, Andrea Michi, Anton Zhernov, Marco Gelmi, Marco Selvi, Cosmin Paduraru, Edouard Leurent, Shariq lqbal, Jean- Baptiste Lespiau, Alex Ahern, Thomas Köppe, Kevin Millikin, Stephen Gaffncy, Sophie Elster, Jackson Broshear, Chris Gamble, Kieran Milan, Robert Tung, Minjac Hwang, Taylan Cemgil, Mohammadamin Barckatain, Yujia Li, Amol Mandhane, Thomas Hubert, Julian Schrittieserr, Demis Hassabis, Pushmeet Kohli, Martin Riedmiller, Oriol Vinyals, and David Silver. Faster sorting algorithms discovered using deep reinforcement learning. Nature, 618(7964))257-263, June 2023. ISSN 1476-4687. doi: 10.1038//4 586- 0260022--
[6] Paul Christiano, Jan Leike, Tom B. Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences, February 2023.
[7] Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. March 2022.
[8] Maxinilian Igl, Gregory Farquhar. Jelena Luketina. Wendelin Boehmer, and Shimon Whiteson. Transient non-stationarity and generalisation in deep reinforcement learning. In International Conference on Learning Representations, 2021.
[9] Clare Lyle, Zeyu Zheng, Evgenii Nikishin, Bernardo Avila Pires, Razvan Pascanu, and Will Dabney. Understanding plasticity in neural networks. In Proceedings of the 40th International Conference on Machine Learning, ICML 233 JMLR. org. 2023.
[10] Clare Lyle, Zeyu Zheng, Khimya Khetarpal, Hado van Hasselt, Razvan Pascanu, James Martens, and Will Dabney, Disentangling the causes of plasticity loss in neural networks, 2024.
[11] Matthias Plappert, Rein Houthooft, Prafulla Dhariwal, Szymon Sidor, Richard Y. Chen, Xi Chen, Tamim Asfour, Pieter Abbeel, and Marcin Andrychowicz. Parameter Space Noise for Exploration. In International Conference Learning Representations, February 2018.

---

[12] Jacob Beck. Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, and Shimon Whiteson, A Survey of Meta-Reiinorreement Learning, arXiv preprini arXiv:2301.08028, 2023.
[13] Junhyuk Oh, Matteo Hessel, Wojciech M Czarnecki, Zhongwen Xu, Hado P van Hasselt, Satinder Singh, and David Silver. Discovering reinforcement learning algorithms. Advances in Neural Information Processing Systems, 33:1060-1070, 2020.
[14] Louis Kirsch, Sjoerd van Steenkiste, and Juergen Schmidhuber. Improving generalization in meta rein- forcement learning using learned objectives. International Conference on Learning Representations. 2020.
[15] Chris Lu, Jakub Kuba, Alistair Letcher, Luke Metz, Christian Schroeder de Witt, and Jakob Foerster. Discovered policy optimisation. Advances in Neural Information Processing Systems, 35:11655--1646 2022.
[16] Luke Metz, James Harrison, C. Daniel Freeman, Amil Merchant, Lucas Beyer, James Bradbury, Naman Agrawal, Ben Poole, Igor Mordatch, Adam Roherts, and Jascha Sohl-Dickstein. VeL.:: Training Versatile Learned Optimizers by Scaling Up. arXiv preprint arXiv:221 1. 1.09760, 2022.
[17] Qingfeng Lan, A. Rupam Mahmood, Shuicheng Yan, and Zhongwen Xu. Learning to Optimize for Reinforcement Learning. arXiv preprint arXiv: 2023.
[18] Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. arXiv preprint arXiv: 14126::442.6980, 2017.
[19] Tijmen Tieleman, Geoffrey Hinton, et al. Lecture 6.5-rmsprop: Divide the gradient by running average of its recent magnitude. COURSERA: Neural networks for machine learning, 4(2):26 31, 2012.
[20] Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Yao Liu, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, and Quoe Le. Symbolic Discovery of Optimization Algorithms. arXiv preprint 20Xiv:2302.0667,. 2023.
[21] Michael Matthews, Michael Beukman, Benjamin Ellis, Mikayel Samvelyan, Matthew Jackson, Samuel Coward, and Jakob Foerster. Craftax: A lightning- benchmark for open-ended reinforcement learning. In International Conference on Machine Learning (ICML), 2024.
[22] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal Policy Optimization Algorithms. arXiv preprint arXiv: 1xxiv:/7000 06.400 2017.
[23] Marcin Andrychowicz, Misha Denil, Sergio Gomez Colmenarejo, Matthew W. Hoffman, David Pfau, Tom Schaul, and Nando de Freitas. Learning to learn by gradient descent by gradient descent. Neural Information Processing Systems, 2016.
[24] Luke Metz, Niru Maheswaranathan, C. Danicl Freeman, Ben Poole, and Jascha Sohl-Dickstein. Tasks, stability, architecture, and compute: Training more effective learned optimizers, and using them to train themselves. arXiv preprini vxXiv:200..1224: 2020.
[25] Ingo Rechenberg. Fvolutionsstrategie Optimicrung technischer systeme nach prinzipien der biologischen evolution. 1973.
[26] Tim Salimans, Jonathan Ho, Xi Chen, Szymon Sidor, and Ilya Sutskever. Evolution Strategies as Scalable Alternative to Reinforcement Learning. arXiv preprint 1rrxiv:703.03380 2017.
[27] Paul Vicol, Luke Metz, and Jascha bahl-Dickstein. Unbiased gradient estimation in unrolled computation graphs with persistent evolution strategies. In International Conference on Muchine Learning, pages 10553-10563. PMLR, 2021.
[28] Felipe Petroski Such, Vashisht Madhavan, Edoardo Conti, Joel Lehman, Kenneth O. Stanley, and Jeff' Clune. Deep Neuroevolution: Genetic Algorithms Are Competitive Alternative for Training Deep Neural Networks for Reinforcement Learning. arXiv preprint urXiv:1712.00566, 2018.
[29] Nikolaus Hansen and Andreas Ostermeier. Completely derandomized self-adaptation in evolution strategies. Evolutionary Computation, 1(2):599-195 2001. doi: 10.1162/0033655611100199999
[30] Daan Wierstra. Tom Schaul, Tobias Glasmachers. Yi Sun. Jan Peters. and Jürgen Schmidhuber. Natural evolution strategics. Journal of Machine Learning Research, 15(27):949 980, 2014.
[31] John Geweke. Antithetic acceleration of Monte Carlo integration in Bayesian inference. Journal of Econometrics, 38(1):73 89, 1988.
12

---

[32] Matthew Jackson, Chris Lu, Louis Kirsch, Robert Lange. Shimon Whiteson, and Jakob Foerster. Discover- ing lemporally-aware reinforcement learning algorithms. In Second Agent Learning in Endeedsssss Workshop, 2023.
[33] James Bradbury. Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclau- rin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao Zhang. JAX: Composable transformations of Python+NumPy programs, 2018.
[34] Robert Tjarko Lange. Evosax: JAX- based Evolution Strategies. arXiv preprint arXiv:2212.041800 2022.
[35] Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nati Srebro, and Benjamin Recht. The marginal value of adaptive gradient methods in machine learning. Advances in neural information processing systems, 30, 2017.
[36] Sashank Reddi, Satyen Kale, and Sanjiv Kumar. convergence of adam and beyond. International Conference on Learning Representations, 2018.
[37] Peter Henderson, Joshua Romoff, and Joelie Pineau. Where Did My Optimum Go?: An Empirical Analysis of Gradient Descent Optimization in Policy Gradient Methods, arXiv preprint 1rX.v:1810.02525, 2018.
[38] Emmanuel Bengio, Joelle Pineau, and Doina Precup. Correcting Momentum in Temporal Difference Learning. arXiv preprint (tti::/00603955. 2021.
[39] Richard S. Sutton. Learning to predict by the methads of temporal differences. Machine Learning, 3:9-44, 1988.
[40] Mchmet Sarigül and Mutlu Avei. Performance comparision of different momentum techniques on deep reinforcement learning. 2017 IEEE International Conference on INnovations in Intelligent SysTems and Applications (INISTA), pages 302-306, 2017.
[41] D.M. Bossens, N.C. Townsend, and A.J. Sobey. Learning to learn with active adaptive perception, Neural Networks, 1 115:30-49, 2019. ISSN 0893-6080. doi: n..0016i.eeumet22219.000
[42] Li and Jitendra Malik. Learning to optimize. In International Conference on Learning Representations, 2017.
[43] Diogo Almeida, Clemens Winter, Jie Tang, and Wojciech Zaremba. A Generalizable Approach Learning Optimizers, arXiv preprint arXiv:2106. 00958, 2021.
[44] Olga Wichrowska, Niru Maheswaranathan, Matthew W. Hoffiman. Sergio Gómez Colmenarejo, Misha Denil, Nando de Freitas, and Jascha Sohl-Dickstein. Learned optimizers that scale and generalize. In Proceedings of the 34th International Conference on Machine Learning Volume 70, 1CML 17, pages 3751-3760. Jogg, 2017.
[45] Luke Metz, C Daniel Freeman, James Harrison, Niru Maheswaranathan, and Jascha Sohl- Dickkstinn Practical tradeoff's between memory, compute, and performance in learned optimizers. In Conference on Lifelong Learning Agents, pages 142-164. PMLR, 2022.
optimization
[46] Erik Gärtner, Luke Metz, Mykhaylo Andriluka, C. Daniel Freeman, and Cristian Sminchisescu. Transformer-based learned optimization. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 11970-11979, June 2023.
[47] Luke Metz, Niru Maheswaranathan, Brian Cheung. and Jaseha Narain Sohl-Dickstein. Meta-learning update rules for unsupervised representation learning. arXiv preprint arXiv:1804 100.:0020400022,, 2018.
[48] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
[49] Ghada Sokar, Rishabh Agarwal, Pablo Samuel Castro, and Utku Evci. The dormant neuron phenomenon in deep reinforcement learning. In Pruceedings of the 40th International Conference on Machine Learning, ICML '23, JMLR. org, 2023.
[50] Clare Lyle, Mark Rowland, and Will Dabuey. Understanding and preventing capacity loss in reinforcement learning. In International Conference on Learning Representations, 2022.
[51] Hojoon Lee, Hanseul Cho, Hyunseung Kim, Dachoon Gwak, Joonkee Kim, Jaegul Choo, Se- Young Yun, and Chulhee Yun. PLASTIC: Improving Input and Label Plasticity for Sample Efficient Reinforcement Learning. In Seiert- Conference Neural Information Processing Systems, November 2023.
13

---

[52] Guozheng Ma. Lu Li, Sen Zhang, Zixuan Liu, Zhen Wang, Yixin Chen, Li Shen, Xueqian Wang. and Dacheng Tao. Revisiting Plasticity Visual Reinforcement Learning: Data, Modules and Training Stages. arXiv preprint urXiv:2310.07418. 2023.
[53] Zaheer Abbas. Rosie Zhao, Joseph Modayil, Adam White. and Marlos C. Machado. Loss of Plasticity in Continual Deep Reinforcement Learning. arXiv preprint arXh::2303.07500) 2023.
[54] Kyunghyun Cho, Bart van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-Decodee approaches. In Proceedings of SSST-8, Eighth Workshop on Syntex, Semantics and Structure in Statistical Translation, pages 103-111, 2014.
[55]
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer Normalization. arXiv preprint 2016.
[56] Shengyi Huang, Rousslan Fernand Julien Dossa, Antonin Raffin, Anssi Kanervisto, and Weixun Wang. The implementation details of proximal policy optimization. In ICLR Blog Track, 2022. URL URL https //iclr- Biclr-blog-i track github. io 22//2222222222222 25ppp/ htte--immmeeeettioneeeeeeeeee httthh//.oo0020002222/ppmmemeeep
https:/iclr-bloog
details./
[57] Ilya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts, In Interna- tional Conference on Learning Representations, 2017. URL httphttp::/// /tattpe://opereeviiee net/forum?id
Skq89Scxx
[58] Leslic N. Smith. Cyclical learning rates for training neural networks. In 2017 IEEE Winter Conference on Applications of Computer Vision (WACV), pages 4644472, 2017. doi: 10. 1..10///A...20...888...
[59] Evgenii Nikishin, Max Schwarzer, Pierluca D'Oro, Pierre-I Bacon, and Aaron Courville. The primacy bias in deep reinforcement lcarning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 16828-16847. PMLR, July 2022.
[60] Evgenii Nikishin, Junhyuk Oh, Georg Ostrovski, Clare Lyle, Razvan Pascanu, Will Dabuey, and Andre Bar- reto. Deep reinforcement learning with plasticity injection. In Workshop on Reincarnating Reinforcement Learning at ICLR 2023, 2023.
[61] Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A Survey Zero-shot Generalisation in Deep Reinforeement Learning. Journal of Artificial Intelligence Research, 06:201-264, January 2023.
[62] Kenny Young and Tian Tian. MinAtar: An atari- atari-inspired testbed for thorough and reproducible reinforce- ment learning experiments. arXiv preprint /93.v:/903.03/76 2019.
[63] Robert Tjarko Lange. gymnax: A JAX-based reinforcement learning environment library, 2022.
[64] C. Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax diffcrentiable physics engine for large scale rigid body simulation, 2021.
[65] Emanuel Todorov, Tom Frez, and Yuval Tassa, MuleCo: A physies engine for model-based control. In 2012 IEEE/RSJ International Conference Intelligent Robots and Systems, pages 5026-5033. IEEE, 2012.
[66] Matthew Thomas Jackson, Minqi Jiang, Jack Parker-Holder, Risto Vuorio, Chris Lu, Gregory Farquhar, Shimon Whiteson, and Jakob Nicolaus Foerster. Discovering general reinforcement learning algorithms with adversarial environment design. In Advances Neural Information Processing Systems, volume 36, 2023.
[67] Maxime Shevalier-Boisvert, Bolun Dai, Mark Towers, Rodrigo de Lazcano, Lucas Willems, Salem Lahlou, Suman Pal. Pablo Samuel Castro, and Jordan Terry. Minigrid & Miniworld: Modular & Customizable Reinforcement learning Environments for Goal- Chal-OOiented Tasks, June 2023.
[68] Danijar Hafner. Benchmarking the spectrum of agent capabilities. arXiv preprint arXiv:2109.0 0670009999....00 2021.
[69] DeepMind, Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley. Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman. George Papamakarios. John Quan, Roman Ring. Francisco Ruiz, Alvaro Sanchez, Laurent Sartran, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Miloš Stanojević, Wojciech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. htL
cow/google deepmind
14

---

[70] Rishabh Agarwal. Max Schwarzer. Pablo Samuel Castro, Aaron Courville, and Mare G Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in Neural Information Processing Systems, 2021.
[71] Ian Oshand, Yotam Doron, Matteo Hessel, Aslanides, Eren Sezener, Andre Saraiva, Katrina McKinney, Tor Latrimore, Csaba Szepesvári, Satinder Singh, Benjamin Van Roy, Richard Sutton, David Silver, and Hado van Hasselt. Behaviour suite for reinforcement learning. In International Conference on Learning Representations, 2020. URL https:://ppenrevies net htt///oormmiddddee wow.ggggkwwttt..
[72] Matteo Gallici, Mattic Fellows, Benjamin Ellis, Bartomeu Pou, Ivan Masmitja, Jakob Nicolaus Foerster, and Mario Martin. Simplifying deep temporal difference learning, 2024. URL https://aaxxiii org/
aba/2407.0482
[73] Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Andrei A. Rusu, Joel Veness, Mare G. Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattic, Amir Sadik, loannis Antonoglou, Helen King. Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Naure, 518(7540))529-535 February 2015.
[74] Michael Dennis. Natasha Jaques, Eugene Vinitsky. Alexandre Bayen, Stuart Russell, Andrew Critch, and Sergey Levine. Emergent complexity and zero-shot transfer via unsupervised environment design. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS '20, pages 13049-13061, Red Hook, NY, USA, December 2020. Curran Associates Inc. ISBN 978-1-71382- 954-6.
[75] Tuomas Haarmoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with Stochastic Actor, August 2018.
[76] Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous Methods for Deep Reinforcement Learning, June 2016.
[77] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High-Dimensional Continuous Control Using Generalized Advantage Estimation, Octoher 2018.
[78] Kavosh Asadi, Rasool Fakoor, and Shoham Sabach. Resetting the optimizer in deep RL: An empirical study, arXiv preprint arXiv:2306. 17833, 2023.
[79] Johan Slando-Ceron Mare G. Bellemare, and Pablo Samuel Castro. Small batch deep reinforcement learning. arXiv preprint arXiv:2310.03882 2023.
[80] Pierluca D' D'Oro, Max Schwarzer. Evgenii Nikishin, Pierre-Lue Bacon, Mare G. Bellemare, and Aaron Courville. Sample-Efficient Reinforcement Learning by Breaking the Replay Ratio Barrier. In The Eleventh International Conference on Learning Representations, September 2022.
[81] Max Schwarzer, Johan Samir Obando Ceron, Aaron Courville, Mare G Bellemare, Rishabh Agarwal, and Pablo Samuel Castro. Bigger, better, faster: Human-level atari with human-level efficiency. In International Conference on Machine Learning, pages 30365-30380. PMLR, 2023.
[82] Ronen I. Brafman and Moshe Tennenholtz. R-MAX- A General Polynomial Time Algorithm for Near- Optimal Reinforcement Learning. Journal of Machine Learning Research, 3:213-231, Oct 2002.
[83] Michacl Keamss and Satinder Singh. Near-Optima Reinforeement Learning in Polynomial Time. Machine Learning, 29(2)2209-232, November 2002.
[84] Peter Auer. Thomas Jaksch, and Ronald Ortner. Near-optimal regret bounds for reinforcement learning. In
D. Koller, D. Schuurmans, Bengio, and L. Bottou, editors, Advances in Neural Information Processing Systems, volume 21, 2008.
[85] Haoran Tang, Rein Houthooft, Davis Foote, Adam Stooke, OpenAl Xi Chen, Yan Duan, John Schulman, Filip DeTurck, and Pieter Abbeel. #Exploration: A study of count- based exploration for deep reinforcement learning. Advances in neural information processing systems, 30, 2017.
[86] Georg Ostrovski, Mare G Bellemare, Aäron Oord, and Rémi Munos. Count-based exploration with neural density models. In International Conference on Machine Learning, pages 2721-2730. PMLR, 2017.
[87] Bradly C. Stadie, Sergey Levine, and Pieter Abbeel. Incentivizing Exploration In Reinforcement Learning With Deep Predictive Models, arXiv preprint arXiv:1507.0000 000144 2015.
15

---

[88] Sainbayar Sukhbaatar, Zeming Lin. Ilya Kostrikov, Gabriel Synnaeve, Arthur Szlam, and Rob Fergus. Intrinsic motivation and automatic curricula via asymmetric self-play. In International Conference on Learning Reprexemtations 2018.
[89] Rein Houthooft, Xi Chen, Xi Chen, Yan Duan, John Schulman, Filip De Turck, and Pieter Abbeel. VIME: Variational information maximizing exploration. In Lee. Sugiyama, U. Luxburg, Guyon, and R. Gareett, editors, Advances in Neural Information Processing Systems, volume 29, Curran Associates, Inc., 2016.
[90] Luke Metz, Niru Maheswaranathan, Jeremy Nixon, Daniel Freeman, and Jascha standing and correcting pathologies in the training of learned optimizers. In Kamalika Chaudhuri and Ruslan Salakhutdinov. editors. Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 4556 PMLR, June 2019.
Under-
16

---

A Additional Details: Difficulties in RL
(Problem 1) Non-stationarity Igl al. [8] highlights different sources of non-stationarity across range of RL algorithms. For PPO, non-stationary arises due to changing state visitation distribution, the target value function V*8 (s) depending on the updating policy and bootstrapping from the generalized advantage estimate [77]. Since PPO batches rollouts, con-stationarity occurs between each batch; within the same batch, the agent solves 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 struggle in non-stationary settings 9 and techniques to deal with non-stationarity have been proposed. Asadi et al. 778 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.
Plasticity loss Plasticity loss, which refers to an inability to fit to new objective over the course of training, is an important problem RL. As an agent learns, both its input and target distributions change (i.c., 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.
(Problem 2)
[53] find that
different activation functions can help prevent plasticity loss, Obando-Cert Cendo 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: 882444, name a few. While there are simple, heuristic approaches, like e-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 [T] involves noising each parameter in agent to enable exploration across different rollouts while behaving consistently within a given rollout. This inspired our approach to exploration since algorithm-agnostie and can be implemented directly in the optimizer.
17

---

B Optimizer Details
B.1
Algorithm
We detail an example update step when using OPEN. We use notation from Appendix [..2 and use =GRU to refer the GRU hidden state update In this algorithm, assume the dormaney for each neuron has been tiled over the relevant parameters. In practice we build our optimizer around the library fromMetz al. 45
Ft
Algorithm 1: Example update step from OPEN.
Given: opt, Nbtchh Minnibbtches L, H. B, a1 a22 a3
Input: gt. m,,1, Pt, t, dormancies, h, GRU
Output: (p++1, m,, 2
/+ Compute new momentums
for each B do
m
8
m (1 3) x
end
/* Compute training and batch proportions
TP [I. Nmintbatches)) //ateceessshhhttt Nhatch
BP = ((t//Nmiiiattche) mod L./L
Compute first stage of update
for each parameter i in agent do
input= ) pomn BP, hu.
GG = opt GRU (input, GGRU GRU out, (t.2)
(r-1,i))
m4 6/ 6, opt MLP (GRU
û; 0m exp 022
if in actor then
end
end
updates and apply them for each parameter agent do 0; E, [
/*
P(++1.)) P((.)) - 4
end
return (Pr+1, m,, 2t
GRU
B.2 Input Features to the Optimizer
Our full list of features is concatenated and fed into the optimizer, as algorithm D below. We some notation from Table These features are:
Gradient 9. processed with the transformation > (sigg()),loo(g + c)}-
calculated with the following B coefficients: [0.1, 0.9, 0.99, 0.999, 0.9999]. This processed with the transformation
Momentum
m, + (sign(m.) log(m; + ))} for each B.
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 Noatch (T/ (TT/NNea sTt//Nttess///// //Nenvs/ this is defined as tt//(L+ * //tttabbatchhe/NNaot.ee where t is the current update iteration. We use floor division operator training proportion should be constant through batch.
Batch proportion, or how far the current update is through training with the same batch, caleulated as (tt//Nmmisbbtohhs)) mod L)/L.
18

---

Dormancy, D. calculated for the neuron downstream of the parameter using equation 6
This repeated over parameters with the same downstream neuron.
Proportion of the current layer through the network, h/ H.H.
The gradient, momentum and parameter value features are normalized to have second moment of I across the tensor, following Metz et al. [90] 45), as a normalization strategy which preserves direction. Like Lan et al. [7], we process gradients and momentum > (sign(x)),og(44 }) to magnify the difference between small values.
B.3 Optimizer Architecture Figure
Figure 9 visually demonstrates the architecture of OPEN. Since only the actor is updated with learned stochasticity (factor), multiply the sampled noise with 1/0 mask for the actor/critic.
Gradient
Momentum
V
Features
gactor
Figure 9: A visualization of the architecture used by OPEN. t are parameters of a GRU and 0 are parameters of an MLP. Since gactor only applies the actor's update expression, we set the noise to zero for the critic to keep its updates deterministic.
B.4 Single- Task and Gridworld Optimizer Architecture Details
The details of the layer sizes of OPEN in both the single-task and gridworld settings are in Table |
Table I: Optimizer layer sizes for OPEN in the single-task and gridworld experiments.
Layer Type
Dimensionality
GRU
19,8
Fully Connected
16
Layeroorm
Fully Connected
16, 16]
Layernorm
Fully Connected
[16,3] 3]
In our update rule, we let 011 022 03 0.001. We mask out the third output in the critic by setting the noise (E) to zero; otherwise, updates to our critic would be non-deterministic.
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 [17 for Optim4RL, which includes two networks, each comprising size 8 GRU, and two size 16 fully connected layers.
B.5 Multi-Task Optimizer Architecture Details
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 Table 2 masking out the noise for the critic by setting doubling the GRU sizes to 16 and the hidden layers to size 32.
0. We also enlarge No Features and Optim4RL by
Table 2: Optimizer layer sizes for OPEN in the multi-task experiment.
Layer Type GRU
Dimensionality
19, 16
Fully Connected
16, 32)
Layernorm
Fully Connected
[32, 32]
Layernorm
Fully Connected
[32, 3]
19

---

C Hyperparameters
C.1 PPO Hyperparameters
PPO hyperparameters for the environments included our experiments are shown in Table 3 For our gridworld experiments, we learned OPEN and tuned Adam [8] 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.
Table 3: PPO hyperparameters. All MinAtar environments used common PPO parameters, and are thus under one header.
Environment
Hyperparameter
MinAtar
Ant
Gridworld
Number of Environments Nenvs
64
2048
1024
Number of Environment Steps N eteps
128
10
20
Total Timesteps T
1 1x107
5x107
3 x 107
Number of Minibatches Nmiinsbatch
32
Number of Epochs L
4
4
2
Discount Factor
0.99
0.99
0.99
GAE
0.95
0.95
0.95
PPO Clip e
0.2
0,2
0.2
Value Function Coefficient C1
0.5
0.5
0.5
Entropy Coefficient 02
0.01
0
0.01
Max Gradient Norm
0.5
0.5
0.5
Layer Width W
64
64
16
Number of Hidden Layers H Activation
2
ReLU
tanh
tanh
C.2 Optimization Hyperparameters
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. many cases, we found the optimizers to be fairly robust to reasonable hyperparameters. For Lion (20). we search over learning rate range from 3-10x smaller than Adam and RMSprop as suggested by Chen et al. 20) For Optim4RL [T7] hyperparameter tuning is too expensive; instead, we set learning rate to be 0.1 x LR following the heuristic from 20
For Adam, RMSprop and Lion,
For Adam, RMSprop and Lion, we also test whether annealing learning rate from its initial value to over the course of training would improve performance.
full breakdown of which hyperparameters are tuned, and their value, is shown the following tables. We use notation and implementations for each optimizer from optax 669
20

---

Table 4: Optimization hyperparameters for MinAtar environments. values used in our hyperparameter tuning.
'Range' covers the range of
Hyper
Environment
Optimizer
Range
Parameter
asterix
freeway
breakout
space invaders
LR
0.003
0.001
0.01
0.007
{0.001, 0.01
Adam
Bi
0.9
0.9
0.9
0.9
(0.9.0.999]
B2
0.999
0.99
0.99
0.99
(0.9.0.999)
Anneal LR
(V.X)
LR 0.002
0.001
0.002
0.009
{0.001, 0.01]
RMSprop
Decay
0.99
0.999
0.99
0.99 /
{0.9, 0.999)
Anneal LR
(V.X)
LR 0.0003
0.0003
0.0008
0.0008
{0.0001, 0.001
Lion
31 0.99
0.9
0.9
0.9
{0.9, 0.999)
B2 0.9
0.9 V
0.9
0.99
{0.9, 0.999} (V.X)
Anneal LR
Optim4RI.
LR
0.0003
0.0001
0.001
0.0007
(0.1 x LRAdam
Table 5: Optimization hyperparameters for Ant. hyperparameter tuning.
'Range' covers the range of values used in our
Environment Ant
Optimizer
Hyperparameter
Range
LR
0.0003
0.00001, 0.001
Adam
81
0.99
(0.9. 0.999]
32
0.99
(0.9, 0.999)
Anneal LR
V
[V,X]
LR
0.0008
{0.0001, 0.005]
RMSprop
Decay
0.99
{0.9, 0.999} [V,X]
Anneal LR
X
LR
0.00015
{0.00001, 0.0005]
Lion
31
0.9
{0.9, 0.999}
32
0.9
{0.9. 0.999} |V.X! XX
Anneal LR
/
Optim4RI.
LR
0.00003
(0.1 x LR
Table 6: Optimization hyperparameters for Gridworld, 'Range' covers the range of values used in
our hyperparameter tuning.
Environment Gridworld
Optimizer
Hyperparameter
Range
LR
0.0001
(0.00005, 0.001}
Adam
31
0.99
(0.9. 0.999)
32
0.99
(0.9, 0.999)
Anneal LR
(V. X)
C.3 Craftax-Classic Hyperparameters
For Craftax-Classie (Table C.3 we test O-shot optimizers using the gridworld hyperparameters to
ensure transferring between the settings 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 Bn and
32 values to defaults in optax [69] and tune the learning rate the range (0.00005, 0.0005). All optimizers set to W 64 as to keep the network size close to, but still out of, distribution.
21

---

Table 7: Hyperparameters for Craftax-Classta
Hyperparameter
OPEN (0-shot)
Adam (0-shot)
Adam (Finetuned)
Learning rate
0.0001
0.0005
31
0.99
0.9
32
0.99
0.999
Anneal LR
V
V
Number of Environments Nen Nevvs
1024
1024
256
Number of Environment Steps Nateps
20
20
16
Total Timesteps
3 x 10
3 x 107
3 x 107
Number of Minibatches N Nminsbatch
16
16
8
Number of Epochs L
2
2
4
Discount Factor 7
0.99
0.99
0.99
GAE X
0.8
0.8
0.8
PPO Clip
0.2
0.2
0.2
Value Function Coefficient C1
0.5
0.5
0.5
Entropy Coefficient 02
0.01
0.01
0.01
Max Gradient Norm Layer Width W
0.5
0.5
0.5
64
64
64
Number of Hidden Layers H Activation
2
2
2
tanh
tanh
tanh
C.4 PQN Hyperparameters
For PQN, we use algorithm hyperparameters from rate, 31 and 83 and Anneal LR.
72 For experiments with Adam, we tune learning
Table 8: Hyperparameters for PQN in asterix.
Hyperparameter
OPEN
Adam
Tuning Range
Learning rate
0.0003
(0.0001, 0.001]
3
0.99
(0.9, 0.999)
82
0.999
(0.9, 0.999)
Anneal LR
V
(V.X)
Number of Environments Newvs
128
128
Number of Environment Steps N Nateps
32
32
Total Timesteps T
1 X 107
1 x 107
Number of Minibatches Nmmnibatch
Number of Epochs L
4
4
Discount Factor 7
0.99
0.99
GAE X
0.65
0.65
Max Gradient Norm
10
Layer Width W
128
128
Number of Hidden Layers H
Activation
relu
relu
Normaization
LayerNorm
LayerNorm
C.5 ES Hyperparameters
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 hyprraarameters 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 number of generations, but enables efficient use of computational resources by not continuing training far past any performance gains can realized. For the multi-task setting, due to the extreme cost of training (i.c., having to evaluate sequentially on four Min environments), train each optimizer
different
22

---

with an equivalent level of compute (i.e., 300 generations) and pick the best performing generation in that period.
We periodically evaluate the learned optimizers during training to find the one which generation performs best small in-distribution validation set. Table roughly how many generations we trained each optimizer for, and the generation used at test time, which varied for each environment. We find that some environments, Optim4RL made no learning progress from the beginning of training, with its performance being practically identical to its initialization.
Table 9: ES hyperparameters for meta-training. For MinAtar, the number of generations corresponds (Asterix, Breakout, Freeway, Spacelnvaders). We show the max generations for each optimizer (i.e., how long it was trained for), as well as the generation used (i.e., which training generation was used at inference time).
Environments
Hyperparameter
MinAtar
Ant
Multi-Task
Gridworld
Finit
0.03
0.01
0.03
0.01
decay
0.999
0.999
0.998
0.999
Learning Rate
0.03
0.01
0.03
0.005
Learning Rate Decuy
0.999
0.999
0.998
0.990
Population Size
64
32
64
64
Number of Rollouts
1
1
1
1
Max Gens (OPEN)
[350, 500, 700, 450}
700
300
350
Gen Used (OpeN)
(336, 312, 648, 144]
168
234
333
Max Gens (Optim4RL)
(500, 550, 300, 400)
~ 700
300
N/A
Gen Used (Optim4RL)
(288, 552, 216, 24)
480
288
N/A
Max Gens (No Features)
{800, 850, 500, 600)
400
300
N/A
Gen Used (No Features)
(456, 816, 408, 264) 24
240
270
N/A
Evaluation Frequency
24
9
9
23

---

D Gridworld Details
We train OPEN on gridworlds by sampling environment parameters from a distribution. Here, we use distribution implemented by [66 32 which detailed in Table I0 In the codebase of Jackson et al. [66], this is referred to as 'all'.
Table 10: Distribution parameters for Gridworld. In training, the values for the experienced environ- ment sampled from the value range.
Parameter
Value [range]
Max Steps in Episode
[20, 750)
Object Rewards
1 1.0
Object p(terminate)
00.01, 1.0
Object p pereppwwn)
0.101,001
Number of Objects
[1,
Grid Size
11]
Number of Walls
15
We also evaluate OPEN on three distributions from Oh et al. [13], and three mazes from Chevalier- Boisvert al. [67]. which were the distribution of 'all', We run all experiments on 64 seeds gridworld. These are all detailed below, with the mazes also visualized.
Table 11:
Distribution parameters for Gridworld. Curly brackets denote a list of the true values, corresponding to each object, used in testing.
Name
Parameter
Values 500
Max Steps in Episode Object Rewards
11,0,
Object p(terminate)
dense
Object respawn)
(0.05, 0.1.0.
Number of Objects
4
Grid Size
11
Number of Walls
0
Max Steps in Episode Object Rewards
50
(-1.0, 1.0)
Object p p(terminate)
(1.0,
rand_sparse
Object pl prespaw))
(0.0,0.0)
Number Objects
2
Grid Size Number of Walls
13
0
Max Steps in Episode Object Rewards
1000
(1.0, 1.0,-1.0,
Object p(terminate)
0.5.
rand LangJonn
Object pl respawn)
(0.01, 0.01,, 1.0, 1.0)
Number Objects
4
Grid Size
11
Number of Walls

---

Table 12:
Distribution parameters for Gridworld. At inference, true values are sampled from the given range, such that each seed evaluated in slightly different setting.
Name
Parameter
Value
Max Steps Episode
[25, 50
Object Rewards
[0.0, 1.0
Object p(terminate)
[0.01, 1.0
standard maze
Object pl (reppawn)
[0.001,
Number of Objects
3
Grid Size
13
Max Steps in Episode
[25, 50]
Object Rewards
[0.0, 1.0]
sixteen_room
Object p(terminate)
0.01, 1.0
Object p(respawn)
[0.001, 0.1]
Number of Objects
3
Grid Size
13
Max Steps Episode
[25, 50]
Object Rewards
[0.0, 1.0)
labyrinth
Object p (tterminate)
0.011, 1.0
Object p(respawn)
[0.001,0.1]
Number of Objects
3
Grid Size
13
We visualize the three mazes in Figure 10
YLfy
Figure 10: The three mazes from minigrid 67 used for our OOS tests. From left to right, the mazes are: 'standard maze', 'sixteen_ rooms' and 'labyrinth'.
25

---

E Single Environment Return Curves
In Figure which shows return curves over training for the five 'single-task' environments (section OPEN significantly beats 3 of the baselines, performs similarly to Lion (20] in space invaders
O).
and performs comparably hand-crafted optimizers and 'No Features' in ant. OPEN clearly able
to learn strongly performing update rules in a range of single-task settings.
freeway
ant
50
Rmmm
OP
20
Adam Hattlrop
10
Frallane
0.0
1.0
Frames
1
Frames
1o7
107
breakkout
spaceinvaders
60
Retu
100
20
0.0
0.4
08
0.5
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Frames
167
Figure 11: RL training curves comparing learned optimizers and all other baselines, each trained or tuned on a single environment and evaluated on the same environment. We show mean return over training with standard error, evaluated over 16 seeds.
26

---

F Analysis
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 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, attempt 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 a lack of interpretability, is potential drawback of the method in the context of analysis.
All plots included below pertain to the single task regime, besides section cally on deep sea.
F5 which focuses specifi-
F1
Dormancy
Figure 12 shows the 0 dormancy curves during training for each of the MinAtar 662 63
environments. These environments were selected as the only ones where the agent uses ReLU activation functions; ant 665 64 used tanh activation function for which 0 dormancy is not applicable.
Freeway
breakoutt
spaceinvsdders
6-5
0.3
3.09
3.66
6.4
0.4
2.64
0.1
5.0
8:00
3.00
0.25
0.50
0.75
3.00
0.25
0.50
0.75
3.25
0.50
6.75
1,02
125
Usdatas
Updares
Updinm
Figure 12: T 0 dormancy curves for all optimizers on the four MinAtar environments, evaluated over 16 seeds each. We plot the mean dormancy with standard error.
In all environments excluding freeway, dormancy of all agents quickly converged to O. 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 dormancy in 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
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 dormancy regular- ization 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 [ it is likely that the inclusion each additional feature has given OPEN the capability minimize plasticity loss.
27

---

F22 Similarity Between Update and Gradient/Momentum
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 3m' + (1 3)g². where S: the gradient time t. In OPEN, inspired by [24] we input momentum using 3s in 10.1, 0.5, 0.9, 0.99, 0.999, 0.9999). For the ith timescale, we denote the momentum as mi
in Figure 13 For the gradient, we include an additional comparison against Adam.
LFC
spaccisaaors
0.33
0.53
604
0.20
0.00
0.0
ealarh
1:25
413
640 0.10 0.0
123
Teeway
0.00
1.0
512
141
Treway
182
50 soaahteeee
415
613
4.44
0.23
4.4
4.44
414
0.05
404
420
408
800 530 010
0.500
0.3 0:0
3/0 3.55
KCM
130
1.2
Figure 13: Curves showing the cosine similarity between the updates from OPEN and 'No Features' with the gradient or momentum at different /3 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 3 [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 3
we can conclude that:
28

---

For both OPEN and 'No Features", alignment with momentum seems to be maximized at similar (though not always the same) timescales as the B 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 learned momentum, more as training progresses.
E3 Non-Stochastic Update
In the following experiments (i.e., Appendices E.3 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.
E4 F.5
we find it useful to divide updates
In Figure 14|
explore how the magnitude of the non-stochastic update changes over time (i.e., the update before any learned stochasticity is applied, defined as ú, = a,m; exp (2e)). This is normalized with respect to the parameter value, and so calculated (u/pl E,
km
3
km CR 000
100
Figure 14: Plots of the average normalized, non-stochastic update magnitude with respect to time, with standard error over 16 seeds.
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.
F.4 Stochasticity
Figure 15 explores how the weight of the learned stochasticity, 8ucter, changes with respect to time. This stochasticity is incorporated into the update rules as ârettur iiuetor + drgocttor,, with e ~ (0, 1). As in the previous plot, this is normalized with respect to the parameter value, and is calculated as
randomness/p E,
Emeaay
Ennskret
5/00
R20
313
Figure 15: Plots of the normalized stochasticity weight with respect to time. We plot mean values with standard error over 16 seeds.
29

---

As expected, this decreases significantly over time, suggesting the optimizers promotes less explo- ration over time. This makes sense; beneficial to explore early in training, when there time continue searching for better local or global optima, but later in training it is important to converge to the nearest optimum.
We also note the significantly different scales of noise which are applied to different environments, In particular, uses very small noise, suggesting amount of exploration. This may be one reason why OPEN performed similarly to handcrafted optimizers, and 'No Features'.
an environment which needs significant
F,5 Deepsea
Finally, we consider how the randomness used by OPEN changes during training in environments
of different sizes from Deep sea 771 63 This randommess 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 take in to receive large
goctor
reward. We look at randomness// / R;
as in Appendix F.4
332
$007
4:00
Figure 16: Mean normalized randomness applied by OPEN for different sized Deep Sea environments, over 128 seeds, with standard error. Notably, the scale for larger environments is bigger, such that OPEN applies more noise (i.e., more exploration) in larger environments despite having no awareness of the environment size.
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 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.
30

---

G Multi Environment Return Curves
Figure shows return curves the four MinAtar 62 environments included in 'multi-task' training 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 those which they can strongly outperform Adam 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, addition to consistently matching or beating the other learned baselines.
freeway
asterix
25
60
50
20
Return
Retur
30
RMSprop
20
Lion
VeLO
5
OpünviRL
No Features
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Frames
Frames
1e7
breakout
spaceinvaders
150
50
40
Return
50
10
0
0.0
0.2
0.0
0.2
0.4
0.6
0.8
1.0
Frames
1e7
Frames
Figure 17: RL training curves for the multi-task setting. We show mean return training with standard error, evaluated over 16 seeds.
31

---

H
Additional Gridworld Experiments
Similar Figure which looked how OPEN generalized to different agent sizes, we consider how OPEN adapts to different training lengths in Figure I8 In particular, we note that OPEN was only ever trained for gridworlds ('all", Appendix DD which ran for 3e7 training steps, so lengths outside of this are OOS tasks. We also consider that in [ the return of OPEN often started lower than other optimizers and increased only later training. As such, the lifetime conditioning (i.e., training proportion) of OPEN needs to be leveraged to be able to perform well at shorter training lengths. If OPEN did not make use this feature, may end training without having converged to the nearby optima, instead focusing exploration.
We find that OPEN able to beat Adam in all of the grids considered at the in-distribution training length. For OOS training lengths, OPEN generalizes better than Adam, and only in a couple of environments (rand dense, standard _aaee) does not outperform Adam in the shortest training length. For every other training length, OPEN performs better than Adam in every environment, further demonstrating its capability to generalize to OOS environments and training settings.
in distribution
Adarn Normllzed Relurn
1.10
1.05
1.00
2
3
6
Total Timestops
1e7
rand_dense
rand_sparse
rand_ jon__oon
1.5
16
1.
b)
1.0
standard maze
sixtean rooms
Isbyrinth
tour
Tosal
1:7
Figure 18: IQM of return achieved by OPEN, normalized by the return of Adam, against a range of different training lengths with 95% stratified bootstrap confidence intervals for 64 seeds. a) shows performance on the distribution which OPEN was trained on, and Adam was tuned in. b) shows performance OOS gridworlds, with the top row coming from [13] and the bottom row inspired by mazes from 67 all implemented by [66. We mark 3e7 timesteps as the in-iistritution training length for both OPEN and Adam.
32

---

I Ablation Details
L1 Feature Ablation Setup
We train 17 optimizers for each feature ablation, and evaluate performance 64 seeds per optimizer, using a shortened version of Breakout. We use the same PPO hyperparameters for Breakout as in Table besides the total timesteps T, which shorten to 10%. Each optimizer takes 60 minutes to train on one GPU, and we train total of 119 optimizers.
We train each optimizer for 250 generations, keeping all other MinAtar ES hyperparameters from Table g
1.2
Deep Sea Expanded Curve
We show the final return against size for the Deep Sea environment [71] 65 in Figure T6 for all
optimizers; both shared and separated, with stochasticity and without. The one without can 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 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 does not consistently reach maximum performance. In Deep Sea, the final return scale is between 0.99 and -0.01.
Stochastic Shared
Stochastic Separated
Ablated Shared
Ablated Separated
Return
0.4
0.2
0.0
10
20
40
50
Deep Sea Size
Figure 19: IQM of return against size of environment in Deep Sea, with stratified bootstrap 95% confidence intervals for 128 random seeds. The only optimizer which is able to generalize has separate parameters between the actor and critic and incorporates learned stochasticity into its update.
33

---

1.3 PQN Return Curve
Figure 20 shows the performance of OPEN and Adam when training a PQN agent 72 Asterix [62 63,. While OPEN outperforms Adam in this setting, we note that it exhibits similar behavior to the other return curves (Appendix G in having much lower return 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.
asterix (PQN)
40
35
30
25
Return
20
15
10
OPEN Adam
0.0
0.2
0.4
0.6
0.8
1.0
Frames
1e7
Figure 20: Training curves for PQN when optimised by OPEN and Adam asterix. We show mean return with standard error, over 64 seeds.
34

---

Experimental Compute
J.1 Runtimes
We include runtimes (inference) for our experiments with the different optimizers on 4 L40s GPUs. We note that, while OPEN takes longer to run compared to handcrafted optimizers, it offers significant speedup over VeLO and takes similar runtime Optim4RL and No Features in addition to significantly stronger performance. Importantly, this suggests the additional computation used in OPEN to calculate features, such as dormancy, does not translate to significant overhead compared to other learned optimizers.
Table 13: Runtime for each optimizer. We evaluate each over 16 seeds, on L40s GPUs. These take advantage of Jax's ability to parallelize over multiple GPUs (33.
RunTime (s)
Optimizer
freeway
breakout
spaceinvaders
asterix
ant
OPEN (single-tas)
138
87
94
137
453
OPEN (mulli-task)
154
98
109
148
N/A
Adam
53
58
105
333
RMSprop
84
38
52
97
322
Lion
83
44
50
96
303
VeLO
190
135
145
185
581
Optim4RL (single-taske
127
85
90
131
404
Optim4RL, (multi-task)
168
99
116
157
N/A
No Features (single-task
123
79
88
132
416
No Features (multi-task)
154
93
106
147
N/A
J.2 Training Compute
We used range of hardware for training: Nvidia A40s, Nvidia L40ses, Nvidia GeForce GTX 1080Tis, Nvidia GeForce RTX 2080Tis and Nvidia GeForce RTX 3080s. These are all part of an internal cluster. Whilst this makes it difficult to directly compare meta-training times per-experiment, find training took roughly between ~ [16, 32] GPU days of compute per optimizer in the single-task setting, and 60 GPU days of compute for multi-task training.
J.3
Total Compute
For all of our hyperparameter tuning experiments, we used total of 16 GPU days of compute, using Nvidia A40 GPUs.
For each learned optimizer 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 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 GPU days 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; total multi-task training used another 1 GPU year of compute, taking into account the omitted results.
To train distribution of gridworlds, used 7 GPU days of compute Nvidia GeForce RTX 2080 Tis.
Our ablation study used ~ 8 GPU days of compute, running Nvidia A40 GPUs.
Each optimizer took 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
negligible cost (on the order of seconds) in this paper.
Totaling the above, all experiments this paper used approximately 2. GPU years to run, though is difficult aggregate results properly due the variety of hardware used.
35

---

We ran 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 Optim4RL and No Features, we did not have any failed runs and so included all experiments in the paper.
36

---

K Code Repositories and Asset Licenses
Below include a full list of assets used in this paper, in addition 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
2pache-2.0
gymnax
gymnax
Apache-2.0. 2.0
brax
brax
Apache-2.0
l
learned optimization
45
learned optimization
Apache 2.0
groove (gridworlds)
66
groove
Apache-2.0
purejaxrl
15
purejaxrl
Apache-2.5. 2.0
optax
69
optax
Apache-2.0
rliable
70
rlibble
Apache-24
craftax
21
craltax
MIT
purejaxql (PQN)
72
purejaxql
Apache-2.0
37

---

NeurIPS 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 including
empirical results in range of environments (Section 6 Appendices 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 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.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
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 curricula used in multi-task settings, and discuss potential avenues for future work in this regard or in terms of considering other areas of RI. 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 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 vell-specifiiattio, 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 few datasets with 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, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails 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 faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
38

---

3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: We do not make theoretical claims as this is 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 cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofss provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Question:
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 (.2) This was all detailed in the main body of the paper, in Section We provide a full algorithm in Appendix B.I which can 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 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 dataset and/or model, the authors should describe the steps taken 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 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 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 large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
39

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We release all code used for OPEN in a publicly available repository
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips..o cc/
cublic/guidese httpe//guiiiee/eoeuuiisssssess Coodesubiisssooooooiiioo
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
7
www.ssssicniiiiyicc for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: We give moderate detail of the setup for training and testing in our results (Section 6 and ablation (Section 77. We go into significantly more detail of the hyperparam- eters in Appendix C and the architectures used in B We include details for our gridworld experiment Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
40

---

Justification: For all experiments in the main body of the paper, report 95% stratified bootstrap confidence intervals, using the implementation of rliable by Agarwal al. [70]. For return curves in Appendices F 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 [6
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, confi- dence 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, overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should 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.
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 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 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.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
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, 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 (c.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform. in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: This paper does not violate any part of the NeurlPS Code of Ethics; there are no potential harms caused by our research process, nor are there are any potential societal consequences.
41

---

Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: As foundational piece of research, this work has no direct paths to negative societal impacts besides that which can arise from improved reinforcement learning systems, improved machine learning systems more generally. We briefly discuss the general potential positive impacts of effective reinforcement learning systems in our introduction (Section
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (c.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeet of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We do not believe that OPEN has any potential for misuse, and so there is no need to describe safeguards for the models.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided relcasing unsafe images.
42

---

We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We take care to cite every environment used throughout the paper, with due care for who created original environment addition to who implemented the Jax-based version we used. We build our code upon the library from 45]. and are clear to mention this in appendix K All licenses can be found the relevant citation, and are included in Appendix K
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperssi.ppwiiiii thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors are encouraged reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We provide an anonymous link to the code and models used in this paper here and in appendix K
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/code/mmodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not use crowdsourcing experiments, nor does it use experi- ments with human subjects.
Guidelines:
43

---

The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation,
or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not crowdsourcing experiments, nor does it experi- ments with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such as the institution conducting the review.
44


Reviews:
[{"review_id": "KdWpwqGgZ8", "reviewer": "Reviewer_sKXo", "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.", "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?\n- Are there alternate methods for preventing parameter growth from the exploration term?", "limitations": "None", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This is a good paper on the application of learned optimization to RL. Of particular note:\n- I like the focus on learning capacity via single-task learning. This is a good sanity check for learned optimizers. \n- The out-of-distribution experiments are valuable. \n- Discussion of acceleration via parallelization in JAX is an extremely important engineering detail.\n- 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?\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear sKXo,\n\nWe would like to thank you for your positive review of our paper. In particular, we are grateful that you appreciate our **contribution based on conditioning directly on dormancy and other theoretically-grounded inputs** and recognize that this opens up huge potential opportunities in learned optimization. Below, we detail specific answers to the queries raised in your review.\n\n\n# Zero-Mean Updates Seems Counterintuitive\n\n> it seems quite different compared to eg momentum or other normalization strategies (such as adam-style normalization)\n\nIt is worth noting that OPEN **does use momentum** (as detailed in Appendix B.1, which discuss that different timesteps of momentum are included at input to OPEN). \n\n> It makes sense that adding random noise at each timestep could destabilize the model\n\n\nIt is important to highlight that **the instability is not caused by the inclusion of randomness**. We previously ran experiments with and without randomness, and found that the instability persisted until we incorporated the zero-meaning component. In fact, the ICLR24 rebuttals for Optim4RL are openly available online (I am not permitted to post a link per NeurIPS rebuttal guidelines) and **discuss a similar effect with their \"linear optim\" baseline, where they had NaN issues when learning in Brax** (quote: \"Note that LinearOptim in Ant fails due to NaN errors\"). They have little discussion of why this occurs, but we believe their issues arise due to the same problem as we discuss. \n\nBased on our analysis, we discuss the root cause of this issue in our paper (line 173); learned optimizers generally seem to produce large parameters. In domains with continuous action spaces, like Ant, the actor has an output parameter corresponding to the $\\log(\\sigma)$ of the actions (which gets exponentiated). Given that Brax also trains on long horizons, there is ample opportunity for parameters to grow and cause overflow errors.\n\n>  Better ablations on this term would be helpful\n\nRather than the interesting examples of normalization you propose, such as Adam-style (similar to Optim4RL and heavily constrains the model expressiveness) or contractive terms in the noise (which is insufficient, as the noise does not cause the instability), we have previously tested a regularization penalty in the fitness (i.e., $\\text{Fitness} = \\text{Final Return} - \\lambda||\\theta||_2$, where $\\theta$ is the agent's parameters). However, this introduces an additional, very expensive to tune hyperparameter ($\\lambda$) and frequently failed to escape the local minimum at $\\theta\\approx0$. Since these results are trivial, and tuning $\\lambda$ for a learned optimizer is beyond our capabilities, we **do not believe these results need to be included in the paper**. Instead, we have added mention of them in the main text to aid others and prevent dead-ends in their future research.\n\nWe found that zero-meaning did not significantly limit the flexibility of the update rule, does not require careful hand-design, and empirically had no impact in environments where exponentiation does not cause issues while making training possible in those it does. As such, we hope that you agree that **no additional ablations are needed on this term**, though we have now mentioned in the paper that this instability occurs even without the randomness.\n\n# Clarity in Eq(8)\n\n>The notation in eq (8) is a bit hard to parse\n\nWe appreciate that our current notation in eq(8) caused confusion; we wanted to keep our notation consistent with the other equations. Is it easier to interpret if we write equation 8 as $\\hat{u}_i^{\\text{actor,new}} := \\hat{u}_i^{\\text{actor}} + \\alpha_3 \\delta_i^{\\text{actor}} \\epsilon$? \n\n# Type of ES used in OPEN\n\n>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?\n\nWhat we call \"OpenAI ES\" is considered \"standard ES\" in many circles (and is widely used in learned algorithms/learned optimizers, e.g., [1,2,3]). We use the name from evosax to prevent ambiguity. We did run preliminary experiments with CMA-ES and a genetic algorithm, but found performance insufficient. \n\nWhile alternatives to ES, like PES/NRES, exist, **they introduce extra bias** which can cause issues. In particular, PES and NRES update the optimizer in an *online* fashion (i.e., during the innter-training loop), and are seen as alternatives to *truncated* ES, whereas OpenAI ES completes a full inner-training loop before updating. While PES/NRES are unbiased *in theory*, their implementation of updating the meta-parameters online introduces hysteresis effects *in practice*, since the state of a system depends on its history. Given how difficult learned optimization in RL has proven historically, we chose to avoid this additional source of bias which could introduce instability into the system. We found that, in practice, OpenAI ES performed very well and we did not require truncations as training was fast enough thanks to the speed of JAX and fast convergence of OPEN compared to other methods. Thus, we decided to pursue OpenAI ES in the publication. \n\n# Conclusion\n\nWe hope that we have addressed any concerns you might have had about our paper. If there is anything else we can respond to, please do let us know.\n\n---\n\n[1] Chris Lu et al., Discovered policy optimisation. 2022.\n\n[2] Luke Metz et al., VeLO: Training Versatile Learned Optimizers by Scaling Up. 2022.\n\n[3] Luke Metz et al., Tasks, stability, architecture, and compute: Training more effective learned optimizers, and using them to train themselves. 2020.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "krbhUFs1AI", "reviewer": "Reviewer_Fr37", "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.", "questions": "Can the authors say something about the computational cost of their method versus the other methods used in their evaluation?", "limitations": "Related to my question, maybe a small section about computational differences should be added to the main text.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Very clear and well-written paper. \n- Comprehensive Introduction and Related Work.\n- Strong evaluation using Adam-Normalized scores. \n- 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.\n\n- An evaluation on another RL algorithm or a more complex environment (Atari?) would strengthen the results. (The PureJaxRL Github Repo [3] also has DQN)\n\nSmall remarks ##:\n\n- Line 73: \"Evolution algorithms\" -> 'Evolutionary algorithms'\n- Line 75: \"This population is use\" -> 'This population is used'\n- Line 250: \"significantly outperform\" -> 'significantly outperforms'\n\n\n\n[1] Sokar et al: The Dormant Neuron Phenomenon in Deep Reinforcement Learning\n\n[2] Obando-Ceron et al: In value-based deep reinforcement learning, a pruned network is a good network\n\n[3] Lu et al: Discovered policy optimisation", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Fr37,\n\nThank you very much for your review. We are glad that you appreciate our **clear writing and strong evaluation** throughout the paper, and were happy to see this reflected in your recommendation of acceptance. Below, we detail answers to your queries.\n\n# Insufficiency of Dormancy\n\n> The use of ReLU dormancy is, in my opinion, not a viable metric for plasticity in RL.\n\nAs we cite in our paper, there exists a long list of work which use dormancy as a proxy measure for plasticity loss; e.g., [1,2,3,4], to name a few. This would suggest that, though dormancy and plasticity are not *directly* equivalent, dormancy is a useful and easily computed metric emblematic of plasticity. While the works you cite are certainly interesting, it is important to recognise that some of these findings (about e.g., the replay ratio) **do not apply in our setting**, where we have thus far focused on PPO. \n\nMore importantly, however, is the fact that we run our *own experiments* in the paper which demonstrate the **significant benefit of providing dormancy directly to the optimizer**; in our ablation (Figure 5), which is evaluated over a large number of optimizers, we find that dormancy was the **most significant individual feature** input to OPEN. In addition to causing a ~20% drop in final return, the percentage of dormant neurons tripled when dormancy was not included at input. Considering the full-scale experiments, OPEN also generally produced fewer dormant neurons despite optimizing *only for final return* (Appendix F1).\n\nThese results directly lead to 2 conclusions:\n- Including dormancy at input to the optimizer significantly improves final return.\n- Since OPEN (with awareness of dormancy) produced significantly less dormant neurons when optimizing for final return, dormancy must have a direct link to performance in this setting. \n\nWhile we believe, as grounded in literature, that the primary reason for these two effects comes down to the effect of plasticity, it may be that there are additional effects beyond our understanding that causes improved performance by reducing dormancy. However, our experiments make clear that **dormancy is an important input feature to OPEN which provides significant benefit**. \n\nBased on the above, we hope that you agree that there is a close link between dormancy and performance in our setting, and that there is existing literature beyond our work by which we can support the claims made in our paper.\n\n# Considering Harder Environments and Other Algorithms\n\n> evaluation on another RL algorithm or a more complex environment (Atari?) would strengthen the results\n\nIn response to your question, we have run experiments focusing on generalization from gridworlds to a **much more complex environment**. In our global response, we introduce an additional generalization experiment on Craftax-Classic [5], a recent reimplementation of Crafter in JAX. In this setting, we find that OPEN is able to **generalize significantly better than Adam** from the gridworlds (i.e., if we also transfer the Adam hyperparameters tuned in the gridworld). Also, the 0-shot generalization performance of OPEN is *only slightly worse* than an in-domain, finedtuned implementation of Adam.\n\nWe refer the reviewer to the global response and PDF for more details.\n\nUnfortunately, due to how long it takes to train a learned optimizer, as well as server issues, we have been unable to finish training OPEN with another RL algorithm in the allotted rebuttal time; **this experiment is therefore still ongoing**. \n\nWe are currently training OPEN with PQN, a vectorized version of DQN. After extensive finetuning, Adam + PQN achieves ~36.5 return on Asterix against ~34.9 return for OPEN (16 seeds). However, **we anticipate OPEN will eventually overtake Adam since the meta-training performance is still consistently increasing**. We will update you with a comment later in the week, pending the experiment's conclusion.\n\n# Computational Cost\n\n> Can the authors say something about the computational cost of their method versus the other methods used in their evaluation?\n\nWe include a very thorough breakdown of evaluation and training cost in appendix J. They key takeaway is that OPEN is significantly quicker than VeLO and not that much slower than Adam in test-time training, with the discrepancy shrinking in more expensive environments where the simulation becomes the bottleneck. In meta-training, OPEN generally converges faster than *No Features* or Optim4RL; we include details of how many training generations each one used in Appendix C.3.\n\nWe apologize for not having included sufficient reference to this in the main body of the text. We have added additional references in an updated manuscript, alongside a sentence explaining the computational differences between OPEN and other optimizers.\n\n---\n\n[1] Clare Lyle et al., Disentangling the causes of plasticity loss in neural networks, 2024.\n\n[2] Hojoon Lee et al., PLASTIC: Improving Input and Label Plasticity for Sample Efficient Reinforcement Learning. 2023.\n\n[3] Guozheng Ma et al., Revisiting Plasticity in Visual Reinforcement Learning: Data, Modules and Training Stages. 2023.\n\n[4] Zaheer Abbas et al., Loss of Plasticity in Continual Deep Reinforcement Learning. 2023.\n\n[5] Michael Matthews et al., Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning. 2024.\n\n[6] Matteo Gallici et al., Simplifying Deep Temporal Difference Learning. 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wazY7S28F0", "reviewer": "Reviewer_KVk8", "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.", "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.\n\nYour 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?\n\nYou 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: 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.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "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.\n\nClarity: 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.\n\nSignificance: 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.\n\nOriginality: 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.\n\nThey 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer KVk8,\n\nThank you so much for your very positive review and recommendation of strong acceptance to the conference. While you have **highlighted many positives of our work, such as the quality of our writing, experiments and ablations**, you do still discuss some weaknesses of our paper. We address these below.\n\n# Novelty of the work\n\n> I wish it was a little more clear how this work is novel (from the authors perspective)\n\nWe list below some key novel concepts we introduce in our paper, and refer to the general response for a more full list. We have now included a couple of sentences into the introduction clearly detailing them:\n\n- We use ES to learn to optimize *for RL*, which has only recently been made possible thanks to JAX acceleration.\n- We propose a way to design a black-box meta-learner in a *very unconstrained way* using input features. This lets the meta-learner *learn what's important* for the goal we care about, such as final return, rather than having regularization penalties for things such weight magnitude. Our method ensures we maximize the true objective, rather than an augmented one, while still accounting for other problems (like plasticity loss) as a byproduct.\n- We highlight a set of *theoretically grounded* features which can be combined to reduce the effect of some specific difficulties in RL. In addition to showing the importance of these features, via ablation, in learned optimization for RL, we believe that similar approaches can and should be applied to a plethora of learned algorithms for RL, and that our features should be generally universal.\n- We introduce a new meta-learned exploration technique which boosts exploration *directly through the learned optimizer*. We demonstrate that this is useful in a hard-exploration environment, DeepSea.\n- We introduce a way to stabilize learned updates in domains with continuous action spaces by zero-meaning the updates, which prevents overflow issues from exponentiation. This does not hurt empirical performance in already stable environments, and makes training possible in environments with continuous action spaces. According to the Optim4RL ICLR24 rebuttals (which can be found online), this was also an issue in a naive baseline they used and was one of the reasons they heavily constrained their update expression, limits its expressibility.\n- We are the first example of learned optimization in RL **which significantly outperforms all baselines in a range of domains**, including beating Adam consistently in generalization to different environments and to different architectures. In our general response and below, we provide even stronger results from our rebuttal experiments, which demonstrate very strong generalization from gridworlds to a harder environment (Craftax).\n\n# Gridworlds and Generalization\n\n\n> I would like to know if this generalizes to more complex tasks\n\nIn our paper, we do outline our hopes of more complex generalization on line 285. It is our belief that the combination of our multi-task experiments, which show that OPEN can fit to multiple harder and diverse environments, and the generalization experiments, which look at generalizing from easy gridworlds to significantly harder gridworlds, demonstrate that more complex generalization results will arise from larger scale training. \n\nAs part of our rebuttals we have considered transfer to *Craftax Classic* [1], a recent JAX reimplementation of Crafter. This is a **significantly harder environment than those tested in our paper**. Notably, since we are considering generalization rather than training directly in craftax, **we do not learn any new optimizers** for this experiment; instead, we evaluate the optimizer trained in section 6.4 of the paper on Craftax, 0-shot. We refer to the general response for full experimental details and results.\n\nIn short, we find that OPEN **generalizes significantly better than Adam** from the gridworlds (i.e., if we also transfer the Adam hyperparameters tuned in the gridworld). Additionally, the 0-shot generalization performance of OPEN is *only slightly worse than an in-domain, finedtuned implementation of Adam*.\n\nThis experiment suggests that OPEN is able to generalize very far out of distribution, and **offers the potential beginnings of truly automated RL** - in particular, meta-training can be done on simpler and quicker environments without the need for any finetuning on the hard/real-world environment!\n\n> is it possible to train it once on a wide variety of environments and then freeze it and use it for basically any environment\n\nAs an academic lab, it is beyond our capabilities to thoroughly explore the question of large-scale training. However, given the sheer levels of generalization we have seen from just our small-scale experiments, and the newfound capability to generalize well to Craftax-Classic, **it is our belief that OPEN could be trained once on a variety of environments and then be applied near-universally in RL**. We think that this would be a natural and highly impactful next step to OPEN and would highly recommend others with greater compute capabilities to take this forward, making use of our open source codebase. We have included a statement to this effect in the paper.\n\n# Training Speed\n\n> 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\n\nThank you for your suggestion about better communicating the speed of training with OPEN. We have included a sentence highlighting that OPEN does not drastically slow down training; in fact, since the main bottleneck in training RL is generally the environment interaction, it is likely that the slowdown from OPEN will decrease as it scales to more complex settings.\n\n---\n\n[1] Michael Matthews et al., Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning. 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cjL0plzhtL", "reviewer": "Reviewer_qjkS", "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.", "questions": "l.75 is use --> is used\n\nThe 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.\n\nAppendix E. Why is Ant being evaluated for more time steps than the other conditions?\n\nIt 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.\n\n------\n\n[1] D.M. Bossens, N.C. Townsend, A.J. Sobey (2019). Learning to learn with active adaptive perception, Neural Networks, \n115, 2019, Pages 30-49.", "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.\n\nIn 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.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\nThe 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.\n\nThe approach works well compared to state-of-the art and provides a scalable approach. \n\nUsing the approach, the exploration parameter is increased automatically for large environments, which is intuitively what you would expect.\n\nThe 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.\n\nThe review of directly related techniques is limited, even though I am aware that this is a niche domain.\n\nOne 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.\n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer qjkS,\n\nThank you for your in-depth review and recommendation of acceptance. We greatly appreciate that **you find our approach performant, scalable and novel**. We address your concerns below.\n\n# Lack of Empirical Detail\n\n>The techniques used for the empirical comparisons are not described in enough detail.\n\nWe appreciate your feedback, but wondered if you could please clarify which details in particular are missing from our experimental desciptions? In our work, we:\n- Explain each training environment and task setting thoroughly.\n- Discuss a range of baselines, with hyperparameters and tuning details (Appendix C).\n- Present results in a standardized framework (rliable), reporting 95% stratified bootstrap CIs over many seeds.\n- Provide conclusions in text which can be drawn directly from figures, without extraneous claims.\n- Include extra details in the appendix, including return curves (Appendix E, G), analysis (Appendix F) and more results (Appendix H).\n- Ablate each component of OPEN, considering the impact on final return and the agents' dormancy, with more details in Appendix I.\n- Ensure all sections are properly cross-referenced and adequately cited.\n\n# Review of Directly Related Techniques\n\n> The review of directly related techniques is limited, even though [...] this is a niche domain.\n\nPlease could you highlight what specifically is missing from our review of related techniques? We thank you for, and have included into our related work, your suggested reference on \"Active Adaptive Perception\". However, regarding the overarching point, we:\n- Include extensive background, covering evolution and optimization from both a handcrafted and learned perspective.\n- Have a detailed related work section introducing meta-learned algorithms in RL, such as LPO, LPG and MetaGenRL, and learned optimization in non-RL (e.g. VeLO, Lion). We also describe Optim4RL.\n- Dedicate a section to three difficulties of RL, with ample links to literature. This is expanded in Appendix A.\n- Motivate many design decisions in our optimizer from prior literature.\n- Compare OPEN against Optim4RL, VeLO and Lion, which are all \"directly related techniques\" involving learned optimization. We also compare against Adam and RMSProp, which are commonly used in RL.\n\n# Comparison to Parameter Space Noise\n\n> It is not clear how the paper compares to approaches for parameter noise\n\nDirectly 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.\n\nHowever, 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.\n\nInterestingly, Bigger, Better, Faster [1] actually removes noisy nets, which are similar to learnable stochasticity in OPEN, due to over-exploration. Our approach, which meta-learns how much noise to apply training, improves performance, **suggesting OPEN resolves this issue**.\n\nWe have now made clear in the paper the advantages of OPEN over parameter space noise.\n\n# Minor Contribution From Individual Elements\n\n>the contribution of each individual part is limited and the technique is a \"combination\" of techniques rather than a new technique altogether.\n\nWe 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:\n\n- Learning to optimize in RL *using evolution strategies*, which has only recently made feasible thanks to JAX acceleration.\n- 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.\n- Introducing *meta-learned exploration* directly in the optimizer, using learnable stochasticity.\n- Stabilizing training with zero-meaning. This does not excessively limit the expressiveness of the optimizer but enables meta-training in environments with continuous actions.\n\nWhile 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].\n\n# Limiting to Gridworlds\n\n>It would be interesting to see how the technique scales to high-dimensional problems\n\nIn response to your concerns regarding the gridworld experiments, we have added new results in a much harder domain: Craftax-Classic, a JAX reimplementation of Crafter. We provide results and extra detail in the general response. \n\nIn short, we demonstrate **better 0-shot generalization** than Adam from our distribution of gridworlds to Craftax, and **perform only slightly worse, 0-shot, than Adam when it is finetuned in domain**.\n\n# Specific Questions\n\nThank you for spotting a typo, we have updated the paper. For Ant, our hyperparameters based on prior work which runs in Brax ($5\\times10^7$) for longer than MinAtar ($1\\times10^7$). Training for more timesteps on Brax than MinAtar is common.\n\n---\n\n[1] Max Schwarzer et al., Bigger, Better, Faster: Human-level Atari with human-level efficiency. 2023", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Dear reviewers,

# Thank you

We are very grateful to have received four high quality reviews, and appreciate the time and effort you all spent with our paper. We are glad that there was consistency among reviewers that our paper was easily understood and presented strong results, and that **every reviewer recommended acceptance to the conference**. There were, however, some points raised by multiple reviewers. These are addressed below as supplement to individual responses.

We have color-coded our general response for convenience:

$\color{red}{R1: sKXo}$

$\color{gold}{R2: Fr37}$

$\color{green}{R3: KVk8}$

$\color{pink}{R4: qjkS}$

# Harder Environments

Reviewers $\color{gold}{Fr37}$, $\color{green}{KVk8}$ and $\color{pink}{qjkS}$ suggested considering more complex environments than gridworlds. Therefore, we examine the generalization of OPEN to a **significantly harder environment**: Craftax-Classic [1], a JAX reimplementation of Crafter.

We consider three optimizers:
1. OPEN, learned on the gridworld distribution from section 6.4, transferring *0-shot* to Craftax-Classic. **We do not retrain OPEN for this experiment.**
2. Adam, using the hyperparameters tuned for the gridworlds in section 6.4, transferring 0-shot.
3. Adam, where we tune the learning rate in-domain and use other hyperparameters tuned for Craftax from [1]. This gives a rough idea of the ceiling on optimizer performance, as Adam is *finetuned directly on the task*.

We outline results and all hyperparameters, many of which come from [1], in the PDF. All optimizers are run on 32 seeds for $3\times10^7$ timesteps. In particular, we find:
- **OPEN drastically outperforms Adam when both generalize** from the much simpler gridworld environment distribution.
- OPEN's 0-shot transfer is only slightly worse than Adam finetuned directly in domain (expensive).

This result suggests OPEN generalizes **even better than was originally demonstrated in the paper**, and lays the foundation for truly automated RL. In particular, we show that OPEN can be trained on quick and simple environments, before transferring to more complex environments **with no additional samples or hyperparameter tuning**. We have included this result, which significantly strengthens our original claims, into our updated paper.

# Additional Algorithms

$\color{gold}{Fr37}$ asked whether OPEN worked with algorithms other than PPO. Therefore, we are *currently* training OPEN with PQN [2], a variant of DQN designed for vectorized sampling, in Asterix to explore how our approach generalizes between algorithms.

Unfortunately, due to the time taken to train a learned optimizer from scratch, and a server crash during rebuttals, our experiment has not finished within the allotted rebuttal time. We will continue to run this, and **will update OpenReview pending the conclusion of our experiment**. However, we present some initial, indicative results *from the midst of training* below.

Adam receives a return of ~36.5 (16 seeds) after finetuning, whereas OPEN *currently* achieves ~34.9 return (16 seeds). Importantly, **the meta-training performance is still consistently increasing, so we believe that OPEN will overtake Adam**. We see this as a very positive sign that OPEN is general enough for a range of RL algorithms. We will add this to our paper when the experiment has finished.

# List of Novelty

Both $\color{green}{KVk8}$ and $\color{pink}{qjkS}$ raised concerns about novelty in OPEN. However, we believe there is *significant* novelty to our work. It is worth noting that only one other work has attempted learned optimization in RL, and they fail to significantly outperform all baselines in any of their tests. On the contrary, OPEN **consistently outperforms all baselines in a range of domains** with a *completely different method*. We have consolidated a list of some of the key components introduced in OPEN, and have added some sentences to our introduction to emphasise this novelty:

- We use ES to learn to optimize for RL. This is only recently possible thanks to JAX acceleration.
- We propose a very unconstrained way to design a black-box meta-learner based on input features. This enables learning on only the ultimate goal, such as final return, rather than having handcrafted and hyperparameter-dependent regularization penalties for things such as dormancy or weight magnitude. As a result, other possible problems, such as plasticity loss, are dealt with as a byproduct. This contrasts with Optim4RL, which relies on excessive structure that limits performance.
- We highlight theoretically grounded features which can be combined to tackle a number of issues in optimization for RL. Our features are thoroughly ablated.
- We develop learnable stochasticity, a new and performant meta-learned exploration technique which parallels a learned version of parameter space noise (with some implementation details). This interacts with the different inputs of OPEN (e.g. dormancy, training proportion) to produce complex stochastic behaviour.
- We introduce a way to stabilize learned updates in domains with continuous action spaces by zero-meaning the updates, to prevent overflow issues from exponentiation. According to the Optim4RL ICLR24 rebuttals (which can be found online), this was an issue in one of their naive baselines too, and was one of the reasons they heavily constrained their update expression. Our method resolves this problem.
- We are the first example of learned optimization in RL which is able to **significantly and consistently outperform handcrafted baselines**, including beating Adam in generalization to different environments and to different architectures. This is bolstered by our additional experiments in generalization to Craftax and training with PQN.

---

[1] Michael Matthews et al., Craftax: A Lightning-Fast Benchmark for Open-Ended Reinforcement Learning. 2024.

[2] Matteo Gallici et al., Simplifying Deep Temporal Difference Learning. 2024.

Meta-review: 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.
Decision: Accept (spotlight)