Discovered Policy Optimisation
Abstract
Tremendous progress has been made in reinforcement learning (RL) over the past decade.
Most of these advancements came through the continual development of new algorithms, which were designed using a combination of mathematical derivations, intuitions, and experimentation.
Such an approach of creating algorithms manually is limited by human understanding and ingenuity.
In contrast, the meta-learning provides a toolkit for automatic machine learning method optimisation, potentially addressing this ﬂaw.
However, black-box approaches which attempt to discover RL algorithms with minimal prior structure have thus far not outperformed existing hand-crafted algorithms.
Mirror Learning , which includes RL algorithms, such as PPO, offers a potential middle-ground starting point: while every method in this framework comes with theoretical guarantees, components that differentiate them are subject to design.
In this paper we explore the Mirror Learning space by meta-learning a “drift” function.
We refer to the immediate result as Learnt Policy Optimisation (LPO).
By analysing LPO we gain original insights into policy optimisation which we use to formulate a novel, closed-form RL algorithm, Discovered Policy Optimisation (DPO).
Our experiments in Brax environments conﬁrm state-of-the-art performance of LPO and DPO, as well as their transfer to unseen settings.
Introduction
Recent advancements in deep learning have allowed reinforcement learning algorithms [35, RL] to successfully tackle large-scale problems [36, 34].
As a result, great efforts have been put into designing methods that are capable of training a neural-network policies in increasingly more complex tasks [33, 31, 24, 9].
Among the most practical such algorithms are TRPO [31] and PPO [32] which are known for their performance and stability [2].
Nevertheless, although these research threads have delivered a handful of successful techniques, their design relies on concepts handcrafted by humans, rather than discovered in a learning process .
As a possible consequence, these methods often suffer from various ﬂaws, such as the brittleness to hyperparameter settings [31, 12], and a lack of robustness guarantees.
The most promising alternative approach, algorithm discovery , thus far has been a “tough nut to crack”.
Popular approaches in meta-RL [30, 7, 4] are unable to generalise to tasks that lie outside of their training distribution.
Alternatively, many approaches that attempt to meta-learn more general algorithms [25, 16] fail to outperform existing handcrafted algorithms and lack theoretical guarantees.
Recently, Mirror Learning [18], a new theoretical framework, introduced an inﬁnite space of provably correct algorithms, all of which share the same template.
In a nutshell, a Mirror Learning algorithm is deﬁned by four attributes, but in this work we focus on the drift function .
A drift function guides the agent’s update, usually by penalising large changes.
Any Mirror Learning algorithm provably achieves monotonic improvement of the return, and converges to an optimal policy [18].
Popular RL methods such as TRPO [31] and PPO [32] are instances of this framework.
In this paper, we use meta-learning to discover a new state-of-the-art (SOTA) RL algorithm within the Mirror Learning space.
Our algorithm thus inherits theoretical convergence guarantees by construction.
Speciﬁcally, we parameterise a drift function with a neural network, which we then meta-train using evolution strategies [29, ES].
The outcome of this meta-training is a speciﬁc Mirror
Learning algorithm which we name Learnt Policy Optimisation (LPO).
While having a neural network representation of a novel, high-performing drift function is a great ﬁrst step, our next goal is to understand the relevant algorithmic features of this drift function.
Out analysis reveals that LPO’s drift discovered, for example, optimism about actions that scored low rewards in the past—a feature we refere to as rollback .
Building upon these insights we propose a new, closed-form algorithm which we name — Discovered Policy Optimisation (DPO).
We evaluate LPO and DPO in the Brax [8] continuous control environments, where they obtain superior performance compared to PPO.
Importantly, both LPO and DPO generalise to environments that were not used for training LPO.
To our knowledge, DPO is the ﬁrst theoretically-sound, scalable deep RL algorithm that was discovered via meta-learning.
Related Work
Over the last few years, researchers have put signiﬁcant effort into designing and developing algorithmic improvements in reinforcement learning.
Fujimoto et al.
[9] combine DDPG policy training with estimates of pessimistic Bellman targets from a separate critic.
Hsu et al.
[15] stabilise the, previously unsuccessful [32], KL-penalised version of PPO and improve its robustness through novel policy design choices.
Haarnoja et al.
[13] introduce a mechanism that automatically adjusts the temperature parameter of the entropy bonus in SAC.
However, none of these hand-crafted efforts succeeds in fully mitigating common RL pathologies, such as sensitivity to hyperparameter choices and lack of domain generalisation [4].
This motivates radically expanding the RL algorithm search space through automated means [27].
Popular approaches in meta-RL have shown that agents can learn to quickly adapt over a pre-speciﬁed distribution of tasks.
RL 2 equips a learning agent with a recurrent neural network that retains state across episode boundaries to adapt the agent’s behaviour to the current environment [4].
Similarly, a
MAML agent meta-learns policy parameters which can adapt to a range of tasks with a few steps of gradient descent [7].
However, both RL 2 and MAML usually only meta-learn across narrow domains and are not expected to generalise well to truly unseen environments.
Xu et al.
[40] introduce an actor-critic method that adjusts its hyperparameters online using metagradients that are updated with every few inner iterations.
Similarly, STAC [42] uses implementation techniques from IMPALA [5] and auxiliary loss-guided meta-parameter tuning to further improve on this approach.
Such advances have inspired extending meta-gradient RL techniques to more ambitious objectives, including the discovery of algorithms ab initio .
Notably, Oh et al.
[25] succeeded in meta-learning an
RL algorithm, LPG, that can solve simple tasks efﬁciently without explicitly relying on concepts such as value functions and policy gradients.
Similarly, Evolved Policy Gradients [14, EPG] meta-trains a policy loss network function with Evolution Strategies [29, ES].
Although EPG surpasses PPO in average performance, it suffers from much larger variance [14] and is not expected to perform well on environments with dynamics that differ greatly from the training distribution.
MetaGenRL [17], instead, meta-learns the loss function for deterministic policies which are inherently less affected by estimators’ variance [33].
MetaGenRL, however, fails to improve upon DDPG [21] in terms of performance, despite building up on it.
Neither EPG nor MetaGenRL have resulted in the discovery of novel analytical RL algorithms, perhaps due to the limited interpretability of the loss functions learnt.
Lastly, Co-Reyes et al.
[3], Garau et al.
[10] and Alet et al.
[1] discover and improve standard
RL conventions by evolving, symbolically, algorithms represented as graphs, which leads to improved performance in simple tasks.
However, none of those trained-from-scratch methods inherit correctness guarantees, limiting our certainty of the generality of their abilities.
In contrast, our method, LPO, is meta-developed in a Mirror Learning space [18], where every algorithm is guaranteed convergence to an optimal policy.
As a result to this construction, meta-training of LPO is easier than that of methods that learn “from scratch”, and achieves great performance across environments.
Furthermore, thanks to the clear meta-structure of Mirror Learning, LPO is interpretable, and lets us discover new learning strategies.
This lets us introduce DPO—an efﬁcient algorithm with a closed-form formulation that exploits the discovered learning concepts.
Background
In this section, we introduce the essential concepts required to comprehend our contribution—the
RL and meta-RL problem formulations, as well as the Mirror Learning and Evolution Strategies frameworks for solving them.
3.1 Reinforcement Learning
Formulation We formulate the reinforcement learning (RL) problem as a Markov decision process (MDP) [35] represented by a tuple (cid:104)S , A , R, P, γ, d (cid:105) which deﬁnes the experience of a learning agent as follows: at time step t ∈ N , the agent is at state s t ∈ S (where s 0 ∼ d ) and takes an action a t ∈ A according to its stochastic policy π ( ·| s t ) , which is a member of the policy space Π .
The environment then emits the reward R ( s t , a t ) and transits to the next state s t +1 drawn from the transition function, s t +1 ∼ P ( ·| s t , a t ) .
The agent aims to maximise the expected value of the total discounted return,
[Equation]
The agent guides its learning process with value functions that evaluate the expected return conditioned on states or state-action pairs
[Equation]
function),
[Equation]
the state-action value function) .
The function that the agent is concerned about most is the advantage function , which computes relative values of actions at different states,
[Equation]
Policy Optimisation
In fact, by updating its policy simply to maximise the advantage function at every state, the agent is guaranteed to improve its policy, η ( π new ) ≥ η ( π old ) [35].
This fact, although requiring a maximisation operation that is intractable in large state-space settings tackled by deep
RL (where the policy π θ is parameterised by weights θ of a neural network), has inspired a range of algorithms that perform it approximately.
For example, A2C [24] updates the policy by a step of policy gradient (PG) ascent
[Equation]
estimated from a batch of B transitions.
Nevertheless, such simple adoptions of generalized policy iteration [35, GPI] suffer from large variance and instability [43, 33, 32].
Hence, methods that constrain (either explicitly or implicitly) the policy update size are preferred [31].
Among the most popular, as well as successful ones, is Proximal Policy Optimization [32, PPO], inspired by trust region learning [31], which updates its policy by maximising the PPO-clip objective,
[Equation]
, (4) where the clip ( · , 1 ± (cid:15) ) operator clips (if necessary) the input so that it stays within [1 − (cid:15), 1 + (cid:15) ] interval.
In deep RL, the maximisation oracle in Equation (4) is approximated by a few steps of gradient ascent on policy parameters.
Meta-RL The above approaches to policy optimisation rely on human-possessed knowledge, and thus are limited by humans’ understanding of the problem.
The goal of meta-RL is to instead optimise the learning algorithm using machine learning.
Formally, suppose that an RL algorithm alg φ , parameterised by φ , trains an agent for K iterations.
Meta-RL aims to ﬁnd the meta-parameter φ = φ ∗ such that the expected return of the output policy, E [ η ( π K ) | alg φ ] , is maximised.
3.2 Mirror Learning
A Mirror Learning agent [18], in addition to value functions, has access to the following operators: the drift function D π k ( π | s ) which, intuitively, evaluates the signiﬁcance of change from policy π k to π at state s ; the neighbourhood operator N ( π k ) which forms a region around the policy π k ; as well as sampling and drift distributions β π k ( s ) and ν ππ k ( s ) over states.
With these deﬁned, a Mirror
Learning algorithm updates an agent’s policy by maximising the mirror objective
[Equation]
If, for all policies π and π k , the drift function satisﬁes the following conditions:
It is non-negative everywhere and zero at identity D π k
( π | s ) ≥ D π k ( π k | s )
= 0 , 2.
Its gradient with respect to π is zero at π = π k ,
then the Mirror Learning algorithm attains the monotonic improvement property, η ( π k +1 ) ≥ η ( π k ) , and converges to the optimal return, η ( π k ) →
η ( π ∗ ) , as k → ∞ [18].
A Mirror Learning agent can be implemented in practice by specifying functional forms of the drift function and neighbourhood operator, and parameterising the policy of the agent with a neural network, π θ .
As such, the agent approximates the objective in Equation (5) by sample averages, and maximises it with an optimisation method, like gradient ascent.
PPO is a valid instance of Mirror Learning, with the drift function:
[Equation]
While it is possible to explicitly constrain the neighbourhood of policy update [31], some algorithms do it implicitly.
For example, as maximisation oracle of PPO (see Equation (4)) has a form of N steps of gradient ascent with learning rate α and gradient clipping threshold c , it implicitly employs a neighbourhood of an Euclidean ball or radius Nαc around θ k .
Different Mirror Learning algorithms can differ in multiple aspects such as sample complexity and wall-clock time efﬁciency [18].
Depending on the setting, different properties may be desirable.
In this paper, we optimise for the return of the K th iterate, η ( π K ) .
3.3 Evolution Strategies
Evolution Strategies [28, 29, ES] is a backpropagation-free approach to function optimisation.
At their core lies the following identity, which holds for any continuously differentiable function F of φ , and any positive scalar σ
[Equation]
(7) where N (0 , I ) denotes the standard multivariate normal distribution.
By taking the limit σ → 0 , the gradient on the left-hand side recovers the gradient of ∇ φ F ( φ ) .
These facts inspire an approach of optimising F with respect to φ without estimating gradients with backpropagation—for a random sample (cid:15) 1 , . . .
, (cid:15) n ∼ N (0 , I ) , the vector 1 nσ (cid:80) ni
=1 F ( φ + σ(cid:15) i ) (cid:15) i is an unbiased gradient estimate.
To reduce variance of this estimator, antithetic sampling is commonly used [26].
In the context of meta-RL, where φ is the meta-parameter of an RL algorithm alg
φ , the role of F ( φ ) is played by the average return after the training, F ( φ ) =
E [ η ( π K ) | φ ] .
As oppose to the meta-gradient approaches described in Section 2, ES does not require backpropagation of the gradient through the whole training episode—a cumbersome procedure which, often approximated by the truncated backpropagation, introduces bias [38, 39, 25, 6, 23].
[Figure]
Methods
Our overall approach is to meta-learn a drift function to perform policy optimisation over a ﬁxed episode length K .
Hence, our meta-objective is the expected ﬁnal return,
[Equation]
4.1 Drift Function Network
The drift function that we learn takes form D π k ( π | s ) =
E a ∼ π k
[ f φ ( x ) | s ] , where f φ ( x ) is a fullyconnected neural network parameterised by φ .
Our drift network is a function of the probability ratio between a candidate and the old policy, r = π
( a | s ) /π
k ( a | s ) , and of the advantage
A =
A π k ( s, a ) (which we assume to be normalised across each batch).
To ease learning complicated mappings, we include non-linear transformations of these arguments (we perform ablations in E), ultimately forming the following input x
[Equation]
[Equation] [Equation]
r )
[Equation] [Equation]
[Equation]
In order to guarantee that the neural network is a valid drift function, it must be the case that f φ ( x r,A ) = 0 and ∇ r f φ ( x r,A ) =
whenever r = 1 , and f φ ( x r,A ) ≥ 0 everywhere.
As in our model, x r,A = 0 whenever r = 1 , the former condition is guaranteed by excluding bias terms from the network architecture.
To meet the latter two conditions, we apply the ReLU activation at the last layer with a slight shift, x (cid:55)→ ReLU ( x − ξ ) , where ξ = 10 − 6 .
To alleviate the difﬁculty of the meta-training, the main variant of our implementation initialises the drift function near the PPO one.
That is, before we pass the last layer’s output to shifted ReLU, we add to it the PPO drift function (more precisely, the input of the expectation in Equation 6).
We have found that this operation leads to a better performance and generalisation across different environments.
Furthermore, such a setting directly tests for the optimality of the drift function of
PPO.
If PPO was optimal, the network could simply learn the zero mapping.
In principle, the drift function is not restricted to operate with probability ratios and advantages only; it could also accept other arguments, such as algorithm hyperparameters, statistics measuring the progress of training, or even task information.
However, many of these vary greatly in the dimensionality or scale across different instances or types of environments.
Furthermore, a large number of arguments would impede the analysis of the learnt drift function, and comparison to PPO, which are also goals of our work.
Hence, we work with simple and transferable ratios and normalised advantage estimates.
Nevertheless, for speciﬁc applications, it may be beneﬁcial to consider other, possibly domain-speciﬁc, information.
4.2 Meta-Training the Drift Function Network
The meta-objective we optimise is the performance of the learner policy at the end of training:
F ( φ ) = E (cid:2)
η ( π θ K ) | alg φ ] , where θ K is the K th (last) iterate of the RL training under the Mirror
[Figure]
[Figure]
Learning algorithm alg
φ .
The expectation is taken over the randomness of the initial parameter θand stochasticity of the environment.
We solve this problem using ES with antithetic sampling.
At each generation (outer loop iteration)
, we sample a batch of perturbations of φ , initialise the policy parameters θ 0 , and then train the policy under alg φ , using the drift function’s parameter φ , for K iterations.
At the end of the inner-loop training, we estimate the return of the ﬁnal policy φ θ K , and use it to estimate the gradient of ∇ φ F ( φ ) as in Equation (7).
We meta-learn the drift function and evaluate policies trained by it in the Brax [8] physics simulator environments.
Brax is designed to take advantage of parallel computation on accelerators, allowing us to roll out thousands of episodes in parallel, and train entire RL agents within minutes.
Thus, it is well suited for this type of optimisation problem.
Furthermore, we increase its efﬁciency by vectorising the policy optimisation algorithm itself , which lets us train hundreds of RL agents per minute using accelerators.
We implement our method on top of the Brax version of PPO, which provides a
Mirror Learning-friendly code template, keeping the policy architecture and training hyperparameters unchanged.
For meta-training we use both evosax [19] and the Learned_optimization [22] libraries.
For full details of meta-training see Appendix A.
Empirical Studies
We consider two different meta-training setups.
First, as described in Subsection 5.1 we attempt to learn a drift function completely from scratch to investigate how similar it is to existing algorithms like PPO.
Second, in Subsection 5.2 we ask whether we can learn a drift function that successfully generalises to multiple environments, if it is initialised near PPO.
5.1 Learning drift functions from scratch
In this setting, f φ is a neural network with two hidden layers of size 256 and a ReLU activation function.
We meta-train it across 5 Brax environments.
We name the resulting algorithm LPO-Zero, and visualise it in Figure 2.
Interestingly, LPO-Zero appears to have learnt a few PPO-like features, as can be observed on Figure 2.
For example, it appears to have learnt to clip the update incentive at a speciﬁc ratio threshold, much like PPO; however, it only does so for negative advantages.
Nevertheless, LPO-Zero largely underperforms with respect to LPO, and possibly requires much more training to catch up.
We present ﬁgures with performance evaluation of LPO-Zero in Appendix B.
5.2 Learning with the PPO Initialisation
In this setting, f φ is a small neural network, with a single hidden layer with 128 neurons, with bias terms removed, and a tanh activation function.
Here, we add PPO to the output of the last hidden layer, before passing it to the shifted ReLU,
[Equation]
where ˜ f φ ( x r,A ) is the output of the last hidden layer of the drift network.
As such, the resulting drift function similar to that of PPO at initialisation.
Surprisingly, we have found that meta-training in a single environment is sufﬁcient to generate drift functions whose abilities transfer to unseen tasks.
Moreover, we found that the learnt drifts generally display similar characteristics.
For readability, we chose the drift function that was trained on Ant, whose induced algorithm we refer to as Learnt Policy Optimisation (LPO).
Visualisations and results for drift functions trained on other environments can be found in the Appendix C.
The results on Figure 6 show that LPO, trained only on Ant, outperforms PPO in unseen environments.
Furthermore, the Brax PPO implementation uses different hyperparameters, such as the number of update epochs and the total number of timesteps, for each of the tasks.
This means that LPO, which was trained on Ant with hyperparameters associated to it, is robust not only against new environments, but also against new hyperparameters.
We visualise the derivative of the LPO loss in Figure 3, which enables us to derive an analytical version of it in Section 6.
Analysis of LPO
In this section, we analyse the two key features that are consistently learnt and contribute most to LPO’s performance.
We then interpret their effect on policy entropy , and the update asymmetry discovered by LPO, through which it differs largely from PPO (recall Figure 3 for visualisation).
Rollback for negative advantage.
In the bottom-left quadrant of the heat map, which corresponds to A < 0 and r < 1 (negative advantage and decrease in action probability), we observe that the ratio derivative of the LPO objective is positive in a large region, roughly corresponding to r < 1 − (cid:15) .
This implies that actions which fall into this quadrant, although seemingly not appealing, are encouraged to be taken by the agent, which can be interpreted as a form of rollback [37].
Hence, LPO learns to decrease r down to 1 − (cid:15) , but unlike PPO, encourages r to stay precisely around that value.
By doing so, LPO prevents the agent from giving up on actions that appear poor at the moment, and encourages it to keep exploring them at a moderate frequency.
[Figure]
Cautious optimism for positive advantage.
The upper-right quadrant corresponds to A > 0 and r > 1 (positive advantage and increase in action probability), which is induced by actions that seem the most appealing to update to.
Nevertheless, LPO is cautious in doing so, gradually decreasing the pace of its update towards them, and eventually abstaining from chasing the most extreme advantage values—these may come from critic errors.
We would like to highlight that this view of LPO on positive advantages in much more sophisticated than that of PPO, which simply removes any incentive from updating towards actions with r > 1 + (cid:15) , and thus can be viewed as optimistic relative to PPO.
Implicit entropy maximisation.
Together, these two central features of LPO encourage the agent to spread its policy probability mass moderately over all actions, thus leading to larger entropy and allowing for richer exploration.
Thus, LPO has implicitly discovered entropy maximisation , which we demonstrate in Figure 4.
We would like to highlight that LPO achieved this with its drift function only, without artiﬁcially augmenting the original RL objective with an entropy bonus [11, 12]
Update asymmetry.
LPO learns asymmetric features that respect a natural asymmetry of behaviour change in RL: increasing r for positive advantage A may encourage exploration of a newly-found action or strengthen a dominant action, whereas decreasing r for negative A will always discourage exploration of that action and strengthen a dominant action.
In this context, the two discussed features of LPO make it completely unlike PPO, which clips the update incentives symmetrically around the origin.
Secondary Features.
LPO, but not LPO-Zero, appears to consistently learn objectives with gradient spikes around r = 1 in the upper left and lower right quadrants.
Nevertheless, adding them to our analytic model of LPO did not improve performance.
We speculate, therefore, that these spikes are mostly artifacts of the network parameterisation.
Discovered Policy Optimisation: A New RL Algorithm Inspired By LPO In this
section, building upon concepts that LPO has discovered, we introduce a novel algorithm—
Discovered Policy Optimization (DPO).
7.1 The Discovered Drift Function Model
Combining the key features identiﬁed in Section 6, we construct a closed-form model of LPO that can easily be implemented with just a few lines of code on top of an existing PPO implementation.
We name the new algorithm Discovered Policy Optimisation (DPO) because we have not derived it—it was instead discovered in the meta-learning process.
DPO is a Mirror Learning algorithm, with a drift function that takes different functional forms, depending on the sign of advantage A , as dictated by the update asymmetry principle from the previous section.
Speciﬁcally, we have found
[Figure]
that the (parameter-free) drift function
[Equation]
faithfully reproduces the key features of LPO (cautious optimism and rollback) for appropriate constants α = 2 , β = 0 .
6 (see Appendix D for veriﬁcation of the drift conditions).
We visualise DPO in Figure 5 and note that even the “crossing-over” of gradient slices of LPO on Figure 3 is faithfully reproduced.
7.2 Results
We compare DPO to PPO and LPO on a variety of Brax environments in Figure 6.
We use the PPO implementation provided by Brax, with an addition of advantage normalisation as we observed it to improve the method’s performance across the majority of the environments.
Our methods also use this implementation technique.
We further demonstrate on out-of-distribution performance by evaluating performance on the Minatar environments [20, 41] in Figure 7.
While evaluating DPO, similarly to LPO, we do not re-tune any hyperparameters that were originally selected for PPO in Brax.
The results on Figure 6 show that DPO matches the performance of LPO and outperforms PPO on the evaluated environments, despite being a two-line analytic model of LPO based on two key features .
This enables RL practitioners to implement DPO as easily as PPO with a performance on par with our best learnt drift function.
[Figure]
Conclusion
In this paper, we approached the problem of algorithm discovery by restricting our meta-learning to the space of valid Mirror Learning algorithms.
Speciﬁcally, we optimised a drift function parameterised by a neural network, which we trained with Evolution Strategies .
We consider this work to be an example of the new, promising paradigm of RL algorithm discovery.
Namely, our strategy was to develop a high-performing RL algorithm by combining theoretical insights with large-scale computational techniques.
As a result of the training, we obtained a theoretically sound method that we named Learnt Policy Optimisation (LPO), which outperforms a state-of-the-art baseline (PPO) in unseen environments, and with unseen hyperparameter settings.
After analysing the learned features discovered by LPO, we introduced Discovered Policy Optimisation (DPO)—a closed-form approximation to LPO.
Our experimental results show that DPO matches LPO in performance and robustness to hyperparameter settings.
However, a possible weakness of this approach is that it could overﬁt to speciﬁc code-level details of the implementation.
For example, LPO is built on Brax’s
PPO implementation, which makes numerous design decisions to optimise for wall-clock time rather than sample efﬁciency.
In the future, we plan to expand the variety of inputs to the learnt drift function, as well as to parameterise and meta-learn other attributes of Mirror Learning.
We expect these advancements to provide more insights into policy optimisation, ultimately resulting in more robust and better performing RL algorithms.
Acknowledgments and Disclosure of Funding
Christian Schroeder de Witt is sponsored by the Cooperative AI Foundation.
Compute for this project was partially run on Oxford’s Advanced Research Cluster (ARC).
This work used the Cirrus
UK National Tier-2 HPC Service at EPCC funded by the University of Edinburgh and EPSRC (EP/P020267/1).
This work was also supported by an Oracle for Research Cloud Grant (19158657).
A Meta-Training Details A.1 LPO-Zero
Our LPO-zero implementation was implemented on top of the learned_optimization library [22].
The drift function is parameterised by a one layer fully connected network with 1 hidden layer and 256 hidden units.
Meta-training is done in a distributed fashion using batched, async metaupdates across 350 workers each of which with one TPUv4i accelerator.
On a centralized learner process we accumulate gradients from these workers until 350 gradients are computed (using a single perturbation from an antithetic ES based gradient estimator).
Once this number is reached, we perform one outer-iteration with Adam using a learning rate of 0.006.
In this experiment, we meta-train over a uniform mixture of the ant, walker2d, halfcheetah and fetch environments.
We take the default hparams for PPO from Brax for each implementation except for the number of epochs trained which we set to 183 to match what was done with the ant environment.
In each worker, for a particular environment, we perform a full PPO training for both a positive, and negative perturbation of the underlying meta-parameters.
At the end of each training, we evaluate 10240 rollouts on the environment with the resulting policy and use these as our ﬁtness function.
Meta-training was done over the course of 2 days and performed approximately 400 outer-updates.
We ﬁnd though performance still increases with increased meta-training time.
A.2 LPO
Our LPO implementation was implemented on top of the evosax library [19].
The drift function is parameterised by a one layer fully-connected network with 1 hidden layer and 128 hidden units.
Meta-training was only done on a single machine with 4 V100 GPU’s with synchronous updates.
Meta-training was done over the course of 2 days and performed approximately 700 outer-updates.
We ﬁnd though performance still increases with increased meta-training time.
[Table] [Table]
[Figure]
D DPO Drift Veriﬁcation
The DPO drift function is given by
[Equation]
The ﬁrst condition for a valid drift is that f be non-negative everywhere, which trivially holds since
ReLU ( x ) ≥ 0 for all x ∈ R .
The second condition is that f be zero at π = π old .
Now r = π/π old = 1 implies r − 1 = 0 and log r = 0 , which combined with tanh(0)
=
imply that f = 0 as required.
The ﬁnal condition is that the gradient of f with respect to π be zero at π = π old .
This is equivalent to having zero gradient with respect to r = π/π old at r = 1 since the gradients are equal up to a constant.
Now writing
[Equation]
and
[Equation]
log
[Equation]
r )
[Equation]
β ) for A ≥ 0 and A < 0 respectively, we have
[Equation]
and
[Equation]
which both evaluate to 0 at r = 1 , since cosh(0)
= 1 .
This implies for A ≥ 0 that
[Equation]
at r = 1 and for A < 0 that
[Equation]
at r
[Figure]