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, 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 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 confirm state-of-the-art performance of LPO and DPO, as well as their transfer to unseen settings.

Reviews and author responses:

 Thanks for the thoughtful feedback!

> On Input Parameterization

That makes sense! However, on our end, these experiments are quite costly (the results in the Appendix are not fully meta-trained to convergence, for example). We decided to go with the input we were most confident would work rather than the most simple one. Removing squared inputs actually happened to be the one ablation that had one of the (relatively) bigger differences in performance (interestingly, it performed better in Walker2D but worse in the other environments). We will consider just including the simplest / minimal drift input for the camera-ready copy though! 

> Extra Environments

We will for sure incorporate this discussion in the camera-ready-copy! Currently, it is a little bit too short notice for us to properly update the submitted manuscript, but we fully agree with the reviewer's points and appreciate the feedback. We will also try to do more analysis to confirm the trend of "multiple environments" vs. "long rollouts" performance.



 Hi authors,

Thank you for performing the extra experiments and including the code! I think all of those are very helpful.

Some comments in light of that:
- Input parameterization: Since the input parameterizations perform similarly, wouldn't it make sense for the main paper to give the simplest representation for $x_{r,A}$ that still performs well? You can then include the higher-order representations (e.g., squared inputs) in the appendix and show they only give marginal improvements at best.
- Extra environments: It's good to see that LPO generalizes to Atari-like domains! I think these plots (or some subset) should be included in the main paper to show that your objective **generalizes**. In light of that, it appears that your objective works well in cases where you have many environments for rollouts and perform short rollouts. If you have a small number of environments and perform long rollouts, the learned objective does not do well. I think it's important to point out that trend (if that is indeed the trend) in the main paper and corroborate the negative results in the Appendix. This will make the community aware that meta-learned objectives can depend on RL hyperparameters like number of environments for rollouts, which can be a good avenue for future research.

 Thanks for the response!

>  I still think that the algorithm design is heavily engineered

We perform ablations on the input parameterisations in the updated Appendix E! In general, it seems as though using different input features produces similar results. 

> evaluation on Mujoco-Brax-based setups only makes it difficult to generalize to all environments where PPO has been well tested

We have included results in Minatar [1] in the updated Appendix F! Interestingly, LPO more consistently outperforms DPO in that setting, implying that there may be more relevant features in LPO that we missed. Both LPO and DPO are comparable to PPO. We chose Minatar because there exist vectorisable implementations that interface nicely with Brax's PPO implementation [2].

[1] Young, Kenny, and Tian Tian. "Minatar: An atari-inspired testbed for thorough and reproducible reinforcement learning experiments." arXiv preprint arXiv:1903.03176 (2019).

[2] Lange, Robert. "{gymnax}: A {JAX}-based Reinforcement Learning Environment Library," http://github.com/RobertTLange/gymnax


 > You should run and include experiments confirming your intuition about other inputs not making much difference.

We have done ablations on the inputs to the drift function in the updated supplementary! The results are in Appendix E. In general, we find that different input parameterisations tend to perform similarly. 

> To verify robustness and improvement of LPO/DPO with respect to PPO, you should run experiments in other domains where PPO is known to work

We have included results in Minatar [1] in the updated Appendix F! Interestingly, LPO more consistently outperforms DPO in that setting, implying that there may be more relevant features in LPO that we missed. Both LPO and DPO are comparable to PPO. We chose Minatar because there exist vectorisable implementations that can interface nicely with Brax's PPO implementation [2].

> I did not find any code in the supplementary.

We have included code in the supplementary!

[1] Young, Kenny, and Tian Tian. "Minatar: An atari-inspired testbed for thorough and reproducible reinforcement learning experiments." arXiv preprint arXiv:1903.03176 (2019).

[2] Lange, Robert. "{gymnax}: A {JAX}-based Reinforcement Learning Environment Library," http://github.com/RobertTLange/gymnax

 >> **Followup for A1** You should run and include experiments confirming your intuition about other inputs not making much difference.

>> **Our Answer** As we elluded in A1, currently we do not know other inputs that would guarantee us invariance across environments other than the ratio $r$ and advantage $A_{\pi}$. As for different transformations of $A_{\pi}$ and $r$, we already tried polynomials in $r$ of higher order, but it did not bring improvements. This suggests that other smooth functions would not bring much improvement either, as such functions can be approximated with polynomials by Taylor's theorem. We are running some experiments to confirm that.

>>**Followup for A4** You can move the results on LPO-Zero to the appendix.

