Learning to Optimize for Reinforcement Learning

Thank you for your detailed feedback, suggestions, and comments. We address your points in the following.

• Add more ablation studies. Please check the general response.
• In Figure 5, you suggest the overlap in the support of the gradients is a predictor of performance. A simple experiment to evaluate this would be retraining the optimizer with rescaled rewards on the grid-world tasks, which would shift the support of the meta-training gradients. Yes, this is exactly what we did. We designed six gridworlds with different reward scales such that the union of all agent-gradients covers a large range. With carefully rescaled rewards, we can significantly improve the performance of the learned optimizer trained in gridworlds for Brax tasks. In fact, simply rescaling the agent-gradients before feeding them into the optimizer also works.
• The inability of LinearOptim to learn anything is very surprising and should be investigated further to ensure it is not erroneous. We checked the implementation again and ensured it was not erroneous.
• Transformation of the gradient output is a major component of gradient processing. An expanded form of the output should be presented. We will add more explanations for this part.
• Add more citations, modify misleading claims, and correct typos. We will add those citations and fix the typos and misleading claims.

Thank you for your response and new experiments.

New ablation studies - Thank you for including these ablations. Regarding pipeline training, the results are somewhat promising but I would like to see them strengthed with the full environment set and labeled significance. Regarding optimizer structure, it is not clear from the experimental setup that Figure 4 shows an ablation of this specifically, since Optim4RL is proposed as a combination of methods. This should be clearly described in the experimental setup or Figure 4. Regarding gradient preprocessing, it is quite significant for the story of this paper that this method does not aid (and possibly harms) performance. I commend the authors for being candid about these results and the paper still contains useful contributions without them, however incorporating this and the other new results would require a significant revision, which is likely out of scope for this rebuttal period.

Gradient distribution experiments - Thank you for pointing this out. Even though I agree with the motivation for this experiment, it is not clearly labeled in the paper (see fig 5, 6 captions and section 6.2 title) and incomplete without replotting the agent-gradient distribution in figure 6. Furthermore, adding training environments (rather than just rescaling the reward magnitude on the original environment) adds a confounder beyond the gradient distribution. I would strongly recommend merging figures 5 and 6, such that figure 5 a-d contains one more curve for the Optim4RL on extended training set and 5 e-h contains the distribution over the extended training set.

LinearOptim results - I was suggesting that the results be extended with further environments, since they are highly surprising and necessitate further investigation. The release of anonymized source code would also be reassuring.

Remaining edits - Thank you for being receptive to these suggestions - I would like to see them included in the revision.

Dear Reviewer,

We did not find the mentioned related work GROOVE (Jackson et al., 2023). Could you please provide more details about this work?

Thanks.

Thank you for your detailed feedback, suggestions, and comments. We address your points in the following.

• Add more ablation studies. Please check the general response.
• Add more citations and training details. We will add those citations and more training details to clarify.
• Have the authors analyzed how different values of m and n impact this? Could reset times be adapted by leveraging optimizer reset properties? Do we require them to be equal all the time? m and n are not required to be equal all the time. In our experiments, n is the number of training environments that depends on the task itself; m is affected by the training steps of each task and it has a similar magnitude as n. We leave the study of the performance impact of m and n as future work.

Thank you for your detailed feedback, suggestions, and comments. We address your points in the following.

• Add more ablation studies. Please check the general response.
• Why can’t the learned optimizer "overfit" to the training task to outperform Adam? Note that we choose a very small network to represent the learned optimizer so that we can train it within the computation and memory budget. This also helps reduce the deployment computation burden, i.e. reducing the inference computation of a fixed learned optimizer. However, the disadvantage of using a small network is that it is hard for the learned optimizer to “overfit“ a relatively complex task, such as the gridworlds which have image inputs. Instead, we did experiments on a relatively simple task — Catch form bsuite. The results indeed showed that the learned optimizer can “overfit” Catch, converging faster than Adam and RMSProp.
• Why do we need to check that the model has enough capacity to represent the identity function (e.g. a single linear layer can represent the identity function easily)? Why do we need an RNN for this problem? Isn’t the input just the current gradient? We want to verify whether a learned optimizer represented by an RNN model can approximate the simplest gradient optimization method — SGD with learning rate 1. In this setting, essentially the RNN model is required to approximate an identity function with agent-gradients as inputs. We use RNN models so that historical gradients can be taken into account and are not completely ignored. For example, historical gradients are important to compute the momentum in RMSProp. Theoretically, a neural network can represent an identity function. However, in practice, it might not be true. Please check this paper for an example.
• In Figure 4 (b) it seems like STAR is starting to learn a bit. Is it possible that with a bit more hyper-parameter tuning it would match the other optimizers in performance? We tuned STAR with more hyper-parameter options and improved its performance (average return) to 10; however, this is still much worse than the performance of our method (average return>20).

