Discovered Policy Optimisation
Abstract
The last decade has been revolutionary for reinforcement learning (RL) — it can now solve complex decision and control problems.
Successful RL methods were handcrafted using mathematical derivations, intuition, and experimentation.
This approach has a major shortcoming— it results in specific solutions to the
RL problem, rather than a protocol for discovering efficient and robust methods .
In contrast, the emerging field of meta-learning provides a toolkit for automatic machine learning method optimisation, potentially addressing this flaw.
However, black-box approaches which attempt to discover RL algorithms with minimal prior structure have thus far not been successful.
Mirror Learning , which includes RL algorithms, such as PPO, offers a potential framework.
In this paper we explore the Mirror Learning space by meta-learning a “drift” function.
We refer to the 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 confirm 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 (RL) algorithms to successfully tackle large-scale problems.
Early deep RL algorithms, such as A2C [22] and DDPG [19] follow the scheme of generalized policy iteration (GPI), where the RL agent alternates between policy evaluation and policy improvement phases [31].
Namely, they train a critic neural network that estimates the value function, and an actor network which models the agent’s policy, and use these to estimate the update direction with policy gradients of the expected return.
Another approach is that of trust-region learning (TRL) methods, which construct a surrogate objective that they optimise within a small region around the current policy [28], aiming at guaranteeing stability of the policy iterates.
Indeed, TRL-inspired algorithms, such as TRPO [28] and PPO [29], have been among the most widely used methods [2] and are known for their performance and stability.
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 flaws, such as the brittleness to hyperparameter settings [28, 11], and a lack of optimality guarantees.
Unfortunately, the most promising alternative approach, algorithm discovery , thus far has been a tough nut to crack.
First of all, techniques that enable optimisation of algorithm components in the outer loop process—meta RL [27, 7]—are at an early stage of development, and are vastly more compute intensive than the RL problem in the inner loop [35].
Second, to discover more complex algorithms, one needs to solve a more sophisticated, and computationally taxing meta-optimisation problem.
As a result, most methods discover simple algorithms at best [23].
Recently, Mirror Learning [16], a new theoretical framework, introduced an infinite space of provably correct algorithms, all of which share the same template.
In a nutshell, a Mirror Learning algorithm is defined 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 [16].
Popular RL methods such as TRPO and PPO are instances of Mirror Learning.
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.
Specifically we parameterise a drift function with a neural network, which we then meta-train using evolution strategies [26, ES].
The outcome of this meta-training is a specific mirrorlearning algorithm which we name Learnt Policy Optimisation (LPO).
However, we take a step further and learn from LPO about important RL concepts that it has discovered by visualising and analysing its optimisation objective.
While having a neural network representation of a novel, high-performing drift function is a great first step, our next goal is to understand the relevant algorithmic features of this drift function.
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.
To our knowledge, DPO is the first theoretically-sound, deep RL algorithm that was discovered via meta-learning.
Related Work
In the last few years, significant effort has been directed at developing effective RL algorithms through both algorithmic and implementational advances.
Fujimoto et al.
[9] combine DDPG policy training with estimates of pessimistic Bellman targets from a separate critic.
Hsu et al.
[14] stabilise the, previously unsuccessful [29], KL-penalised version of PPO and improve its robustness through novel policy design choices.
Haarnoja et al.
[12] 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.
Indeed, recent work explores applications of meta-RL techniques in guiding RL algorithm discovery and design.
RL 2 equips a learning agent with a recurrent neural network conditioning on transitions between tasks, and adapts the agent’s behaviour to the current environment [4].
Similarly, a MAML agent learns policy meta-parameters which can adapt to any task with a few steps of gradient descent [7].
Neither RL 2 nor MAML, however, go beyond improving the robustness of ultimately handcrafted algorithms.
To overcome this, Xu et al.
[35] introduce an actor-critic method that adjusts its hyperparameters online using meta-gradients that are updated with every few inner iterations.
Similarly, STAC [37] uses implementation techniques from IMPALA [5] and auxiliary loss-guided meta-parameter tuning to further improve on FRODO.
Such advances have inspired extending meta-gradient RL techniques to more ambitious objectives, including the discovery of algorithms ab initio .
Notably, Oh et al.
[23] succeeded in meta-learning an
RL algorithm, LPG, that can solve simple tasks efficiently without explicitly relying on concepts such as value functions and policy gradients.
Similarly, Evolved Policy Gradients [13, EPG] meta-trains a policy loss network function with Evolution Strategies [26, ES].
Although EPG surpasses PPO in average performance, it suffers from much larger variance [13].
MetaGenRL [15] instead meta-learns the loss function for deterministic policies which are inherently less affected by estimators’ variance [30].
MetaGenRL, however, fails to improve upon DDPG [19] 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 [16], 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 efficient 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, respectively.
3.1 Reinforcement Learning
Formulation We formulate the reinforcement learning (RL) problem as a Markov decision process (MDP) [31] represented by a tuple ⟨S , A , R, P, γ, d ⟩ which defines 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]
the state value
[Equation]
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 ) [31].
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 [22] 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 [31, GPI] suffer from large variance and instability [38, 30, 29].
Hence, methods that constrain (either explicitly or implicitly) the policy update size are preferred [28].
Among the most popular, as well as successful ones, is Proximal Policy Optimization [29, PPO], inspired by trust region learning [28], which updates its policy by maximising the PPO-clip objective,
[Equation]
where the clip ( · , 1 ± ϵ ) operator clips (if necessary) the input so that it stays within [1 − ϵ, 1 + ϵ ] interval.
In deep RL, the maximisation oracle in Equation (4) is implemented through 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 optimise the algorithm of the learning agent.
Formally, suppose that an RL algorithm alg ϕ , parameterised by ϕ , trains an agent for K iterations.
Meta-RL aims to find the meta-parameter ϕ
= ϕ ∗ such that the expected return of the output policy, E [ η ( π K ) | alg ϕ ] , is highest.
3.2 Mirror Learning
A mirror-learning agent [16], in addition to value functions, has access to the following operators: the drift function D π k ( π | s ) which, intuitively, evaluates the significance 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 defined, a mirror-learning algorithm updates an agent’s policy by maximising the mirror objective
[Equation]
If, for all policies π and π k , the drift function satisfies the following conditions:
It is non-negative everywhere and zero at identity D π k
( π | s ) ≥ D π k ( π k | s ) =
[Equation]
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 → ∞ [16].
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.
It can be shown that PPO is a valid instance of Mirror Learning, with the drift function:
[Equation]
Although in its formulation, PPO puts no explicit constraints on its update size [29], as its maximisation oracle (see Equation (5)) is 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 display different properties, such as sample complexity, world-clock time efficiency, or they can aid fulfilling other purposes [16].
That is to say, depending on a preference, a different algorithm is optimal, and PPO is not necessarily the one.
In this paper, we specify the optimality criterion to be the return of the K th iterate, η ( π K ) , and discover a mirror-learning algorithm that is superior to PPO with respect to it.
3.3 Evolution Strategies
Evolution Strategies [25, 26, ES] is a backpropagation-free approach to optimisation of stochastic functions.
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 ϵ 1 , . . .
, ϵ n ∼ N (0 , I ) , the vector 1 nσ (cid:80) ni
=1 F ( ϕ + σϵ i )
ϵ
i is an unbiased gradient estimate.
To reduce variance of this estimator, antithetic sampling is commonly used [24].
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 [33, 34, 23, 6, 21].
[Figure]
Methods
Our overall approach is to meta-learn a drift function to perform policy optimisation over a fixed episode length K .
Hence, our meta-objective is the expected final return,
[Equation]
In Subsection 4.1 we describe the parameterisation of our learnt drift function, while in Subsection 4.
we provide a detailed description of our meta-learning approach.
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 ) as 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, ultimately forming the following input
[Equation]
We perform ablations on different variants of the drift input in Appendix ?? . In general, we found that the performance of the learned drift function seems to be fairly robust to different input parameterisations.
In order to guarantee that the neural network is a valid drift function, it suffices to impose that f ϕ ( x r,A )
=
and ∇ r f ϕ ( x r,A ) = 0 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 difficulty 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 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 or statistics measuring the progress of training.
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 specific applications, it may be beneficial to consider other, possibly domain-specific, information, and we do encourage it. [Figure] [Figure]
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 mirrorlearning 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 final 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 efficiency 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 [17] and the Learned_optimization [20] 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 specific 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 figures 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 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 sufficient 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 , we observe that the ratio derivative of the LPO objective is positive in a large region, roughly corresponding to r < 1 − ϵ .
This implies that actions that 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 [32].
Hence, LPO learns to decrease r down to 1 − ϵ , but unlike PPO, encourages r to stay precisely [Figure] 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.
Cautious optimism for positive advantage.
The upper-right quadrant corresponds to A > 0 and r > 1 , 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 want 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 + ϵ , 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.
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, advancing 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 identified 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 [Figure] 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.
Specifically, we have found 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 verification 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, which we enhanced with advantage normalisation as we observed it to improve 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 [18, 36] 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 reveal 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.
Specifically, 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 overfit to specific 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 efficiency.
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/P
A Meta-Training Details A.1 LPO-Zero
Our LPO-zero implementation was implemented on top of the learned_optimization library [20].
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 fitness function.
Meta-training was done over the course of 2 days and performed approximately 400 outer-updates.
We find though performance still increases with increased meta-training time.
A.2 LPO
Our LPO implementation was implemented on top of the evosax library [17].
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 find though performance still increases with increased meta-training time.
[Table]
Unroll Length
Number
[Table]
Number
[Table]
[Table]
Number
[Table]
Environments
[Table]
[Figure]
D DPO Drift Verification
The DPO drift function is given by
[Equation]
The first 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 final 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]
r
[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 = 1 .
Taken together
[Figure]