>>**Our Answer** Another reason for which we would like to keep it is that its objective is similar to that of LPO (see Figures 2 & 3). We think it is an important finding that shows that our meta-learning method indeed leads to a discovery of RL concepts.

>>**More Experiments** To verify robustness and improvement of LPO/DPO with respect to PPO, you should run experiments in other domains where PPO is known to work. Also, what do you mean by that you struggle to make Brax's PPO work?

>>**Our Answer** Following the reviewer's advice, we implemented DPO in PyTorch and tested it in MuJoCo. Unfortunately, it did not work that well (at least with PPO hyperparameters). Although it is possible that more extensive hyperparameter tuning would improve the performance, it is likely that in the meta-learning process LPO has learnt an algorithm (and discovered DPO) that is optimal for the specific implementation and testing frameworks (Brax and MuJoCo). Similarly, PPO for Brax works poorly in MuJoCo. We chose JAX and Brax because they enabled us to train multiple inner-loop agents fast and in parallel [Freeman et al., 2021]. Once other development tools, like PyTorch and MuJoCo, allow for it, our method can be easily applied in them. 

[Freeman et al., 2021] C. Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, Olivier Bachem; "Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation". 2021.


 Hi authors, thank you for answering and clarifying all the issues and questions.

I am convinced that the direction of discovering better algorithms as proposed in the paper holds merit, and the replies have answered most of my questions. I also acknowledge the critical insights drawn from the analysis of PPO and LPO.

However, I still think that the algorithm design is heavily engineered, and limited evaluation on Mujoco-Brax-based setups only makes it difficult to generalize to all environments where PPO has been well tested. Further work on introducing a more generalizable discovery of algorithms rather than using evolutionary strategies for maximizing performance will be a more substantial contribution.

I have improved my score following the author's justifications and changes in the Appendix references.

 Hi authors, thank you for your response!

**A1**: Your explanation seems reasonable, but I think it's important to include those ablations to confirm your intuition.

**A4**: You can always refer the reader to the Appendix if LPO-Zero is moved there as well as briefly mention in the main paper that it did not work very well. Keeping it in the main paper is admittedly not a big deal.

**More evaluations**: Evaluating your learned/crafted objectives on a wider range of environments is crucial. One thing that makes PPO so appealing is that its clipped objective works out of the box on a large range and diversity of environments (including Brax). If your LPO/DPO objective performs poorly outside of Brax, that limits the impact of your work.

When you say you are "struggling to get Brax's PPO implementation to work in [the other environments]," do you mean that porting the code of the objective function to other environments is difficult or that the resulting RL experiments have poor results?

 We thank the reviewer for their feedback!

>There are no additional results in Appendix D as mentioned in the paper

Apologies! This was from a previous version of the paper in which we did not have space to present all of the Brax results. We have since fixed the paper.

>there is no justification for why an initial drift function of PPO helps.

Indeed we do not have a formal theoretically justification for why an initial drift function of PPO helps. However, intuitively, it makes sense that initialising close to a well-known performing algorithm would perform better than initialising from a completely random one.

>overall Appendix E is missing

Apologies for the mislabeling of sections! “Appendix E” in the paper referred to “Appendix C” in the supplementary, where we plot the results of training LPO in different Brax environments. We have since fixed the references

>**Q1** It would be interesting to visualise and analyse what does the drift function $\tilde{f}_{\phi}$ converge to, and contrast it with PPO.

 **A1** We believe that this question is largely addressed in the paper already. Although we did not explicitly visualise $\tilde{f}_{\phi}$, we visualised the optimisation objective of LPO, which involves both advantage and the converged drift function (see Figure 3). Specifically, we plot the derivative with respect to the policy ratio, which explicitly informs how the objective affects learning. We believe that such a plot is more informative than a naive plot of the objective, as that would potentially involve "noisy" components that are independent of the policy variable. The same visualisation was conducted for PPO (Figure 1), whose objective was contrasted with LPO in Section 6.

>**Q2** How would a fine-tuned PPO perform in comparison to a fine-tuned DPO?

**A2** We have not yet fine-tuned DPO. However, as we used Brax's fine-tuned PPO, and simply adopted the same hyperparameters to DPO and obtained better performance, we expect that fine-tuning DPO would perform even better. We are currently in the process of collecting these results, but we don’t expect the results to fundamentally change the core takeaways from the paper.

>**Q3** Can the authors add the corresponding PPO heat map to Figure 5?

**A3** That heat map has already been plotted in Figure 1!

