Debiasing Meta-Gradient Reinforcement Learning by Learning the Outer Value Function
Abstract
Meta-gradient Reinforcement Learning (RL) allows agents to self-tune their hyperparameters in an online fashion during training.
In this paper, we identify a bias in the meta-gradient of current meta-gradient RL approaches.
This bias comes from using the critic that is trained using the meta-learned discount factor for the advantage estimation in the outer objective which requires a different discount factor.
Because the meta-learned discount factor is typically lower than the one used in the outer objective, the resulting bias can cause the meta-gradient to favor myopic policies.
We propose a simple solution to this issue: we eliminate this bias by using an alternative, outer value function in the estimation of the outer loss.
To obtain this outer value function we add a second head to the critic network and train it alongside the classic critic, using the outer loss discount factor.
On an illustrative toy problem, we show that the bias can cause catastrophic failure of current meta-gradient RL approaches, and show that our proposed solution fixes it.
We then apply our method to a more complex environment and demonstrate that fixing the meta-gradient bias can significantly improve performance.
Introduction
Recently, reinforcement learning (RL) methods have embraced ideas from meta-learning either to train an agent to quickly adapt to new tasks [6, 12] or to improve the learning process online within a single task [27, 29, 23, 26, 28].
RL algorithms are known to be highly sensitive to their hyperparameters such as the discount factor [1].
Thus, self-tuning these hyper-parameters online can greatly improve performance and has yielded state-of-the-art performance for model-free RL [7].
Meta-gradient RL is the backbone underlying self-tuning and relies on an inner loss function used to update the agent’s parameters, as well as an outer loss to evaluate (or ground) the updated parameters in order to compute a meta-gradient.
This meta-gradient is then used to update the meta-parameters (e.g. the discount factor).
In this paper we focus on the outer loss used by current meta-RL algorithms and show that its estimation has a bias.
This bias comes from using the critic that is trained using the meta-learned discount factor for the advantage estimation in the outer loss.
Meta-gradients obtained using this biased estimate can lead to a failure of self-tuning, potentially resulting in a myopic policy.
We propose a computationally cheap approach that eliminates the meta-gradient bias and we observe that it solves catastrophic failure in a toy environment and improves performance on a more challenging deep RL task.
There have been several works focusing on bias and variance in meta-gradients.
(Bonnet et al., 2021)
[5] focuses on bias-variance trade-off within multi-step meta-gradients.
(Vuorio et al., 2022) [24] studies a different bias, namely the sampling bias, and (Liu et al., 2021)
[15] addresses a [Figure] compositional bias.
Although these aforementioned works are complementary to ours, the bias we highlight is different as it is related to the outer loss of meta-gradient algorithms.
Background
Reinforcement Learning : The goal of an RL agent is to learn the policy π that maximizes the γ -discounted expected return: E π | s t “ s
[Equation]
‰ , where r t and s t are the reward and state at timestep t .
Actor-critic algorithms [13, 17] learn a policy π θ p with parameters θ p and a critic V π,γθ c with parameters θ c that approximates the value function V π,γ “ E π | s t “ s “ř [Equation] k
[Equation]
The critic is used to inform the direction of the policy update.
More specifically, the policy is trained using the policy gradient [25, 20], ˆ g pg , estimated by Monte Carlo using
[Equation]
where τ “ t s t , a t , r t , s t ` 1 u are transitions.
One could use the advantage [20] instead of the action value function in the estimate to reduce variance [8, 18], but we stick to the action value function for simplicity.
The full actor-critic update combines the policy gradient with entropy regularization [21, 17, 10, 9] and a critic update using the mean squared error of the TD error [22].
We provide further discussion and definitions in Appendix A.
Self-Tuning RL : We refer to the aforementioned actor-critic update as the inner update with loss L inner p θ, η q where η are the hyper-parameters.
Self-tuning meta-RL purposely exposes η as trainable meta-parameters and adapts them online during training by computing meta-gradients.
This makes use of the fact that the inner update is differentiable with respect to η .
The meta-gradient ∇ η L meta is defined as the gradient of a meta-loss with respect to the meta-parameters.
It is computed in two phases (see figure 1): (1) The agent takes one (or multiple) parameter update(s) by running the actor-critic algorithm with meta-parameters η to obtain updated parameters θ 1 p η q .
The meta-loss is computed from the updated parameters θ 1 p η q to derive the meta-gradient ∇ η L meta p θ 1 p η qq .
One can distinguish two kinds of meta-loss (see figure 1).
Firstly, meta-gradients (MG) [27, 26, 28, 23] compute the gradient of an outer loss L outer applied to updated parameters θ 1 p η q .
Secondly, bootstrapped metagradients (BMG) [16, 7] further update the parameters by taking one (or several) gradient step(s) to obtain a target ˜ θ from which to compute a matching loss between θ 1 p η q and ˜ θ .
The target updates proposed by [7] can be done with respect to the inner loss L inner for all but the last step, which is computed with an outer loss L outer .
Both methods of computing meta-gradients necessitate an outer loss, L outer , that is defined to be an actor-critic loss similar to L inner but uses different hyper-parameters η 1 than the inner loss.
In particular, the policy gradient term of the outer loss uses a different discount factor γ 1 , leading to the following policy gradient estimate ˆ g 1 pg .
[Equation]
Solution to the meta-gradient bias
In this section, we highlight a bias in the meta-gradient estimation of current meta-RL methods and propose a simple and computationally-efficient way to solve it.
Bias in the meta-gradient : Both MG and BMG methods rely on the estimation of the outer loss to compute the meta-gradient (see figure 1).
As seen in section 2, the policy gradient terms from the inner and outer loss differ in terms of what hyper-parameters they use.
While the inner loss necessitates estimating the value function V π,γ , the outer loss requires V π,γ 1 .
Since the critic is trained to estimate the value function V π,γ , a bias consequently arises when using the critic to estimate the advantage in the policy gradient term of the outer loss:
[Equation]
The right-hand side of equation 3 is the biased estimate of the outer policy gradient used by current self-tuning methods.
Removing the bias with an outer-critic head : We propose to learn another critic that we call the outer-critic , whose goal is to estimate the outer value function, i.e. the value of the policy with (constant) discount γ 1 .
More precisely, the outer-critic V π,γ 1 θ ¯ c , parameterized by θ ¯ c , is trained to approximate the outer value function : V π,γ 1 p s q “ E π | s t “ s ”řk
[Equation]
1 k r
[Equation]
ı .
This way, one can use the outer-critic to remove the aforementioned bias in the estimation of the outer-loss: ´ r ` γ 1 V π,γ 1 θ ¯ c p s 1 q ¯ log π θ 1 p p a | s q .
To implement the outer-critic, we modify current critic architectures by adding a second head whose goal is to learn the outer value function V π,γ 1 .
As a consequence, our way of removing the bias in the outer loss advantage estimation is computationally cheap and easy to implement on top of current self-tuning methods to improve meta-gradient estimation.
Further details are provided in appendix B.
Experiments
We demonstrate the catastrophic failure of current meta-gradient approaches in a toy problem, then we scale to a deep RL experiment where we observe that our method improves performance over biased meta-gradients.
It is important to note that for these self-tuning experiments, we initialize γ to a low value (0.95 and 0.8), in fact too low for a converged policy to be optimal (both experiments require optimizing over long horizons). This is to accentuate the effect of the meta-gradient bias.
We open-source the code to reproduce all the experiments, see https://github.com/instadeepai/ outer-value-function-meta-rl .
4.1 Discounting Chain
For our first experiment, we provide a simple illustration of the effects of the bias caused by using the inner critic in the outer loss.
For this purpose, we experiment on the Gymnax [14] JAX implementation of the toy Discounting Chain environment originally proposed in bsuite 2 [19].
The Discounting
Chain environment, shown in Figure 2, has a structure designed to highlight the effect of different discount horizons and is therefore highly sensitive to the discount factor.
Since we have access to the true value function for this environment, we use it in the policy gradient advantage when training the policy.
Having the true value function allows us to focus on the bias induced by using the wrong value of γ in the outer loss.
We train the MG and BMG agents with and without the outer critic, with an initial γ “ 0 . 95 .
We apply a modified sigmoid activation between 0 . 9 and 1 . 0 to ensure the discount remains bounded.
We provide further details on the experimental setup in Appendix C.1.
In figure 2, the standard MG and BMG agents fail catastrophically, as they are unable to increase the discount factor and therefore converge to the myopic policy.
The reason is that if the inner-critic is used in the outer-loss advantage, then myopic policies are preferred (because V π,γ ă V π,γ 1 ), which then causes the meta-gradient to push γ down instead of up.
On the other hand, the MG and BMG agents equipped with the outer critic are able to successfully self-tune the discount factor to converge to the optimal policy and solve the task.
4.2 Snake
Next, we scale our experiments to the deep RL setting using the Snake environment from Jumanji [4, 5].
This environment is sensitive to the discount factor, as the snake agent must plan ahead to maximize return without bumping into itself.
We compare an MG algorithm with and without the outer critic.
Unlike in the Discounting Chain experiment, where we had access to the true value function, we now train a value function for the inner and outer critic.
We give γ a relatively poor initialization of 0 . 8 , as this makes the self-tuning more challenging.
We provide further details for the setup of the experiment in Appendix C.2.
In figure 3, the agent with the second critic shows a stronger meta-gradient signal, which results in γ being increased faster.
Although both agents asymptotically converge to the maximum return, fixing the meta-gradient bias results in a significant improvement in training performance.
Conclusion
We have shown that current meta-gradient RL algorithms suffer from a meta-gradient bias induced by having the critic, trained using the meta-learned discount factor, estimate the advantage of the outer objective which necessitates a different discount factor.
In our experiments, we have demonstrated that this meta-gradient bias can lead to catastrophic failure of meta-gradient RL algorithms, which we fix with a simple, yet efficient solution.
We augment the critic network with an outer critic head whose goal is to estimate the value function discounted by the outer discount factor.
Furthermore, we have shown that our method results in performance improvements in more complex environments.
Bootstrapped meta-gradient [7], the state-of-the-art meta-gradient RL agent, has worse than human performance on Bowling, Solaris, and Skiing from the Atari ALE benchmark.
These environments are known to be challenging due to delayed reward/credit assignment [3, 2].
Thus, correcting the bias we highlight in this work may be important for improving meta-gradient RL approaches in these increasingly challenging environments, which we hope to study in future work. Acknowledgements
Research supported with Cloud TPUs from Google’s TPU Research Cloud (TRC).
We would like to thank Nathan Grinsztajn for their helpful discussions
A Background
The goal of an RL agent is to maximize the γ -discounted state value function: V π,γ p s q “ E π | s t
[Equation]
s
[Equation]
The action value function is analogously defined below.
[Equation]
The policy gradient theorem [25, 20] in equation 4 provides a way to derive the gradient of the expected discounted return using the value function of the policy.
[Equation]
Actor-critic algorithms [13, 17] learn a parameterized value function (the critic) and a parameterized policy which is updated in a direction informed by the critic.
More precisely, the critic learns the value function V π,γθ c with critic parameters θ c , and policy π θ p with policy parameters θ p .
Estimated action value function ˆ Q π,γ p s, a q “ r ` γV π,γθ c p s 1 q where r and s 1 are respectively the reward and the next state obtained from taking action a in state s .
The goal of an RL agent is to learn the policy π that maximizes the γ -discounted expected return: E π | s t “ s
[Equation]
‰ , where r t and s t are the reward and state at timestep t .
An actor-critic update (A2C)
[17] is composed of 3 terms:
• Policy gradient:
[Equation]
• Critic gradient:
[Equation]
• Entropy regularization: ∇ θ p H p π θ
[Equation]
p
[Equation]
s qq
Therefore, the full actor-critic update is: θ Ð θ ` αf p θ, η, τ q where θ “ t θ p , θ c u and η “ t γ u or η “ t γ, λ u if TD( λ ) is used for advantage estimation, and f p θ, η, τ q “ c PG
[Equation]
B Outer Critic
To learn the outer value function, we augment the critic network with a second head that outputs the value with discount γ 1 .
To avoid conflicting gradients between the two value functions being learned, we use a stop-gradient between the common torso and the outer-critic head (see figure 4).
This way, the critic learning dynamic is not changed by adding the outer-critic.
To learn the outer value function, we proceed akin to the standard value function with a TD loss l outerTD with γ 1 instead of γ .
[Equation]
C Experiments C.1 Discounting Chain Experiment Details
We initialize γ “ 0 .
95 such that the value function will initially favor the myopic policy.
For simplicity, we use only a policy gradient loss for both the inner and outer loss functions.
Furthermore, for BMG we do a single gradient step with the outer loss to obtain the target, which we use for policy matching.
We optimize the parameters of the actor for MG and BMG with SGD and use Adam [11] to optimize the discount factor γ .
For all agents, we use a policy network composed of a single linear layer.
We generate a new batch of data online for both the inner and outer loss functions.
Using an adaptive optimizer (e.g. Adam or RMSProp) to optimize the parameters would change the meta-gradient for BMG since there would be some inertia from the inner update biasing the direction of the target update.
As we would like our inner update step to be in the direction of the unbiased meta-gradient, we use SGD instead of Adam.
More details on the hyper-parameters used in this experiment can be found in table 1.
To solve the environment, the self-tuning agent has to increase the discount factor to a sufficient level so that it can receive a signal from the long-term effect of the first action.
However, if the meta-objective is sufficiently biased, the meta-gradient will not encourage the agent to increase the discount.
C.2 Snake Experiment Details
We use the environment "Snake-6x6-v0" from Jumanji [4] by calling Jumanji’s registry: jumanji.make("Snake-6x6-v0") .
In this environment, a snake agent navigates a grid in order to eat apples, while avoiding bumping into its own body which causes it to die.
Because good policies require the agent to zigzag and reason about long horizons, the discount factor plays a very important role in the learning dynamics and is well suited to studying any meta-gradient biases.
The hyper-parameters used in the experiments are provided in table 1.
D Outer Loss Advantage D.1 Quantify the Meta-Gradient Bias
One way to quantify the meta-gradient bias we highlight in this paper is to observe the advantage in the outer loss. By definition, the advantage is centered, i.e. its expectation over actions is 0. However, we can see in figure 5 that current self-tuning methods (MG, BMG) suffer from a bias in their advantage estimation. Using the outer value function in the outer loss fixes the bias and the advantage remains centered.
[Table]
Discounting
[Table]
Snake
[Table] [Figure]
D.2 Normalizing Advantages
We show that the advantages in the outer loss are not centered, which leads to a meta-gradient bias. One could argue that normalizing the advantages across a batch would fix this issue. We show in this section that this is not the case.
Indeed, in a batch, all the samples’ advantages are biased.
Removing this offset on the batch level does not fix the direction of the meta-gradient. This leads to the normalized advantages being centered but the meta-gradient still keeps its bias as shown in figure 6.