Learning to Optimize for Reinforcement Learning

TL;DR

The paper tackles the challenge of learning optimizers for reinforcement learning, showing that optimizers trained for supervised learning fail in RL due to non-IID gradient distributions and high bias/variance from stochastic interactions. It introduces Optim4RL, which uses pipeline training (training multiple agents in parallel with regular resets) and an inductive-bias update structure that mirrors adaptive optimizers to stabilize meta-learning and enable learning from scratch. Empirically, Optim4RL generalizes to unseen Brax tasks and outperforms several baselines, including VeLO, while remaining simpler and more robust to training instability. This work provides a practical, data-driven pathway to RL-specific learned optimizers with notable generalization potential and efficiency advantages over prior approaches.

Abstract

In recent years, by leveraging more data, computation, and diverse tasks, learned optimizers have achieved remarkable success in supervised learning, outperforming classical hand-designed optimizers. Reinforcement learning (RL) is essentially different from supervised learning, and in practice, these learned optimizers do not work well even in simple RL tasks. We investigate this phenomenon and identify two issues. First, the agent-gradient distribution is non-independent and identically distributed, leading to inefficient meta-training. Moreover, due to highly stochastic agent-environment interactions, the agent-gradients have high bias and variance, which increases the difficulty of learning an optimizer for RL. We propose pipeline training and a novel optimizer structure with a good inductive bias to address these issues, making it possible to learn an optimizer for reinforcement learning from scratch. We show that, although only trained in toy tasks, our learned optimizer can generalize to unseen complex tasks in Brax.

Figures (5)

• Figure 1: Visualizations of agent-gradient distributions (a) at the beginning of training, (b) in the middle of training, and (c) at the end of training. All agent-gradients are collected during training A2C in big_dense_long , optimized by RMSProp. We compute $\log(|g|+10^{-8})$ to avoid the error of applying $\log$ function to non-positive agent-gradients.
• Figure 2: (a) An example of pipeline training where the reset interval $m=3$ and the number of units $n=3$. All training units are reset at regular intervals to diversify training data. (b) The network structure of Optim4RL. $g$ is the input agent-gradient, $h_i$ and $h'_i$ are hidden states, $\alpha$ is the learning rate, $\epsilon$ is a small positive constant, and $\Delta \theta$ is the parameter update.
• Figure 3: The optimization performance of different optimizers in four RL tasks. Note that the performance of VeLO is estimated based on Figure 11 (a) in metz2022velo. All other results are averaged over $10$ runs, and the shaded areas represent $90\%$ confidence intervals. Optim4RL is the only learned optimizer that achieves satisfactory performance in all tasks.
• Figure 4: Optim4RL shows strong generalization ability and achieves good performance in Brax tasks, although it is only trained in six simple gridworlds from scratch. For comparison, VeLO metz2022velo is trained for $4,000$ TPU-months with thousands of tasks but only achieves sub-optimal performance in Ant. These results demonstrate the generalization ability of Optim4RL in complex unseen tasks, which is a significant achievement in itself, proving the effectiveness of our approach.
• Figure 5: Visualizations of gradient distributions (a) at the beginning of training, (b) in the middle of training, and (3) at the end of training. All gradients are collected during training in MNIST, optimized by RMSProp.