>**Q4** The inputs to the drift function of LPO are cherry-picked.

**A4** We chose policy ratio and the (normalised) advantage function as inputs because they are invariant to many environments. Indeed, we trained such a drift on Ant only, and transferred it to other environments. The nonlinear transformations in line 180, which were chosen arbitrarily, were supposed to help the drift function learn complicated mappings. Designing other, transferable inputs to the drift function is an interesting avenue of future research.

>**Q5** Why DPO does not use $(r-1)^n$ as input?

**A5** DPO is not a meta-trainable algorithm (unlike LPO). It is, instead, an analytical approximation of $LPO$. The functions that it uses (line 302) were chosen to reproduce the intuitions from the LPO opitimisation objective. Mappings like $(r-1)^n$ were not necessary. 

>**Q6** Can you provide some studies on how different values of $\alpha$ and $\beta$ affect the performance of DPO?

**A6** Although we write the variables $\alpha$ and $\beta$ symbolically for clarity, it is important that these take the values $\alpha=2$ and  $\beta=0.6$ to reproduce LPO accurately. We do not consider them hyperparameters to tune (at least, not in this paper), but rather elements of our analytical model. 

We will try to provide these results for the reviewer shortly; however, we do not expect them to significantly change the paper.

>**Q7** There were too few experiments conduceted to justify the generalisation claims about LPO. Can you run some experiments on Procgen? 

**A7** We are looking into evaluating versions of LPO and DPO in other environments. However, we are struggling to get Brax’s PPO implementation (upon which LPO and DPO were learnt) to work well in them. Brax’s implementation makes a number of different code-level design decisions in order to run quickly in terms of wall-clock time (rather than sample-efficiency). For example, they use a very large number of parallel environments, very short environment trajectories and update their GAE calculation after every minibatch. 

However, note that we provided experiments on 8 tasks in total, with a remarkably high unseen:seen ratio ($7:1$). Other works usually use less---for example [Kirsch et al., 2020] use $4$ tasks in total, with the unseen:seen ratio of only up to $2:1$. 

>**Q8** Could you add a properly referenced appendix and include equation numbering?

**A8** Lack of these is a clerical mistake on our behalf. We apologise and have addressed this concern in revision.


 We thank the reviewer for their thoughtful review! We address some of their feedback below

>**Q1** If the inputs to the drift function from line 180 are changed, does the resulting algorithm still work well?

**A1** Our choice of the input arguments was deliberate. The choice of invariant quantities---the policy ratio $r$ and the (normalised) advantage function $A_{\pi}$---enabled us to learn a drift function that is applicable to multiple environments. The nonlinear mappings of these (line 180) helped us learn more complex features. We considered increasing the maximal degree of the polynomial in $r$ (from quadratic to cubic) but it did not bring any improvement. Lastly, finding alternative, invariant inputs to the drift function is an interesting alley of future research. 

We will try to more ablation results on the inputs in an updated version shortly (these experiments are computationally expensive); however, we do not believe that the results will change the fundamental takeaways from the paper.

>**Q2** Is the presence of $\bar{\pi}$ in line 137 a typo?

**A2** Yes, it is. Thank you for pointing it out. We have fixed it in the revision.

>**Q3** The variable $r$ is used for both the reward and the policy ratio. That it confusing.

**A3** Thanks for pointing that out! We have changed the notation in the revision.

>**Q4** I suggest moving the training of LPO-Zero to the appendix. It is not an important part of the paper.

**A4** We agree that LPO-Zero is not core to the paper and does not achieve great results (and frankly including the results invites reasonable criticisms/concerns that Reviewer KEqn points out); however, we still think it’s important to show that it was something we attempted so that future work may be able to overcome the method’s shortcomings.

>**Q5**Despite what the paper says, I did not find any code in the supplementary.

Apologies! We will shortly include code in the supplementary (we realised late that there is still more anonymisation to be done).

>the evaluation environments are limited

We are looking into evaluating versions of LPO and DPO in other environments. However, we are struggling to get Brax’s PPO implementation (upon which LPO and DPO were built) to work well in them. Brax’s implementation makes a number of different code-level design decisions in order to run quickly in terms of wall-clock time (rather than sample-efficiency). For example, they use a very large number of parallel environments and very short environment trajectories and update their GAE calculation after every minibatch. 



 We thank the reviewer for their thoughtful review! We address some of the reviewer's questions below.

>It's straightforward to compare the learned methods with these mirror-learning-based methods, e.g., TRPO