Thank you for your detailed feedback, suggestions, and comments. We address your points in the following.

• Add more ablation studies. Please check the general response.
• The authors show that the agent-gradient distribution is non-IID in RL. However, they do not show that the gradient distribution for normal supervised learning (SL) is IID. We did not claim that non-iid gradient distribution is a unique challenge for learning to optimize for RL. In fact, the gradient distribution in supervised learning is also non-iid. For example, there is a larger portion of zero gradients at the end of the training for both RL and SL. However, we expect the gradient distribution to be more iid in RL tasks since RL tasks are inherently more non-stationary. We will rewrite relevant parts to clarify this.
• The difference between pipeline training and the training style of LPG. The major difference between pipeline training and the training style of LPG is that pipeline training is designed to make the agent-gradient distribution more iid. Pipeline training resets training units in a specific way such that the input agent-gradients are more diverse, spreading across a whole training interval.
• Why is the neural network so small? We choose a small network so that we can train an optimizer within the computation and memory budget. This also helps reduce the deployment computation burden, i.e. reducing the inference computation of a fixed learned optimizer.
• Many of the learned optimizers use ES to train their learned optimizers. Is there any particular reason you decided not to do this? Learned optimizers are mainly trained with the meta-gradient method and evolutionary method. Following previous works, we choose one of the methods to train our learned optimizer, i.e. the meta-gradient method. Using the evolutionary method is an interesting future direction.
• Add more citations and correct typos. We will add those citations as well as fixing the typos.
• Are you retraining STAR and L2LGD2? Yes, we retrain them from scratch.

Thank you for the additional results! Using the WD to the total distribution is a neat way to calculate the iid-ness. It would be good to see these plots and histograms. Are you able to update the paper with them?

Minor nits: It would be good to make sure the problems are similarly-scaled. E.g. would it be possible to do something like doing behavioral cloning or contextual bandits on transitions sampled from an expert. This way, you assure that the input spaces match and the loss functions / objectives are effectively the same, but the problem is now iid.

the non-iid issue in learning to optimize is not explicitly identified in previous works

This doesn't make it non-trivial. My understanding is that it is so commonly understood that it's not something most people would write in a paper. Adding the comparisons between RL and SL make it non-trivial. It's good to see the authors do this.

I hope that the authors will update the paper so that we can see the latest results (that should hopefully replace Sections 4.1 and 4.2) in more detail as soon as possible. My understanding is that the authors can re-write and update the manuscript during the discussion period.

We have finished experiments of meta-training STAR in Ant for Figure 4(c). The performance of STAR in Ant is quite unstable, achieving an average return of 1000 while our method Optim4RL can reach an average return of 7000.

We will also update the draft to incorporate the required modifications by all reviewers with the new results.

Dear Reviewers,

We have updated our draft to incorporate modifications by all reviewers with the new results. Specifically, we

• add section 3.2 to discuss related works about discovering general reinforcement learning algorithms;
• add ablation studies of pipeline training and the inductive bias of our optimizer Optim4RL in section 6.1 (We didn't include gradient processing ablation since this component may be moved to the appendix.);
• add STAR and LinearOptim results in Figure 4(c). Note that LinearOptim in Ant fails due to NaN errors, probably due to highly non-stationary parameter updates;
• add a discussion about m and n in section 5.1;
• add more details about computing $\Delta \theta$ in section 5.3;
• plot the agent-gradient distribution of all six gridworlds in Figure 6 as well as adding the related discussion;
• update section 4.1 and 4.2 to discuss supervised learning setting and include related experiments in appendix E;
• add more training details in appendix C;
• include other minor changes, such as fixing typos and misleading claims.

We thank all reviewers for providing your valuable suggestions. We are grateful for the opportunity to address your concerns and would like to express our sincere appreciation for considering a reevaluation of our work. In case we forget to include any important changes, please let us know.