M ETA -G RADIENTS IN N ON -S TATIONARY E NVIRONMENTS
Yannick Schroecker
DeepMind Tom Zahavy DeepMind Satinder Singh
DeepMind
A BSTRACT
Meta-gradient methods (Xu et al., 2018; Zahavy et al., 2020) are a promising approach to the problem of adaptation of hyper-parameters in non-stationary reinforcement learning problems.
Recent works enable meta-gradients to adapt faster and learn from experience, by replacing the tuned meta-parameters of fixed update rules with learned meta-parameter functions of selected context features (Almeida et al
Flennerhag et al., 2022). We refer to these methods as contextual meta-gradients .
The context features carry information about agent performance and changes in the environment and hence can inform learned meta-parameter schedules.
As the properties of meta-gradient methods in non-stationary environments have not been systematically studied, the aim of this work is to provide such an analysis.
Concretely, we ask: (i) how much information should be given to the learned optimizers so as to enable faster adaptation and generalization over a lifetime, (ii) what meta-optimizer functions are learned in this process and (iii) whether meta-gradient methods provide a bigger advantage in highly non-stationary environments.
We find that adding more contextual information is generally beneficial, leading to faster adaptation of meta-parameter values and increased performance.
We support these results with a qualitative analysis of resulting meta-parameter schedules and learned functions of context features.
Lastly, we find that without context, meta-gradients do not provide a consistent advantage over the baseline in highly non-stationary environments.
Our findings suggest that contextualising meta-gradients can play a pivotal role in extracting high performance from meta-gradients in non-stationary settings
I NTRODUCTION Meta
-gradient approaches to learning adaptive optimizers are a promising complement to handengineering gradient-based optimizers in reinforcement learning (RL).
By adapting relevant optimization hyperparameters or the entire update rule to the current domain, they often outperform well-tuned gradient-based optimizers (Schraudolph, 1999; Mahmood et al., 2012; Andrychowicz et al., 2016;
Xu et al., 2018).
The adaptability of meta-learned optimizers is particularly relevant for non-stationary environments: as the agent continues learning during its lifetime, appropriate hyperparameter values for an optimizer are likely to change over time and can be impossible to determine in advance (Parker-Holder et al., 2022).
Despite this promise, the performance and properties of meta-gradients in the context of reinforcement learning in non-stationary environments have not been systematically studied, which is the focus of our work.
The question we focus on is how does the information provided to the meta-optimizer as well as the rate of environment non-stationarity effect the performance and properties of meta-gradients.
Most of the prior work in this area can be classified within the framework of white and black box optimization (see Related Work in Xu et al.
for a more detailed classification of prior work).
Black-box methods express the entire update rule as a rich parametric function, typically a recurrent neural network, and learn the parameters of this function in an end-to-end manner (Xu et al., 2020;
Oh et al., 2020; Kirsch et al., 2021).
In contrast, white-box methods tune the hyperparameters of the optimization algorithm (Mahmood et al., 2012; Xu et al., 2018; Zahavy et al., 2020).
We limit the scope of our analysis to white-box meta-gradient methods for self-tuning hyperparameters, which have to date shown the greatest empirical gains in RL (Xu et al., 2018; Zahavy et al., 2020;
Flennerhag et al., 2022).
Since white-box methods are almost completely memory-less and only tune several hyperparameters based on their local influence on agent performance, standard white-box meta-gradients are only capable of tracking good solutions (Sutton et al., 2007) and lack the ability to learn and benefit from past experience.
These problems are partly addressed in some recent works in reinforcement (Flennerhag et al., 2020) and supervised (Almeida et al., 2021) learning, which extend the standard white-box formulation by replacing learned hyper-parameters with learned functions of hand-picked context features.
Broadly speaking, context features can be any low-dimensional statistics of optimization, agent, or environment that carry information about suitable hyperparameter schedules.
Examples of such context features in RL are reward histories, temporal difference errors, and task beliefs.
We refer to this approach as contextual meta-gradients .
While sharing some similarities with black-box meta-gradients, contextual meta-gradients differ in two important ways: the space of learned updated rules is limited (the update rule can only be changed through hyperparameters) and they have access to different kind of information (selected context features instead of history of gradients or parameters).
In this paper, we hypothesize that: (1) the addition of context to meta-gradients is particularly wellsuited for non-stationary environments with some repeated structure, as it enables the optimizer to generalize from previously seen contexts and instantly pick good hyper-parameter values; (2) the advantage of all meta-gradients methods (with or without context) over well-tuned fixed hyperparameter schedules increases with the rate of environment non-stationarity, as different environment conditions may require very different learning hyperparameter values.
To examine the first hypothesis, we compare the performance of meta-gradients with and without contextual information across several environments.
Since the addition of contextual information can enable the meta-learner to learn a hyper-parameter schedules to generalize across learning contexts, we look at how increasing the context richness (i.e. amount of information given to the meta-learner) effects the performance.
We find that adding more information to the meta-optimizer most of the time helps training (although some context are much more useful than others), and in some cases, it is necessary for learning quickly adapting meta-parameter schedules.
We complement these findings by examining learned schedules and functions of context, which demonstrate faster adaptation of hyperparameter values and emergence of meaningful functions of context.
Lastly, we investigate the adaptation ability of meta-gradients as the rate of environmental non-stationarity increases, with a particular interest in potential advantages of context in highly non-stationary environments.
We find that without contextual information, meta-gradients provide small to inconsistent advantage over fixed hyperparameter values in highly non-stationary environments.
On the other hand, meta-gradients with contextual information provide a much more consistent advantage.
R ELATED W ORK
The focus of our paper is on white-box methods for learning adaptive optimizers (Bengio, 2000; Maclaurin et al., 2015). These methods tune the meta-parameters (i.e. tunable subset of hyperparameters) of the learning update by computing the gradient of the outer loss with respect to metaparameters, where the outer loss depends on meta-parameters both directly and through the history of last K parameter updates.
In RL, this approach has been used to tune various optimization hyperparameters, including discount and bootstrap parameters (Xu et al., 2018), off-policy corrections (Zahavy et al., 2020), auxiliary rewards and tasks (Zheng et al., 2018; Veeriah et al., 2019; 2021) and weights of rewards in return estimates (Wang et al., 2019).
Among these methods, Flennerhag et al.
and Almeida et al.
(2021) propose learning metaparameters as functions of context features, with the application in reinforcement and supervised learning respectively.
While the focus of Flennerhag et al.
(2022) is on reducing myopia and conditioning problems of meta-gradients by developing an improved outer loss, in their experiments on non-stationary environments, they parameterize meta-parameters as a function of contextual features (in their case, reward statistics).
However, they do not study the importance of including contextual information or provide comparison to alternative contextual features.
In Almeida et al.
(2021), learning meta-parameters as function of features is the primary focus.
In contrast to most prior work in white-box meta-gradients, Almeida et al.
(2021) formulate optimization of meta-parameters as a reinforcement learning problem.
The meta-parameter function is a policy modifying the current values of corresponding meta-parameters, trained on a distribution of source tasks with a designed reward function.
The context features are selected to facilitate transfer of learned optimizers to the target tasks.
In contrast to our work, their focus is on achieving high performance on supervised learning tasks instead of analysis of meta-gradients in non-stationary RL environments.
For related work in black-box meta-gradients and the use of context outside of work in metagradients, see Appendix A.1.
B ACKGROUND
We follow the meta-gradient literature and make a distinction between the parameters θ (e.g. weights of the policy and value networks) and the meta-parameters η (e.g. learning rates, discounts, weights of regularization losses).
The parameters θ are trained by minimizing an inner loss function
L ( inner ) η
( θ, D ) parameterized by the meta-parameters η , whereas the meta-parameters η are trained to minimize the outer loss L ( outer ) ( η, D ) , where in both cases D refers to the rollout data used to estimate the loss.
The parameters and meta-parameters are generally optimized at two different time-scales; we denote by i the inner loop and by k the outer loop iteration step.
The outer loss objective depends on the meta-parameters η through their influence on the learned parameters θ .
This dependence is given by the update rule:
θ i +1 =
f ( θ i , η k , D i ) , where D i refers to the data used to compute the losses at i -th iteration.
For example, for stochastic gradient descent, θ i +
=
θ i − lr ∇ θ L ( inner )
η ( θ ) , which is parameterized by η through L ( inner )
η ( θ ) .
Since backpropagating through the entire history of updates is too computationally demanding and may not provide a good gradient estimate due to non-stationarity intrinsic to RL, the meta-gradients are truncated to the last
K inner loop updates, which we refer to in experimental section as meta rollout length .
The outer losses can be computed on a new rollout data D k or the rollout data used in the inner loop D i .
A more detailed description of common inner and outer objectives can be found in the Appendix A.2.
C ONTEXTUAL M ETA -G RADIENTS
Contextual meta-gradients extend white-box meta-gradients by replacing the meta-parameter variables
η k in the update rule with parameterized functions η k ( · ) of context features c i : f
( θ i , η k ( c i ) , D i ) .
The meta-parameters η of the meta-parameter function η ( · ) are learned by optimizing the outer loss L ( outer ) ( η ) .
To avoid the overloaded notation and terminology, throughout this paper, we use η k and the term meta-parameter to denote two related quantities: the parameters of the meta-parameter network, as well as the predictions of the meta-parameter network, at outer iteration k .
It will be clear from context which one we are referring to.
Also note that depending on how the context features are shared, this change allows for different meta-parameter predictions between individual inner updates or even individual samples contributing to the inner loss.
Similar ideas were explored in Flennerhag et al. and Almeida et al.
The context features can be any learning or environment statistics that capture information relevant for optimal meta-parameter schedules.
We do not assume that context features will be repeated during a lifetime, however choosing context features that do repeat should provide contextual metagradients with a greater advantage over non-contextual meta-gradients.
If such repetition exists, meta-parameter functions can learn and generalize schedules from previously seen contexts.
In our experiments, the features were chosen to indicate changes in task at hand and the agent performance (e.g. statistics of rewards, TD-errors, values), all of which may require different meta-parameter values.
Lastly, we note that gradients do not propagate through the context features.
E NVIRONMENTS
For our analysis, we use two non-stationary environments.
The agent interacts with these environments via a single-stream of experience via the standard RL formulation (there is only one environment and not multiple copies of it, as is common in distributed RL frameworks).
The simplicity of the environments allows for a fine control over the sources of non-stationarity and makes it easier to interpret learned meta-parameter schedules and functions.
More information about environments can be found in Appendix A.4.
Two Colors.
We start with a gridworld environment introduced in Flennerhag et al.
(2022), where the agent is tasked with picking up one of the two items.
Picking up one of the items results in reward[-+
-]
+1 or -1, after which the agent and two objects are randomly re-spawned.
Which object carries positive reward changes every 100,000 steps.
The agents used in our experiments are memory-less, hence the optimal policy has to be re-learned after each task switch, which requires increasing the rate of exploration after the task switches.
Switching MDPs.
In our second non-stationary setting, the agent interacts with a sequence of grid worlds, where a new grid world is sampled every 100,000 steps from a set of N grid worlds.
The reward function and transition function of each grid world are randomly generated.
We use this environment to study how changes in the environment dynamics changes, in addition to the reward function, affect meta-gradient methods.
Additionally, in contrast to the Two Colors environment, two consecutive tasks may not be as different from each other and contain some positive transfer.
This implies that in some cases, the optimal exploration strategy upon a task change may not be to explore too much.
We experimented with two variants of Switching MDPs with different number of grid-worlds N .
In the first case N =4 MDPs so that the agent is repeatedly exposed to the same tasks, and can improve its learning over time.
In the second case, N =1000 MDPs, and the agent is very unlikely to experience the same task twice, requiring constantly learning almost from scratch in each MDP.
This case is interesting, because we can study the generalization of contextual meta-gradients to unseen contexts.
As tasks effectively do not repeat, distribution of context features that are given to the meta-parameter function after each task switch is much more irregular.
For example, we do not observe regular periodic drops and increases in mean reward as we do in Two Colors.
E XPERIMENTS
We designed the experiments with the goal of answering the following questions about the behaviour of meta-gradients in non-stationary environments:
Do meta-gradients benefit from contextual information?
We hypothesise that ability to learn meta-parameter schedules as functions of contextual information will enable faster adaptation in non-stationary environments, as the optimizer can leverage knowledge from previously seen contexts and instantly use good meta-parameter values without needing to slowly tune them.
We compare the performance of contextual meta-gradients and baselines in Section 6.1.
What functions of context are learned in this process?
To explain the performance difference between methods, we examine the meta-parameter schedules during training and the learned metaparameter functions of context in Section 6.2.
What should the contextual information be?
The choice of context features could be essential for learning schedules that generalize over the training.
Here we look into increasingly rich contexts: what is the effect of adding particular candidate contexts and can too much information be detrimental for generalization?
We shed light on these questions in Section 6.3.
How does the performance of meta-gradients depend on the rate of non-stationarity?
As the rate of change of the environment increases, we hypothesized that the ability to adapt the metaparameters becomes more important.
Furthermore, the addition of context to meta-gradients should lead to further advantages, as the they enable faster adaptation.
The advantage of meta-gradients under different rates of non-stationarity is examined in Section 6.4.
E XPERIMENTAL A XES
To ensure that the conclusions we draw are robust, we run the experiments along several axes: (i) agents in inner loop, (ii) outer loop objectives, (iii) context features.
Agents.
We use two different kinds of RL agents: Actor-Critic (AC)
(Sutton et al., 1999) and Q ( λ ) (Peng & Williams, 1994).
The parameters of an AC agent are updated every 16 environment steps, and those of Q ( λ ) agent at each step.
If using AC agent, we tune the coefficient (i.e. weight) of entropy loss in the inner loss, whereas if using Q ( λ ) agent, we tune the ϵ parameter of ϵ -greedy exploration.
The details of the agent architectures and objectives can be found in Appendix A.5.
Outer Loop Objectives.
In experiments combining meta-gradients (MG) with AC agents, the outer loss is a sum of policy loss L π and a target entropy loss L ent (the coefficient of this entropy loss is a fixed hyperparameter α ( outer ) ent ): [Table]
[Equation]
If using Bootstrapped Meta-Gradients (BMG) (Flennerhag et al., 2022) [Table] the outer loss is the KL divergence between the target policy
π ˆ θ and K -step bootstrap π θ K :
[Equation]
where the target policy is the policy reached after K + L − 1 steps of inner loop optimization.
For a longer description of BMG loss and how it can be made differentiable with respect to ϵ when used with Q ( λ ) agents, see Appendix A.3.
Context Features.
We compare two kinds of contextual features, simple features referred to as
Reward (only using reward statistics) and more complex features referred to as Rich (using reward, TD-error and value statistics).
If the context is Reward , each context feature is either a history of the last H observed rewards ( Q ( λ ) agent) or a history of mean rewards observed in each of the last H rollouts (AC agent).
If the context is Rich , we compute the same history for reward, TD-error and values.
In AC experiments, rich context also includes a history of last H standard deviations in addition to means.
We use H = 10 in experiments with AC agents, and H = 100 in experiments with Q ( λ ) agents.
Each context feature is normalized and scaled to be in the [-1, 1] range, before being concatenated and passed down to the learned meta-parameter function.
More detailed description of context features used in our experiments can be found in Appendix A.6.
In each experiment, we report only the results for the best hyper-parameter setting .
For the baselines (which do not adapt the meta-parameters), we include the meta-parameter of interest into the hyperparameter sweep.
For MG objectives, we sweep over the meta learning rate, meta rollout length ( K ) and target entropy loss coefficient ( α ( outer ) ent ).
For BMG methods, we sweep over meta learning rate, meta rollout length ( K ) and target rollout length ( L ).
For fairness, the size of hyperparameter sweeps is the same for all meta-gradient methods.
For more details on the hyper-parameter sweeps and values, see Appendix A.5.
6. D O M ETA -G RADIENTS B ENEFIT FROM A DDING C ONTEXTUAL I NFORMATION ?
We start the experiments by testing how adding contextual information to meta-gradients effects training.
If our hypothesis is true, we expect to see the advantages of adding contextual information across different agents and outer loop objectives
Two Colors: Actor Critic Agent.
We first look at the AC agents trained on Two Colors (described in Section 5).
In Table 1 (left), we report the mean and standard deviation of the total reward after 10M environment steps.
We find that all meta-gradient methods outperform the baseline with fixed meta-parameters (AC) and, consistent with reports in Flennerhag et al. (2022), BMG performs better across methods.
More importantly, the addition of context (see columns with Context Features set to Reward or Rich) is crucial for obtaining significant improvement over the baseline, with the richer features performing slightly better.
Two Colors: Q ( λ ) Agent.
Next, we look at the Q ( λ ) agents on Two Colors.
In Table 1 (right), we again report the mean and standard deviation of total returns after 10M environment steps.
We tune the ϵ parameter of ϵ -greedy exploration with only BMG objective, as the updated value function in this case is not a differentiable function of ϵ and consequently can not be straightforwardly optimized with regular MG objectives. [Table]
We find that meta-gradients significantly outperform the non-adaptive baseline and the addition of context further boosts the performance.
Switching MDPs: Q ( λ ) Agent .
Next, we look at Q ( λ ) agent trained on Switching MDPs.
In Table 2, we report the mean and standard deviation of total rewards after 20M environment steps.
We report only the experiments with Q ( λ ) agents since AC agents were performing very poorly on this environment.
Again, we tune only the ϵ parameter.
For these experiments, we also vary the number of different MDPs available in the environment (see Section 5 on meaning and importance of these environment variations).
We compare the results under the two environment variants of interest (4 and 1000 MDPs) in Table 2.
The addition of context was necessary to obtain significant improvements with meta-gradients, with
Rich context resulting in the biggest improvement.
The improvement is also much more significant in the regime where the MDPs effectively do not repeat (1000 MDPs).
Because the consecutive tasks will typically be more similar compared to Two Colors, we hypothesize there is less need to re-learn in the regime with a small number of tasks (4 MDPs), requiring less adaptation of meta-parameters.
6. W HAT M ETA -P ARAMETER S CHEDULES AND F UNCTIONS ARE L EARNED ?
In this section, we wish to examine the predictions of the learned meta-parameter functions throughout training.
For the ease of analysis, we focus on BMG with Reward context features (BMGReward) while only tuning one meta-parameter.
Learned Schedules.
We look at the relationship between the learned meta-parameter schedule and the observed rewards for AC agents with BMG trained on Two Colors.
In Figure 1a, meta-gradients do not rely on context features
[Figure]
and in Figure 1b, meta-gradients use Reward context features.
The two curves in both figures are: entropy loss coefficient during training (orange curve) and mean reward over rollout (blue curve).
When meta-gradients do not have access to reward context (Figure 1a) [Figure] the mean reward takes longer to recover after the drop following each task switch.
Furthermore, entropy coefficient is almost constant.
Note that this can not be explained by too low meta learning rate: the results shown here [Figure] [Figure] o e ff i
c i e n t
[Figure]
[Figure] [Figure] are under the best hyper-parameter configuration and the best performing meta learning rates were never the highest values in the sweep. In contrast, when the reward context is added (Figure 1b), the entropy coefficient rapidly increases after each task switch to allow the agent to explore.
As a result, the mean rewards of the agent are better as we observed in the previous section, suggesting that the addition of context is crucial for obtaining fast adaptation [Figure]
Learned Functions.
Next, we look at how the context to meta-parameter mapping itself changes during a lifetime.
Our methodology is inspired by visualization of model trajectories in Erhan et al.
Due to low dimensionality of inputs and outputs of meta-parameter function, we can design a small number of representative context inputs and for each, track the corresponding meta-parameter prediction during training.
Note that the meta-parameter function is not trained on these context inputs, we just log the output for each while the meta-parameter function is trained as usual.
The five context inputs were designed to represent qualitatively different learning situations: constantly high context values ( ”high” ), monotonically increasing from lowest to highest value ( ”increasing” ), constant zero ( ”zero” ), monotonically decreasing between the highest and lowest value ( ”decreasing” ) and constantly low context ( ”low” ).
For example, for rewards bounded in [ − 1 , 1] and H = 3 , the context inputs are: ”high” =
[1 , 1 , 1] , ”increasing” =
[ − 1 , 0 , 1] , ”zero” =
[0 , 0 , 0] , ”decreasing” =
[1 , 0 , − 1] and ”low” =
[ − 1 , − 1 , − 1] .
In our experiment, H = 10 and the range is [ − 1 , 1] due to pre-processing of context features.
Note that some of the context inputs have the same mean (e.g., ”increasing” , ”zero” and ”decreasing” all have mean zero), however they capture qualitatively different behaviors.
For example, the ”decreasing” input corresponds to a context likely observed just after switching the task, while the ”increasing” probe is likely observed as the agent starts improving in a new task.
We visualize the learned meta-functions of AC agent trained with BMG-Reward on Two Colors in
Figure 2a, and those of Q ( λ ) agent trained with BMG-Reward on Switching MDPs (1000 MDPs variant) in Figure 2b.
First, we can observe that different inputs results in very different metaparameter predictions and that the meta-parameter functions seem to converge during training (the small local changes are likely due to non-stationary training distribution), hence addition of context enables convergence instead of just tracking good meta-parameter values as for vanilla MG (Sutton et al., 2007).
Next, we look at the values of predicted meta-parameters for each of the five context inputs.
We can observe that the differences in predictions are sensible: in Figure 2a, the entropy loss coefficient is highest when the rewards are at low or decreasing, and the lowest when the rewards are high or decreasing; whereas in Figure 2b, the exploration is highest when the rewards are low, decreasing or zero and lowest when the rewards are high.
The predictions are different for context inputs with the same mean value (”increasing”, ”decreasing”, ”zero”), indicating learned function responds to more complex patterns than just mean values of features.
Lastly we note that for Switching MDPs, we see more variability in learned functions across different training seeds.
This is likely due to a more complex interaction between encountered tasks and stochasticity in generating and sampling tasks (i.e. how much transfer there is between consecutive tasks will depend on which two tasks are sampled). [Figure]
6. H OW I MPORTANT IS THE C HOICE OF C ONTEXT ?
In the following set of experiments, we inspect how making the context more rich effects performance.
For this goal, we focus on Two Colors with the AC agent and BMG objective.
We consider two variants; one that tunes the entropy coefficient (Figure 3a) and one that tunes the coefficients of both the entropy and the L 2 loss (Figure 3b).
The L 2 loss in our experiments is defined over the weights of both policy and value functions.
See Appendix A.5 for a longer description of the terms used in the inner loss.
The x-axis of each figure corresponds to the degree of richness of the context.
On the left, we start with no context, and then add one-by-one the statistics of following quantities as context features: value, reward, TD-error, action probabilities, cosine distance between last two gradients, past metaparameter predictions and state statistics.
For each of these quantities (except cosine distance and past meta-parameter predictions), the features are a history of means and standard deviations calculated over the last H rollouts.
For cosine distance and past meta-parameter predictions, the features are just a history of their last H values.
Since the number of features is large, we decrease the history length H from 10 to 4.
If we were to use H = 10 , the dimension of richest context input would be 660, which could make the optimization of meta-function too challenging (see Appendix A.6 for description and dimensions of each feature).
The only two context features that seem to harm the performance (Figure 3b) are value features and state features.
The drop seen at the end by adding state visitation statistics is likely due to overfitting.
State features have much higher dimensionality compared to other features (see Appendix A.4 and A.6 for more detailed description of state features), yet they are not informative of the agent performance and hence good meta-parameter values.
We do not observe as much overfitting for other features.
Generally speaking, richer features lead to better performance as there are more potentially useful sources of signal for the meta-parameter function.
6. H OW DO M ETA -G RADIENTS P ERFORM UNDER D IFFERENT R ATES OF N ON -S TATIONARITY ?
Lastly we look at how the degree of non-stationarity in the environment effects the performance of meta-gradients. We control the degree of non-stationarity by changing how often the tasks switch – shorter change periods correspond to higher rates of non-stationarity.
In all of the Figures included in this section, the medians and interquartile ranges are computed over 10 random seeds.
Two Colors.
In this experiment, we compare the relative improvement of different meta-gradient methods over the AC baseline for various non-stationarity rates (Figure 4).
The relative improvement is defined as percent increase or decrease in performance (as measured in total reward after 10M steps) over the mean performance of the baseline (with the same nonstationarity rate).
The absolute (non relative) values of all the methods can be found in the Appendix A.7.
We use BMG as an outer loop objective and tune entropy rate coefficient.
The compared methods are: BMG (no context), BMG-Reward and BMG-Rich.
As before, we report only the results for the best hyper-parameter configurations.
We find that when the environment changes too rapidly, at around 50,000 steps, the performance of meta-gradients deteriorates possibly due to their myopic nature.
We expected to find that improvement brought on by meta-gradients is greater when the environment changes more rapidly.
This expectation was met up to a point: BMG with a rich context were less affected in addition to outperforming contextless meta-gradients for any rate of change, indicating that providing more information can provide enough signal to enable learning even in this regime.
Switching MDPs.
Figure 5a and Figure 5b present the relative improvement of meta-gradients over the Q ( λ ) baseline in the Switching MDPs environment (for N = 4 and N = 1000 respectively).
Absolute (non relative) values can be found in the Appendix A.7.
We find that the use of context becomes more effective as we increase the rate of environment nonstationarity and the number of different tasks, where fast adaptation of meta-parameters becomes more relevant.
C ONCLUSIONS
We studied the performance and properties of white-box meta-gradients in non-stationary environments.
To study the effect of adding contextual information to the learned optimizer, we focused on formulations of meta-gradients where the learned meta-parameter values are functions of selected context features.
We found that adding more contexual information is almost always beneficial for lifetime performance.
Inspection of learned meta-parameter schedules and functions provides evidence of faster adaptation for meta-gradients with contextual information and convergence of meta-parameter functions over training.
When looking at the effect of increasing the rate of nonstationarity, we find that the meta-gradients without context, in contrast to meta-gradients with context, do not offer a large consistent advantage over fixed meta-parameter schedules.
An interesting avenue for future research are studies of contextual meta-gradients in non-stationary environments with less repeatability of contexts and continual supervised learning.
[Figure]
A A PPENDIX A. R ELATED W ORK
An alternative approach to white box methods for learning adaptive optimizers are the black box methods, which parametrize the entire update rule as a neural network, and learn it from scratch by training on a distribution of supervised (Andrychowicz et al., 2016) or reinforcement learning (Kirsch et al., 2020;
Oh et al., 2020; Kirsch et al., 2021) tasks.
In RL, most of these methods require training over multiple tasks or lifetimes, which makes them not directly applicable to our setting of interest.
The most relevant work is Xu et al.
(2020), which develops a black-box method that trains an optimizer over a single lifetime.
In black-box methods, the inputs to the learned optimizer functions are typically a history of parameters and gradients, but more similar to context features explored in our work, they can also include the entire rollout trajectories (Xu et al., 2020;
Oh et al., 2020).
However, these works do not explore the importance of selecting inputs for the learned optimizer and its effects on training in non-stationary environments.
Another related line of work is that of Contextual MDPs (Hallak et al., 2015; Modi et al., 2018;
Belogolovsky et al., 2021).
In this setting, contexts, which define the reward and transition kernel, are sampled from a distribution.
This formulation allows the study generalisation, e.g., we can learn a mapping from contexts to MDPs and ask how well should this mapping generalize to unseen contexts; see (Kirk et al., 2021) for a recent survey on the topic.
Our work, on the other hand, does not make such an assumption about the world and instead focuses on a single, but non-stationary MDP. It might be possible to think about some specific non-stationary MDPs as special cases of contextual MDPs, but we leave this exciting direction for future work. A. M ETA -G RADIENTS O BJECTIVES
Recall from Section 3 that the parameters θ are trained by minimizing an inner loss function
L ( inner ) η
( θ, D ) parameterized by the meta parameters η , whereas the meta-parameters η are trained to minimize the outer loss L ( outer ) ( η, D ) , where in both cases D refers to the rollout data used to estimate the loss.
The parameters and meta-parameters are generally optimized at two different time-scales; we denote by i the inner-loop and by k the outer-loop iteration step.
Most works in RL rely on entropy-regularized actor-critic algorithms (Xu et al., 2018;
Zahavy et al., 2020), where the inner loss objective L ( inner ) ( θ, D ) consists of the following three terms (policy, value and entropy loss respectively):
[Equation]
where α π , α v , α ent are the weights of individual losses, t are the environment step indices of transition tuples ( s t , r t , a t , s t +1 ) in D , and G t is a return estimate at timestep t .
The hyperparameters required for computation of each of the three terms (for example, discounts γ and bootstrapping parameters λ in n -step returns Sutton & Barto (2018)) or the weights of individual losses, can become tunable meta-parameters.
The outer loss objective generally consists of the same three terms, weighted by hyperparameters α ( outer ) π , α ( outer
[Equation]
v , α
[Equation]
outer ) ent :
[Equation]
where the data D used for estimating the outer loss, could be coming from a separate rollout or the same data as used in the inner loss.
A. B OOTSTRAPPED M ETA -G RADIENTS : Q ( λ ) AGENT
An alternative outer loss is proposed in Flennerhag et al.
(2022): the meta-parameters are trained to minimize a distance to a target which has been bootstrapped from the meta-learner.
This method is referred to as Bootstrapped Meta-Gradients (BMG).
In the best studied version of BMG, the metagradients are propagated through the last K updates, while minimizing the KL-divergence between the policy parametrized with θ K and a bootstrap target
π ˆ θ obtained by optimizing the policy for another L − 1 steps under the meta-learned update rule:
[Equation]
where the hyperparameter L is referred to as bootstrap target length.
The use of target which is
L − 1 steps ahead (without increasing the number of steps the meta-gradients are backpropagating through) reduces the myopia of standard MG objectives, while the use of KL divergence reduces ill-conditioning of outer loop objective.
The resulting adaptations of meta-parameters encourage reaching the target policy in a smaller number of inner optimization steps.
When tuning the exploration rate ϵ of Q ( λ ) agents, we use the following implementation from
Flennerhag et al.
(2022) to make the outer objective differentiable with respect to ϵ .
In equation 8, the stochastic bootstrap policy π θ K ( a | s ) and the target policy π ˆ θ ( a | s ) are defined as:
[Equation]
( s, a ′ ) 0 else. (9) , where q θ K ( s, a ) is the learned value function, the parameters ˆ θ are once again obtained by taking another L − 1 update steps, and | A | is the number of different actions.
The resulting objective does
[Figure]
not require differentiation through the update-rule, hence as in Flennerhag et al.
(2022), we use
K = 0 .
In this section, we provide additional information about the environments used in this paper.
The
Two Colors environment is visualised in Figure 6a and Switching MDPs in 6b.
A.4. T WO C OLORS
The dimension of the grid in Two Colors environment is 5 × 5 .
The observation space is constructed by concatenating one-hot encodings of x and y coordinates of positions of agents and two other objects.
The total dimension of resulting observation space is 3 × 2 × 5 = 30 .
A.4. S WITCHING MDP S
The reward function for each MDP is generated in the following manner: for each state-action pair there is a 50% chance of zero reward, 20% chance of reward of +1, 20% chance of reward -1, and 10% chance of a random reward sampled uniformly from the interval [-1, 1].
The transition function for each MDP is a standard two dimensional grid world—states are characterized by an ( x, y ) coordinate, and there are four actions that move the agent up, left, right, and down unless the intended cell is the edge of the grid or a wall is present.
The transition function changes across the
N MDPs by the addition of a set of walls randomly placed throughout the grid.
Up to 15 walls are placed per-MDP, in sequence, by randomly sampling unoccupied cells (by goal, agent, or another wall).
At the beginning of an experiment, we generate a set of N different MDPs.
At regular fixed intervals, the next grid world is sampled from this set uniformly with replacement.
The set and the order of samples is uniquely determined by experiment random seed.
If N is large compared to the total number of task switches experienced during a lifetime (e.g. N = 1000 with a change period 100 , 000 and a lifetime of 20 M steps), the probability of agent experiencing the same task multiple times is small.
Note that since we measure the lifetime performance after 20 M environment steps in our experiments, N does not correspond to the number of different MDPs experienced during this lifetime.
A. E XPERIMENTAL S ETUP AND H YPER - PARAMETERS
In this section, we provide further details on the experimental setup and hyper-parameters of agents and meta-learners used in Section 6.
The hyper-parameters used in the experiments with AC agents are described in Table 3.
Our implementation of the AC agent implements the softmax policy and the value function as two separate [Table] [Table] [Table]
[Table]
[Table] feed-forward MLPs. The parameters are updated every 16 environment steps: given fixed parameters, the agent interacts with the environment for 16 steps collecting observations, rewards and [Table]
[Table]
actions into a rollout D , which is then used to compute the inner loss.
The inner loss consists of the following four terms (policy, value and entropy and L 2 loss respectively):
[Equation]
[Equation] In most of the experiments [Table] the only meta-parameter is α ent , and α L 2 = 0 (i.e. there is no Lregularization).
However, in Section 6.3 (Figure 3b), we tune both α ent and α L 2 .
There we used two separate MLPs to predict each of the two meta-parameters.
The output of the meta-parameter network predicting α L 2 has been scaled by a fixed quantity ( 10 − 4 ) to prevent instabilities caused by too strong forgetting.
The weights of both policy and value networks were included in L L2 ( θ ) .
The hyper-parameters used in the experiments with Q ( λ ) agents are described in Table 6.
The qfunction is again a feed-forward MLP.
The agent is optimized at each step, i.e. without batching.
To avoid instabilities this could cause, we use a momentum term that maintains an exponentially moving average over gradients, with the discount factor 0 .
We sweep over the learning rate of the inner learner only for the fixed meta-parameter baseline, for the meta-gradient methods, the inner learner’s learning rate is set to 3 · 10 − 5 .
To ensure stable initial predictions, the meta-networks were pre-trained on random context inputs sampled from uniform distribution [ − 1 , 1] , to predict an output in the middle of the possible metaparameter range ( 0 .
5 for α ent and ϵ , 5 · 10 − 5 for α L 2 ).
Note that when using contextual metagradients, we do not sweep over the hyper-parameters introduced by the addition of context (e.g. the size of meta-network and context history H ).
A. C ONTEXT F EATURES
In this section, we describe in more details the construction of context features.
At each update step i , the meta-parameters are computed by feeding context features
c i into a meta-parameter function.
The context features are a concatenation of one or several different types of features (e.g. reward, value, TD-error), the type of features used in each experiments is described in the corresponding experiment’s section.
For example, when the context is specified as Reward , the input to the metaparameter function is c i =
[ c ( reward ) i ] , and when the context is Rich , the input to the meta-parameter function is c i =
[ c ( reward ) i , c ( value ) [Equation]
, c ( td − error )
i ] .
For each different feature type (here denoted with ”feat”), the features are a history of size H :
[Equation]
where ˆ f ( feat ) i is a quantity computed using statistics at i -th update step.
Each of these quantities has been normalized (we used the statistics of context features measured during the first 2M steps of training with the AC baseline agent, but more generally, one could use a running average as an estimate) and passed through a tanh ( · ) function to ensure all features are in the [-1, 1] range.
We will refer with f ( feat ) i to quantities computed before this transformation.
For example, in experiments described with AC-BMG-Reward, each f ( reward ) i is a mean of rewards r t from the rollout D i which was used to compute i -th update.
In experiments described with Q ( λ ) -BMG-Reward, because i -th update for Q ( λ ) agents is calculated using data from only one environment step, f ( reward ) i is the reward r t from only that step.
The summary of all different feature types used in experiments with AC agents, including how f ( feat ) i is defined for that feature type and dimensions of that feature, can be be found in Table 5.
Note though that when the feature is specified as Reward (AC-MG-Reward and AC-BMG-Reward), we only used mean of rewards, and the corresponding feature dimension is H .
The summary of all different feature types used in experiments with Q ( λ ) agents can be found in Table 6.
A. D IFFERENT R ATES OF N ON -S TATIONARITY : T OTAL R EWARDS
Lastly, we supplement the results in Section 6.4 by illustrating how the performance of all methods drops as the rate of environment non-stationarity increases.
In Two Colors experiments (Figure 7), no learning occurs when the environment switches every 10k steps.
In Switching MDPs, note that the performance drop is much more significant when the number of different MDPs is large (Figure 8b), indicating that in this environment, the agents benefit from repeated exposure to the same MDP.
[Figure]