Note that PPO is still considered to be one of the most widely-used methods in RL. In particular, when it was introduced, it was notable for outperforming TRPO on many tasks. Thus, while we could compare to TRPO, for example, we would simply expect it to perform worse than PPO. 

>The experimental results of learning drift functions from scratch underperform LPO and even PPO.

Indeed! We think this is partially a result of our limited compute budget. Training for longer did improve performance, but the experiments already took a very long time to run. We hope that in future work we can scale the method to better learn from “scratch”. 

 We see starting with the PPO initialisation simply as a way to speed up the training. While the initialisation did involve some human handcraft, the end result was still something significantly different from PPO. Note that any design decisions when it comes to the parameterisation of the drift function (e.g. even the usage of specific neural networks architectures, activation functions, and hyperparameters) can also be counted as “human handcraft”, which is inevitable.


> **Q1** Is it possible that the optimal learning mechanism (i.e., the optimal drift function) is different for every environment?
 
 **A1** It is definitely possible, and is likely true. For example, note that new RL algorithms rarely completely outperform previous ones across every environment. However, the idea is that by restricting the expressivity of the meta-parameterisation (e.g. through the input arguments), we can prevent overfitting. We used the Mirror Learning framework to generate a theoretically-justified meta-parameterisation that would allow for generalisation through its convergence guarantees. 

One interesting way to potentially explore this idea in future work would be to allow environment information to be an input to the drift function, thus changing the “learning mechanism” for different environments. 


>**Q2** Can you explain why you use the candidate arguments in the equation in Line 180? Are there any other candidates?

Firstly, we wanted to use inputs which are invariant for all possible environments and policy designs (discrete/continuous). A drift function with such inputs, after being meta-learnt on one task, can be deployed in other ones. The policy ratio can always be obtained when stochastic policies are considered, and the normalised advantage function commonly appears in deep-RL implementations. Information like *state* can vary from one environment to another, thus we did not include it as an input.

Secondly, the input vector is constructed so that it is zero when the old and candidate policies are the same, which helps us meet the drift condition 1. 

Thirdly, we used different non-linear transformations of the ratio and advantage, to help learning more complex mappings.

There are other candidate inputs, such as using higher-order polynomial inputs; however, we found that including them do not seem to significantly improve the performance. We will try to include the results in an updated version shortly (these experiments are computationally expensive); however, we do not believe that the results will change the fundamental takeaways from the paper.

 
This paper proposes a method to learn how to optimize policy under the framework of Mirror Learning. Based on the learning result, they formulate a closed-form policy optimization method for learning the policy. The idea of this paper is interesting. They propose a method that lies between human handcraft and learning discovery. The experimental result is interesting.
I'm not satisfied with the claim and some experimental results (see below for detail).

The author claim that TRPO and PPO are handcrafted by humans rather than discovered by learning. This is overclaimed to me. 
- The key idea of this paper also follows Mirror Learning, which is also discovered by humans.
- Furthermore, learning drift functions from scratch underperforms with respect to LPO. Instead, they learn with the PPO initialization, which is more handcrafted.

I'm not satisfied with the experimental result.
- As the authors claim that discovered by an algorithm is better, and there exits several variants following Mirror Learning, e.g., TRPO and variants of PPO. It's straightforward to compare the learned methods with these mirror-learning-based methods.
- The experimental results of learning drift functions from scratch underperform LPO and even PPO. Although the authors have also tried learning from PPO initialization, this implies it involves more human handcraft. 
 Learning discovery is interesting, and using the fixed learning mechanism is natural. Could that be true if the optimal learning mechanism is different across different environments?

Can you explain why you use the candidate arguments in the equation in Line 180? Are there any other candidates?  

 This paper studies the learning of a “drift function” (a proximal regularization term in policy optimization) for reinforcement learning (RL). By parameterizing the drift function as a neural network and framing the objective as the expected return after K iterations of RL, the authors can find this drift function through black-box optimization. The authors analyze the structure of the learned drift function and compare it to PPO’s drift function. Inspired by insights on the learned function, the authors also present an analytic drift function that contains many of the same properties as the learned function. Experiments show that policy optimization with both novel regularizers regularly outperform PPO on several Brax continuous control problems. The paper studies a highly practical problem in the design of surrogate objectives in policy optimization. The proposed approach is sensible, and the results do appear promising. The writing and presentation are also very clear, and it’s easy to read the paper. I especially like the analysis of the LPO function in Section 6, and the insights to its structure (rollback and cautious optimism) make sense for its role in policy optimization.

I think one major weakness of their work is that the evaluation environments are limited. All environments are continuous control problems in the Brax simulator. To assess how well the LPO/DPO drift function generalizes, I highly suggest evaluating the functions in other physics simulators (e.g., MuJoCo and Isaac) and entirely different domains (e.g., Atari and other discrete problems).
 What happens if you try to learn a drift function without the input $x$ from line 179? Does this drift function also give good RL performance?

Some other comments:\
Line 137: Is it supposed to be $\pi_k$ instead of $\bar\pi$?\
I would suggest moving Sec. 5.1 (Learning drift functions from scratch) to the Appendix since it performs worse than LPO and is not a significant part of the paper.\
Line 177: The variable $r$ is already used for the reward. It may make sense to use a different variable.\
Despite what the paper says, I did not find any code in the supplementary.
 The authors acknowledge that they only study the drift function and hope to expand optimization to other parts of the mirror learning objective.

 The paper is based on the concept of Mirror Learning which includes several RL algorithms like PPO, TRPO, MDPO, etc. The work empirically exploits the flexibility of choosing the "drift" function in the Mirror Learning domain by using meta-learning. Given a neural network drift function, the authors use Evolutionary Strategies to meta-learn a single drift function (generalizable across environments) where the objective is to maximize the return of the policy trained using an update rule composed of the drift function. This is termed the Learnt Policy Optimization (LPO) which is proposed in two variants: one uses PPO clip augmentation at the last hidden layer output and the other doesn't. Eventually, the authors draw meaningful observations from LPO's analysis and proposed Discovered Policy Optimization (DPO) using those observations. The authors verify the approach by showing that DPO's behavior is close to LPO's behavior. All experiments are conducted in brax-based environments. The proposed method holds some strength by improving over a widely-used baseline using:
- parametrized drift functions to capture the importance of action distribution ratio and advantage in policy improvement.
- empirically observed and properly justified learned behavior based on the advantage being positive or negative.
- well-curated and mechanically optimized/fine-tuned algorithm (DPO) using the discovered drift function model.

However, there are the following weaknesses in the proposed framework:
- there are very less ablations on the choice of structure, for example, the non-linear formulation of $x_{r,A}$ and choice of constants for DPO (overall Appendix E is missing)
- there is no justification for why an initial drift function of PPO helps. While it is acknowledged that LPO-Zero underperforms ("possibly" due to less training), there is no evidence that much more training can help LPO-Zero converge to LPO.
- the experimental validation is not sufficient to prove the improvement claimed by the contribution. (there are no additional results in Appendix D as mentioned in the paper)
 It would be nice if the following can be justified and included in the paper:
- Since the overall objective for policy improvement aims to minimize $\mathbb{E}_{\pi_k}\[\mathcal{D}(\pi|s)\]$, I am curious to know where $\tilde{f}_\phi$ converges to and how the final drift function $f_\phi$ is different from the PPO one. This also points towards another interesting phenomenon of DPO (with constants) being an approximated version of PPO with smoothened contours (instead of strict clipping).
- If there is hyperparameter optimization of PPO (just clip factor) with evolutionary strategies to get the best performance (return) for a particular environment, what does the DPO performance comparison look like? 
- Can the authors add the corresponding heat map of PPO as well in Figure 5?
- The discovered drift function seems to be cherry-picked. It will be beneficial if the authors add justifications for their choice of the input vector to the drift function during LPO optimization and also in DPO. Why not $(r-1)^n$ in DPO as that would also satisfy the feasibility conditions? 
- For a single environment, can the authors show the variation of the final performance as a function of constants $\alpha$ and $\beta$ ?
- There should be more experiments on multiple benchmarks to justify the claim since this is an improvement of an algorithm. Can you add some analysis of Procgen experiments?
- Can you add a properly referenced appendix? The references in the paper do not correspond properly to some Appendix sections making it difficult to analyze.
- Also, add equation numbering. The authors have acknowledged that the method shown in the presented work is a step towards discovering more robust and high-performing RL algorithms. In general, the work does not has significant limitations but lacks proper justification of choices which makes it unclear how to conduct algorithm discovery in a proper manner. Further, more empirical evidence and experiments are required to complete the contribution.

Meta-review: This paper proposes a policy-search approach in which a "drift" function is meta-learned in the framework of Mirror Learning.
After reading each other's reviews and the authors' feedback, the reviewers have solved most of their concerns and agree that this paper deserves publication.
The authors need to consider the reviewers' suggestions in preparing the final version of their paper.

Decision: Accept