D ISTRIBUTIONAL M ETA -G RADIENT R EINFORCEMENT L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Meta-gradient reinforcement learning (RL) algorithms have substantially boosted the performance of RL agents by learning an adaptive return.
All the existing algorithms adhere to the same reward learning principle, where the adaptive return is simply formulated in the form of expected cumulative rewards, upon which the policy and critic update rules are specified under well-adopted distance metrics.
In this paper, we present a novel algorithm that builds on the success of metagradient RL algorithms and effectively improves such algorithms by following a simple recipe, i.e., going beyond the expected return to formulate and learn the return in a more expressive form, value distributions.
To this end, we first formulate a distributional return that could effectively capture bootstrapping and discounting behaviors over distributions, to form an informative distributional return target in value update.
Then we derive an efficient meta update rule to learn the adaptive distributional return with meta-gradients.
For empirical evaluation, we first present an illustrative example on a toy two-color grid-world domain, which validates the benefit of learning distributional return over expectation ; then we conduct extensive comparisons on a large-scale RL benchmark Atari 2600, where we confirm that our proposed method with distributional return works seamlessly well with the actor-critic framework and leads to state-of-the-art median human normalized score among meta-gradient RL literature.
I NTRODUCTION
Meta-gradient reinforcement learning (MGRL)
(Xu et al., 2018b; Sutton, 2022) is a family of algorithms that leverage the gradient of the gradient descent update to learn better objectives for reinforcement learning (RL).
The MGRL paradigm has achieved substantial performance breakthroughs and dominated the state-of-the-art model-free RL algorithms in various domains, such as
Atari (Xu et al., 2018b; Zahavy et al., 2020; Flennerhag et al., 2022), DMLab (Zahavy et al., 2020), and DeepMind Control (Zahavy et al., 2020).
The major algorithmic research on MGRL has been developing towards more and more ambitious targets.
For example, it starts from learning two fundamental hyperparameters (i.e., discount factor γ and bootstrapping factor λ )
(Xu et al., 2018b), and then goes to self-tuning 20+ hyperparameters in a high-performing RL agent (Zahavy et al., 2020).
As a powerful tool for discovering RL semantics from the data, MGRL has been utilized to discover intrinsic rewards (Zheng et al., 2018; 2020), auxiliary tasks (Veeriah et al., 2019), options (Veeriah et al., 2021), etc.
Recently, it has even been used to tackle the more flexible form of parameterizing the objective, i.e. using black-box neural networks to learn the objectives from the environment interactions and learning context (Xu et al., 2020;
Oh et al., 2020).
However, MGRL research has been limited to scalar form, where either metaparameters are scalar ones (Xu et al., 2018b; Zahavy et al., 2020), or outputs of the meta networks are scalars (Xu et al., 2020).
This is in line with the traditional definition of returns (Sutton & Barto, 1998) as an expected sum of the discounted rewards.
However, note that MGRL is not limited to providing RL objectives, while it is possible to meta-learn a broader concept, e.g, learning rate (Baik et al., 2020; Pinion et al., 2021) and exploration (Gupta et al., 2018;
Xu et al., 2018a).
In this work, we investigate an essential step of extending the algorithmic scope of MGRL to distributional methods, where RL algorithms take away the expectations in Bellman updates and consider a richer formulation of value distributions (Bellemare et al., 2017).
This is essential for RL updates and even more critical for MGRL which enjoy multiple satisfying properties.
Sparse signals are hard to learn in RL, and bi-level optimization nature of the MGRL makes it extremely difficult to propagate helpful signals to learn meta-parameters.
By modeling value distributions, the RL algorithm can naturally produce a rich set of predictions, which provides dense and informative signals for meta optimization.
Moreover, some distributional RL methods come with new optimization properties in the loss form (not available in conventional non-distributional RL methods), e.g., with KL divergence between discrete distributions (Bellemare et al., 2017) or quantile regression (Dabney et al., 2018b).
Despite the great accomplishment of distributional RL for non-meta policy learning, the attempt to learn a distributional meta objective has not yet been made by any prior MGRL methods.
This paper aims to develop a novel distributional meta-gradient RL algorithm that can discover meaningful value distributions online.
Furthermore, we aim to both methodologically establish a novel distributional meta-gradient framework and find out the intriguing effect of applying distributional meta-gradient techniques on real-world problems, for which we conduct extensive evaluations with both motivating toy domains and large-scale end-to-end policy training problems.
The evaluation results well demonstrate the effectiveness of our method.
Overall, the contributions of our paper are as follows:
(i)
We present a novel meta-gradient variant of the RL method, which considers approximating the distributional return in its value update.
(ii)
Our proposed distributional meta-gradient algorithm is general and can be compatible with almost all the existing meta-gradient RL approaches.
(iii)
We conduct an extensive empirical study on a toy control domain and the large-scale benchmark Atari 2600.
Our results have demonstrated substantial improvements by our method over strong baseline methods on the major evaluation domains.
R ELATED W ORK
In recent decades, meta-learning, or learning-to-learn , has been explored extensively within the machine learning community (Hospedales et al., 2022; Sutton, 2022), e.g., learning a neural network update rule (Bengio et al., 1990), adapting the learning rate (Sutton, 1992; Schraudolph &
Giannakopoulos, 1999), and transferring domain-invariant knowledge (Thrun & Mitchell, 1995).
Notably, it has also driven many advances in the RL problems.
One direction for such attempts is to learn meta -optimizers (Andrychowicz et al., 2016;
Wichrowska et al., 2017) or meta - policies (Duan et al., 2016; Wang et al., 2017) parameterized by recurrent neural networks, which capture the highlevel time-dependent information as meta knowledge.
Another direction is gradient-based metalearning, introduced in Model-Agnostic Meta-Learning (MAML)
(Finn et al., 2017a), as well as its followup works (Finn & Levine, 2018; Finn et al., 2017b; Grant et al., 2018; Al-Shedivat et al., 2018).
Those works focus on learning a good initialization of the model for one or few-shot multitask learning with meta-gradients.
In this paper, we consider solving meta-gradient RL problem from a fresh new angle, i.e., optimizing the adaptive return as distributions.
Policy learning with distributional return has been considered in existing RL literature for off-policy learning only.
All the existing methods fall into the strand of DQN variants which replace the n -step return in DQN with distributions during Q-learning.
In (Bellemare et al., 2017), the return is formulated as a categorical distribution with fixed values for the base, where Bellman operator shifts the base values of the distribution and the target distribution is contracted and projected upon the fixed base.
In (Dabney et al., 2018b), the return is formulated as quantile distributions, where the algorithm tries to learn the values corresponding to each quantile, alleviating the need for distribution projection or value range approximation.
In (Dabney et al., 2018a), the quantiles are approximated by mapping from values drawn from the random distribution (proposals) into quantiles, resulting in an unlimited number of implicit distributions.
In (Wang et al., 2019), the proposals in implicit quantiles are replaced by the output of a learned proposal network.
In contrast to the aforementioned off-policy distributional RL methods, we tackle a novel problem of learning adaptive distributional return in on-policy actor-critic algorithms with meta-gradient.
Compared to existing meta-gradient RL methods with scalar return , our method applies a distributional return, which is advantageous as (1) it can offer a more informative target to approximate, (2) it naturally provides a richer set of hyperparameters to be tuned with white-box meta-learning, which offers more feedback in the meta-gradient, and (3) approximating such a distributional return naturally results in meaningful auxiliary tasks to facilitate more efficient policy and value training.
P RELIMINARY
Quantile Distributional Value Update
In this work, we consider learning the return as a quantile distribution.
One reason for us to choose quantile distribution rather than any other type of distribution is, the Bellman operator for quantile distribution does not need to perform projection (lose precision) or manually define the return boundaries, like C51 (Bellemare et al., 2017) does.
Also, our proposed method employs a deterministic prediction function for quantiles which could be conveniently trained, unlike the stochastic mappings employed by IQN (Dabney et al., 2018a), where the stochasticity would bring noise to the meta gradient and confuse the training of adaptive return .
We denote the distributional return function over each state s as Z θ ( s ) , where θ is its parameter set.
With quantile, the distributional return can be approximated by a uniform mixture of N Dirac functions, i.e., Z θ ( s ) :
= 1 N (cid:80) Nq =1 δ θ q ( s ) , where θ q is the parameter for the q -th Dirac function which can be trained by quantile value update.
For a given quantile distribution Z and a quantile τ , the value of the quantile function is denoted as F − 1 θ ( τ ) .
The quantile distribution can be trained by quantile regression loss L qrτ ( · ) which penalizes overestimation errors with weight τ and underestimation errors with weight 1 − τ , where τ ∈ [0 , 1] :
[Equation]
where
[Equation]
To smooth the quantile loss at zero, we could apply Huber loss over it, i.e.
[Equation]
Thus the Quantile Huber Loss can be derived as follows,
[Equation]
White-box Meta-Gradient RL White-box meta-gradient RL methods use a bi-level optimization scheme to update important high-level parameters to characterize the return and the composed loss function by meta gradient.
We denote the high-level parameter set as η .
In the prior work from (Xu et al., 2018b), η = { γ, λ } ; later, loss weights and importance sampling threshold have been added to η .
We denote the RL loss function as J ( τ, θ, η ) .
During bi-level optimization, the meta-gradient methods first update the model parameter θ with the gradient f ( · ) w.r.t θ :
[Equation]
where ζ is the learning rate and the first-order gradient function f ( · ) is differentiable w.r.t η .
Then, an alternative batch of transitions τ ′ is sampled to optimize the model in the second update stage, where outer objective J ′ computed from the updated model parameters θ ′ and a reference hyperparameters set η ′ is optimized (e.g., default values for η ′ )
w.r.t η :
[Equation]
Since θ ′ is differentiable w.r.t η , the meta gradient is computed as the gradient of the outer objective
J ′ ( τ ′ , θ ′ , η ′ )
w.r.t η .
Thus meta-gradient RL enables the effective high-level parameters in learningto-learn to compute the adaptive return and/or the importance of loss function components.
D ISTRIBUTIONAL M ETA -G RADIENT RL A LGORITHM
Actor-Critic with Distributional λ -Return
The principle of RL is to optimize the agent’s cumulative rewards, also known as the return (denoted as G ), from the sequence of experiences collected from the environment, where the experience at time t id denoted as a tuple τ t
= ( S t , A t , R t , S t +1 ) with S t , A t and R t denoting the state, action, and reward for the state-action pair at step t , respectively.
In practice, the return can be formulated in various forms.
One well-adopted form of return is via bootstrapping .
Let T denote the terminal step and γ the discount factor.
An example n -step return (Sutton & Barto, 1998) is defined as:
[Equation]
where the return is the sum of n step rewards and the bootstrapped value for S t + n .
To better leverage the return G t derived over multiple steps, there is also a λ - return , or TD( λ ) return , which measures the return as the weighted TD error over all time steps (Sutton, 1988), where there is another highlevel parameter λ introduced to characterize the λ - return function.
When λ = 0 , it terminates the bootstrapping from further steps, and when λ = 1 , it promotes pure bootstrapping.
White-box metagradient RL methods attempt to meta-learn the crucial hyperparameters characterizing the λ - return with a bi-level optimization regime.
We consider the training of actor-critic agents with the critic approximating a quantile distributional return predicted by a value function Z θ ( S ) , where S is an input state, θ is the parameter set of value function, and Z θ ( S ) follows the distributional Bellman optimality operator (Bellemare et al., 2017),
[Equation]
where where U D :
= V denotes equality of probability laws, i.e., the random variable U is distributed according to the same law as V. Our proposed distributional adaptive return can naturally result in a larger set of self-tuning hyperparameters.
The critic is updated towards a canonical distributional λ - return denoted as Γ λt following a distributional TD( λ ) algorithm,
[Equation]
where Γ λt is essentially the sum of the predicted return distribution Z ( S t ) and the backpropagated
TD error over all the future steps δ t +1: ∞ , where the later could be computed iteratively and weighted by the product of discounting factor and trace in their exponential forms.
When training the actorcritic with Γ λt under a distributed regime like IMPALA (Espeholt et al., 2018) for efficient learning, the training data would become off-policy and we need to perform off-policy correction to get the unbiased gradient.
Hence distributional λ - return with off-policy correction in the form of v-trace could be defined iteratively as follows,
[Equation]
[Equation]
where c m is the truncated importance sampling (IS) ratio with ¯ c being the clipping threshold, and π ( · ) , µ ( · ) are the present and previous behavior policies, respectively.
To update the distributional actor-critic with off-policy correction, we optimize the following objective function:
[Equation]
[Equation]
where A Z denotes the advantage computed as the sum over quantile error of the distributional return , Γ λ ( q ) ( S ) denotes the q -th quantile in the distributional λ - return , L Z / L π / L H and w Z / w π / w H are functions and weights for the value/policy/entropy losses respectively, sg ( · ) is the stop-gradient function and the first term in L π corresponds to the advantage.
Distributional Meta-Gradient Update
Our proposed distributional Meta-Gradient RL algorithm adopts a white-box bi-level optimization scheme to learn adaptive distributional return Γ ηt .
We denote the meta parameters characterizing the adaptive return Γ λt and important weights in the actorcritic loss function L θ as a set η , where η = { λ, γ, α, w H , w π , w Z } , where α is the weight for the leak-vtrace in STACX (Zahavy et al., 2020) to control the off-policy correction ratio.
Formally, the adaptive return can be formulated as follows,
[Equation]
In other words, we can consider a quantile-specific discounting and bootstrapping operator over the distributions and vectorize the λ and γ in the set η .
Furthermore, we can also specify a vectorial quantile loss weight w Z .
Thus, our method results in a much larger hyperparameter set compared with STACX , i.e., | η | for DrMG , STACX and MG are 3 × n Z + 3 , 7 × n blk , and 2, where n Z and n blk represents the dimension of quantiles and the number of STACX ’s multi-head blocks.
At the inner loop, the inner loss L in θ ( τ, θ, η ) is first computed to update θ towards θ ′ using the gradient ∂ L in θ ( τ, θ, η ) /∂θ , where τ is a batch of data sampled to update the inner loss and η is drawn from the meta parameters.
[Equation]
where the gradient of the distributional value loss is derived as follows,
[Equation]
At the outer loop, we evaluate the high-level parameters η by sampling an alternative batch of transitions τ ′ and computing the outer loss L out η ( τ ′ , θ ′ , η ′ ) , where η ′ corresponds to some reference hyperparameters for η , e.g., the default values predefined like those in other algorithms.
The meta gradient to update the meta parameter η is defined as follows,
[Equation]
where ζ ′ is the learning rate for the outer loss.
The meta gradient evaluated at the outer loss interacts with meta parameters through ∂θ ′ ∂η , which could be further expanded as ∂ 2 L in θ ( τ,θ,η ) ∂θ ∂η based on the chain rule.
With the definition of the distributional return, meta-learning variants become learning the bootstrapping or termination condition over the distributions, rather than scalars, which would bring informative information for return approximation with a greater degree of adaptability.
E XPERIMENTS
We denote our proposed Distributional return Meta-Gradient algorithm as DrMG .
5. T WO
[Equation] G RID -
[Figure] Settings We consider a 7 × 7 grid-world, where the green-colored agent is assigned to collect the key to open the door at the yellowcolored goal.
There are two types of keys; collecting one of them would get a reward of 0.
and the other 0.3.
The agent receives a reward +1 upon successful navigation.
After an episode terminates, we randomly generate locations for both keys and the goal.
We toggle the reward at every 1e5 steps to investigate the agent’s adaptation ability.
Thus a key-agnostic optimal agent could collect +1 .
5 − r pen ∗ T expected return, where r pen and T are the step-wise penalty and episode horizon.
We assume that the agent approximating the distributional return should be able to encode the dynamics in its return so that it can adapt better compared to the baseline agent approximating the scalar return .
Results We compare our proposed DrMG agent with a meta-gradient baseline algorithm
MG as well as a vanilla distributional algorithm QR , where QR is an IMPALA variant with distributional return.
The results are shown in Figure 1.
We configure DrMG , MG , and QR agents by almost identical hyperparameter sets.
Since the RL agents are unaware of the reward toggling during policy training, their adaptation ability is characterized by their efficiency in modeling/approximating the return during the training process.
We notice that after the reward toggle, DrMG achieves peaks with substantially higher return compared to that for MG and QR .
Furthermore, QR slightly outperforms MG , which shows that approximating distributional return benefits the task adaptation.
5. A TARI 2600 200M
Settings We testify our method on Atari 2600 Benchmark (Bellemare et al., 2013) under the 200M setting, where we try best to align the setting with prior works (Xu et al., 2018b; Zahavy et al., 2020).
Our agent is developed upon IMPALA (Espeholt et al., 2018), an efficient distributed actor-critic framework.
The actor-critic module employs a deep ResNet (He et al., 2016) as backbone, which is a feed-forward model without LSTM.
The network takes 4 grayscale images with size 84 × 84 as input.
The backbone is connected to a policy-value head where the value head outputs distributional quantiles with a dimension of N and the policy head outputs policy logits with a dimension of |A| .
The meta module is a joint set of scalars and vectors, with each vector entry representing a distributional hyperparameter in η .
For training, we strictly follow the 200M regime without extra data or experience reuse.
During training, we sample two batches with identical batch size and use them to form two combinations of inner data and outer data to compute the meta-gradient update loop twice.
Each task runs with identical hyperparameters, and device configurations for hardware and software.
We highlight that
DrMG brings very little computational overhead; namely, DrMG only brings | N − 1 | ∗ 256 more parameters compared to MG and it is more computationally efficient than STACX, since STACX introduces multiple policy-value heads as its auxiliary task components.
All the games employ the same wrapper, for which the action repeat is 4, stick action probability is 0, episodic life is false , and maximum episode length is 108,000.
All the tasks adopt the ‘no-ops starts’ protocol, where at the start of each episode, the agent randomly samples a period to take a dummy action ‘0’ up tosteps.
We present the environment wrapper setting in Sec.
A.1.2, and the detailed model/experiment configurations in Sec.
A.2.
For comparison, we evaluate the performance on each game in terms of human normalized score (HNS) defined as HNS
= R ours − R rand R human − R rand , where R ours/ R rand/ R human are the episode rewards for our/random/human expert policies.
Results For evaluation over the 57-task benchmark, we compute the median HNS for each method.
We consider the following five most related baselines: C51 , QR-DQN , Implicit Quantile
Network ( IQN ), Self-Tuning Actor-Critic with auXiliary tasks ( STACX ) and MG .
The first three are off-policy value-based methods which are DQNvariants, while STACX and MG are white-box meta-gradient RL baselines that are IMPALAvariants, the same as DrMG .
There is another recently published meta-grad RL work Bootstrapped
Meta-Gradient ( BMG ), which is excluded here for comparison because the work employs a different
Atari 2600 benchmark with altered task and model
[Figure]
settings, motivated to solve orthogonal problems than our work.
The detailed explanation is presented in Appendix A.1.2.
We show the learning curves for median HNS in Figure 2.
We do not present the curve for STACX , because we find the reported standard for STACX is rather hard to be reproduced.
When carefully reimplementing it with an identical framework, hyperparameters and model architectures, the gap between the reproduced and reported median HNS is > 50% .
Nevertheless, the learning curves for STACX on representative games can be found in Figure 3.
Such a reproduction issue is notorious for Atari 2600 200M evaluation due to the missing official releases of the baselines.
From the median HNS curves shown in Figure 2, we can first observe that DrMG performs substantially better than the conventional distributional RL methods C51 , QR-DQN and IQN , which demonstrates updating the distributional return adaptively with meta-gradient and the actor-critic framework is more effective than updating the distributional returns by off-policy DQN variants.
We also find that DrMG brings noticeable performance improvement over the strong white-box meta-gradient RL method MG .
The median HNS for DrMG is 379 % and that reported in MG is 287 % .
These results show that optimizing the distributional return in our proposed λ -distributional form is more efficient than learning the scalar return .
We present the learning curves of our method, MG and STACX on a set of representative Atari games in Figure 3.
For computational efficiency, the training adopts an early stopping policy, where the training terminates after the agent surpasses a fixed HNS standard (450 % in our evaluation) for ¿2M steps.
We classify the games into three strands { fast , medium , slow } , based on the behavior of the RL agents to hit the designated HNS standard, where fast =
{ Breakout , DemonAttack } , medium = { SpaceInvaders , Robotank , Gopher , Kangaroo } and slow = { BeamRider , Zaxxon , Centipede } .
Games in fast , medium and slow groups consume < 10% , 30% − 60% and ∼ 200M frames
[Figure]
to hit the HNS, respectively.
We examine the fast group to show that learning the distributional return in DrMG would not slow down the learning process compared to its counterparts learning scalar return .
On the classic Atari game Breakout , DrMG converges noticeably faster than MG and
STACX , and on DemonAttack , DrMG learns as fast as MG .
It is also desirable to evaluate learning efficiency for the games from the medium group since different methods behave very differently in making progress in those tasks.
From the results, we can see DrMG progresses significantly faster than the other two methods in Robotank and Kangaroo .
In Kangaroo , our method saves almost 1/ and 1/2 frames compared to MG and STACX .
In SpaceInvaders , the progress of DrMG is comparable to the others.
In the slow group, the RL algorithms strive to progress with slow learning paces.
All the three games are challenging ones for our method as well as the other baselines to achieve high HNS.
So, the performance on them is crucial to derive the median HNS score, i.e., pulling back any one could potentially lower the median HNS.
We can observe that DrMG leads to fairly stable performance among those in slow group.
In BeamRider , Centipede and Zaxxon , DrMG achieve better asymptotic rewards than the baselines.
It is also shown in Figure 3 that the performance of
STACX is not ideal, e.g., irregular performance in DemonAttack and Kangaroo .
5. A BLATION S TUDY
We perform ablation studies to investigate the effects of the following designs in our method: (i) the adaptive distributional return vs. non-adaptive distributional return ; (ii) λ -distributional return vs. n -step distributional return; (iii) vector hyperparameter η vs. scalar
η ; (iv) adopting BMG’s bootstrapped inner steps on top of our method.
Due to space limit, we present the results of ablation study (iv) in appendix A.3.
Adaptive vs. non-adaptive distributional return We investigate the effects of adaptive/nonadaptive distributional return by comparing DrMG with two ablated baselines.
The meta parameter set is split into two groups η = { η 0 , η 1 } , where η 0 = { λ, γ, α } and η 1 = { w H , w π , w Z } .
The first baseline only meta-update η 1 , denoted as DrMG( η 1 ) (complete version is denoted as
[Figure] [Figure]
DrMG( η 0 , η 1 ) ).
The second baseline is QR , a non-meta method equivalent to DrMG ( {} ).
We compare them in Breakout , Krull , and SpaceInvaders .
The results are shown in Figure 4.
We first notice that both versions of non-adaptive distributional λ - return work considerably well with the actor-critic IMPALA framework (outperform their corresponding DQN variants by large margins).
Using QR can learn as smoothly as DrMG in Krull but for the other two domains, it underperforms
DrMG , which shows the effectiveness of using meta-gradient to optimize distributional return .
DrMG( η 1 ) outperforms QR in 2/3 games, but it underperforms DrMG ( η 0 , η 1 ), which demonstrates the superiority of meta-learning adaptive return in distributional meta-gradient algorithms.
λ -distributional return vs. n -step distributional return We investigate the effects of λ -/ nstep distributional return by comparing DrMG with a meta-gradient RL baseline that replaces the λ -distributional return by an n -step return computed over the return distributions, denoted as
DrMG+n-step .
We consider two values for n in evaluation, where n = 3 and n = 5 , i.e., resulting in baselines DrMG+3-step and DrMG+5-step .
We adopt three domains, DemonAttack , SpaceInvaders , and Centipede .
The results are shown in Figure 5.
We can observe that replacing λ -return with n -step return substantially pulls down the rewards even in the fast game Breakout .
This further confirms that backpropagating the distributional TD errors from λ -return is highly efficient to formulate the adaptive distributional return target
. Vector hyperparameters η vs. scalar hyperparameters ¯ η
We then investigate the effects of vector hyperparameters η vs. scalar hyperparameters ¯ η by training DrMG with two kinds of adaptive distributional return parameterizations.
When using scalar hyperparameters ¯ η , the return for each quantile defined in Γ λ shares the same discounting factor γ , trace λ , as well as cost for the quantile regression loss.
When using a vector form of η , each quantile head enjoys a quantile-specific parameterization for the adaptive return .
We evaluate DrMG ( η ) and DrMG ( ¯ η ) in Breakout , SpaceInvaders and Centipede , where the baselines are most distinguishable with DrMG ( ¯ η ).
The results are presented in Figure 6.
We observe that DrMG ( ˆ η ) performs fairly well in Breakout as well as
SpaceInvaders , where other ablated baselines easily fall in the latter case.
However, for the slow game Centipede , DrMG ( ¯ η ) noticeably underperforms its vector version DrMG ( ¯ η ).
C ONCLUSION
We present a meta-gradient algorithm with distributional return named DrMG , which employs a critic that learns adaptive distributional return with parameters to be trained by meta-gradient derived from a white-box bi-level optimization scheme.
We demonstrate that the quantile distributional return can seamlessly work with distributed actor-critic algorithms and learning such a distributional adaptive distributional return can offer more meaningful signals to update the meta parameters.
We show DrMG is able to outperform its most related non-meta-learning baseline IMPALA and several decent RL baselines on a motivating Two-Colors Grid-World domain and the challenging video game benchmark Atari 2600.
Our method is also very general and the distributional return we introduce hereby can be adapted to work with alternative meta-gradient RL methods, with the quantile distributional return being able to be flexibly replaced with other forms of distributions.
We leave the work of integrating more variants of actor-critic algorithms or distributional returns in DrMG as our future work.
It is also worth considering applying DrMG on more practical RL domains, such as pure value-based RL
E THICS AND R EPRODUCIBILITY S TATEMENT
Our work tackles synthetic digital game playing applications and therefore we believe our work is free from ethical problems.
We present details on the distributional framework, Atari wrapper settings, and model hyperparameters in Appendix.
We also include the code for the two important algorithmic components, networks and two-level meta-update step function in our work.
A.1. T WO - COLORS G RID -W ORLD
The two-colors grid-world is a 7 × 7 maze navigation problem.
The agent is tasked to first collect a key and then navigate to the goal location with the key to open the door with the key.
The agent is represented by green color, the two keys are represented by red and blue color, and the goal location is represented by light yellow color in Figure 1.
Reward there are three types of rewards issued to the agent: (ii) key collection rewards for both keys where one is 0.7 and the other is 0.3; (ii) goal reaching result, which is +1 when the agent has collected a key and -1 when the agent come to the goal prior to key collection; (iii) step-wise life penalty 0.01 to promote shortest navigation paths.
Reward toggling the key collection reward values toggle between the two keys after every 1 eenvironment steps, so that the agent needs to learn to identify keys for maximizing the reward, rather than converging to a trivial policy where the agent just learn to collect one fixed type of key without considering another, for learning sub-optimal policy.
Terminal condition the agent will enter a terminal state under the following four scenarios: (i) arrive goal with a key; (ii) arrive goal without collecting a key; (iii) have collected one key and attempt to collect the other key; (iv) have collected one key but attempt to collect the same key again.
State an eight-dimensional vector [ agent-y - goal-y, agent-x - goal-x, agent-y - redkey-y, agent-xredkey-x, agent-y - bluekey-y, agent-x - bluekey-x, blue-collected, red-collected ].
The agent needs to be unaware of whether when or how the key rewards are toggled or specified.
Action set consist of four moving directions: { left, right, up, down } .
When the agent attempts to transit to out-of-boundary locations in the maze, the game does not terminate but keep the agent unmoved.
Analysis the optimal policy is the one which could always identify the correct key rewards association and thus always the collect the higher rewarded key for opening the door.
Thus the optimal episodic return is 1 + 1 = 2 .
A key-agnostic agent is the one that randomly collects one of the key or just stick to one key for collection before reaching the goal location.
Thus the expected episodic return for such key-agnostic agent is 1 + (0 . 7 + 0 . 3) / 2 = 1 . 5 .
The aim of this experiment how fast/well the RL agent behaves within each reward toggling interval.
A.1. A TARI
Wrapper We describe how we implement the wrapper for Atari environment.
We create a base environment from gym.wrappers.atari preprocessing.
AtariPreprocessing , then embed it into a selfimplemented wrapper class to perform the standardized operations like frame skip and no-op start .
Upon reset, the wrapper first selects a random number between 0 and 30 to determine the number of no-op steps.
Then it performs no-op reset by repeatedly taking action 0.
In the step ( · ) function, we receive an action determined by the RL agent and then perform frame skipping and frame stacking to generate observations.
The agent first performs frame skipping by repeat the action by 4 times.
If the episode terminates within the repeats, the agent jumps out of the loop immediately.
After frame skipping , the agent checks life loss and return next observation, cumulative reward within the repeated 4 steps, episode terminal flag and life loss.
Note that an important issue here is no frame pooling is performed.
We found that with frame pooling, the learning curve for the agent hardly progress to above reported MG’s standards in quite a few slow learning games like Centipede ,
BeamRider and Alien .
We use the same wrapper for the meta-gradient RL baselines.
Discussion on why BMG is excluded from our MGRL baseline We hereby explain why BMG is not considered as our comparison baseline from empirical and methodological perspectives.
First, from the empirical perspective, we wish to highlight that the effort of adjusting the setting of
BMG to add it as a baseline for a fair comparison with our work is non-trivial, mainly due to the following factors:
BMG uses a different ALE wrapper, which makes the task different from the one adopted by our work.
The main change is changing the grayscale input (84×84) into RGB input (160×210×3).
BMG uses a much larger ResNet, where the channel depth changed from (16, 32, 32) to (64,128,128,64).
BMG uses replay, wherein in each batch, the replay to online ratio is 2:1, whereas MG/STACX/DRMG do not employ replay.
BMG with replay is also trained under a smaller batch size than MG/STACX/DRMG, where the batch sizes for the former and latter algorithms are 18 and 32, respectively.
The last point is somewhat more troublesome than the first two because it requires heavy engineering effort to fully integrate a reply buffer into the code base of MG/STACX/DrMG when BMG does not have open-source code.
Meanwhile, running the BMG experiment with experience replay and a small batch size would significantly increase the run time for the method.
It also remains unclear whether an IMPALA variant with replay could be considered a fair baseline to compare with the other meta-gradient methods in our work (MG/STACX/DrMG) that do not perform data reuse.
Second, from the methodological point of view, BMG and our work tackle orthogonal directions in meta-gradient RL problem, i.e., BMG is innovated on the solver level, focusing on how to better optimize the bi-level optimization with bootstrapping, whereas our work focuses on modeling the adaptive distributional return which is solved by a standard meta-gradient solver (with two-step inner-outer update).
Nevertheless, we adapt the main bootstrapping techniques of BMG into our work and present comparison results between the bootstrapped DrMG baseline and our method in
Appendix A.3.
A. C ONFIGURATIONS
For distributional computation, we adopt ray for communication.
We adopt batch environment wrapper from an open-source distributed reinforcement learning framework moolib to accelerate actor computation.
The detailed configuration of the distributional framework is as follows: [Table] moolib Num. [Table]
for Evaluator Num.
GPUs for Evaluator 0.
Num.
[Table]
for Actor Num.
GPUs for Actor 0.
Num.
CPU for Learner Num.
CPU for Learner 1.
Num. train itr to push params Num.
rollout step to pull params
[Table]
The hyperparameter configurations for our method is as follows:
[Table]
The versions for sensitive packages are as follows:
Package
Version jax 0.2.
rlax 0.1.
optax 0.1.
gym 0.21.
moolib 0.0.
opencv-python 4.6.0.
python 3.8.
A. A DDITIONAL A BLATION S TUDY R ESULTS
Adapting Bootstrapped Meta Learning (BMG) into DrMG We additionally investigate effects of combining the bootstrapping techniques introduced in the state-of-the-art meta-gradient work
BMG (Flennerhag et al., 2022) with the vanilla version of our method DrMG without bootstrapping.
We denote our method with bootstrapping as DrMG+Bootp .
Note that the main difference between DrMG+Bootp and vanilla DrMG is that DrMG+Bootp considers a less frequent meta-gradient update that considers feedback from a long horizon of optimizer steps, whereas DrMG considers a more frequent meta-gradient update that adapts the meta parameters based on per optimizer step upon per mini-batch of transitions.
For example, when we employ bootstrapping with 3 inner steps, the DrMG+Bootp optimizes η with meta-gradient once with every 3 mini-batches of data (to compute the bootstrapped inner steps with horizon 3), whereas
DrMG optimizers η once upon each mini-batch of data (since the inner step updates the optimizer for a single iteration).
[Figure]
We evaluate DrMG and DrMG+Bootp on six representative games.
We observe that DrMG ( ˆ η ) performs fairly well sometimes, e.g., on Breakout , DemonAttack and Krull , It fails on partial games, e.g., Gopher and Centipede .
We also observed the training instability when using bootstrapping.
To conclude, bootstrapping can be helpful on partial games, but it is also not straightforward to tune/adapt to DrMG.
We leave fully reproducing BMG on Atari 2600 as future work.
A. I MPLEMENTATION A.4. I MPLEMENTATION OF Q UANTILE D ISTRIBUTIONAL L OSS
We present the pseudocode for two important functions in our algorithm as follows.
VTraceOutput =
collections.namedtuple('vtrace_output', ['errors', 'pg_advantage']) def compute_loss( dist_src: Array, taus: Array, dist [Table] [Table] , huber
[Table] param: [Table] , stop_target_gradients
[Table]
bool, )
[Table] > Array: """ [Table] between two
[Table] .""" delta [Table] [Table]
None ,
[Table]
[Table]
[Table]
None ]
[Table]
_neg
[Table]
.).astype(jnp.
[Table]
) delta
[Table]
neg = lax.select(
[Table]
,sg
[Table]
[Table]
_neg) weight = jnp.abs(taus
[ None , :, None ]
[Table]
delta_neg) loss = clipping.huber_loss(delta, huber_param) loss *= weight return jnp.mean(loss, axis=-1) def quantile_regression_td_learning(support:Array, z_logits_tm1:Array,r_t:Array, discount_t:Array,z_logits_t:Array, rho_tm1:Array,eta:namedtuple, stop_target_gradients:bool,huber_param:float, ) -> VTraceOutput: """Distributional TD(lambda) learning."
"" dist_target = r_t
[:, None ] + eta.gamma * discount_t[:, None ] \ * lax.stop_gradient(z_logits_t) dist_target = lax.select(stop_target_gradients, sg(dist_target),dist_target) clipped_rhos_tm1 = sg(jnp.minimum(1.0, rho_tm1))
* eta.alpha + \ (1 - eta.alpha) * lax.stop_gradient(rho_tm1) errors =
[] err_r =td_errors = clipped_rhos_tm1[:, None ] * (dist_target - lax.stop_gradient(z_logits_tm1)) for i in reversed(range(z_logits_t.shape[0])): err_r = td_errors[i] + discount_t[i] * clipped_rhos_tm1[i] * \ eta. [Table] [Table] [Table]
* err_r errors.insert(0, err_r) errors_r =
jnp.array(errors_r)
target_dist = errors_r + lax.stop_gradient(z_logits_tm1) target_dist = jax.lax.select(stop_target_gradients,lax.stop_
[Table]
(target_dist), target_dist) quantile_loss = compute_loss(z_logits_tm1, support, target_dist, huber_param, stop
[Table]
gradients) dist [Table] bootstrap = jnp.concatenate([ eta [Table]
[Table] [Table] * \ lax [Table]
[Table] [Table]
]), ], axis [Table] [Table] estimate
[Table] None [Table] +
[Table] None ]
[Table]
bootstrap import jax from jax import lax import optax from jax import numpy as jnp from typing import Any, Tuple, Optional from collections import namedtuple
EtaOutput = collections.namedtuple(' output', ['
_', 'alpha']) class DRMGAgent : '''Base class for DrMG Agents.'''
def __init__(self, ...): raise NotImplementedError def actor
[Table]
step(self,
[Table]
.): raise NotImplementedError def two
[Table]
steps
[Table]
(self,
[Table]
[Table]
[Table]
opt_count=meta_learner_state.opt_count
+ 1) return new_params, new_meta_params, new_learner_state, \ new_meta_learner_state, logs def _inner_update_and_outer_loss(self,meta_params:Params, params:Params,rng_key:
RngKey, learner_state:OptimLearnerState,batch_base:Transition, batch_val:Transition,network:Network, meta_network:Network,optimizer:Optimizer, ): ######### inner update (_, (logs,)), grads = jax.value_and_grad( self._drmg_loss, has_aux= True )(params, batch_base, rng_key, network, meta_network, True , 'inner_', meta_params, ) grads = jax.lax.pmean(grads, axis_name=_PMAP_AXIS_NAME)
# Gather gradients from all devices.
updates, new_opt_state
= optimizer.update( grads, learner_state.opt_state, params) new_params = optax.apply_updates(params, updates) logs.update({'inner_grad_norm_unclipped': optax.global_norm(grads), 'inner_grad_norm_applied': optax.global_norm(updates), }) ######### outer update outer_loss, (meta_logs) = self._drmg_loss
( new_params, batch_val, rng_key, network, meta_network, False , 'outer_', None ) logs.update(meta_logs) return outer_loss, (new_params, new_opt_state, logs) def _drmg_loss(self,params:Params, batch:Transition,rng_key:
RngKey, network:Network,meta_network:Network, is_inner:
,prefix:str, meta_params:Any
None ,):
[Table]
[Table]
else
[Table]
[Table]
import jax import rlax import haiku as hk from jax import numpy as jnp from typing import Any, Tuple, Optional, Union from collections import namedtuple from functools import partial
RngKey = jnp.ndarray Array =
Union[np.ndarray, jnp.ndarray]
PVOutput = collections.namedtuple('PVOutput', ['policy_logits', 'value']) Action =
Any class AtariPVHead (hk.
Module): def __init__( self, num_actions: int, fc_size: int, orth_linear: bool =
False , num_value_heads: Optional[int] = 1, name: Optional[str] =
None , ):super().__init__(name=name) self._num_actions=
num_actionsself._fc_size=
fc_size self._orth_linear=
orth_linearself._num_value_heads=
num_value_heads def __call__(self, torso_output: Array) -> Tuple[Array, Array]:
Linear = partial( hk.
Linear,w_init=hk.initializers.
Orthogonal(), b_init=hk.initializers.
Constant(0) )
if self._orth_linear else hk.
Linear torso_output = Linear(self._fc_size)(torso_output)
torso_output = jax.nn.relu(torso_output) policy_logits = Linear(self._num_actions)(torso_output) value = Linear(self._num_value_heads)(torso_output) if self._num_value_heads
== 1: value = jnp.squeeze(value, axis=-1) return policy_logits, value class
ResidualBlock (hk.
Module): def __init__( self,num_channels: int, data_format: str = 'NCHW', name: Optional kernel_shape=[3, 3], stride=[1, 1], padding='SAME',data_format=self._data_format, with_bias
=with_bias, ),jax.nn.relu, hk. Conv2D(self._num_channels, kernel_shape=[3,3],stride=[1,1], padding='SAME',data_format=self._data_format, with_ = with_bias, ), ] return branch
+ x class AtariDeepTorso (hk. Module): def __init__( self,block_config: Any = [(16, 2), (32, 2), (32, 2)], data_format: str = 'NCHW', orth_linear: bool =
False , name: Optional[str] =
None ): super().__init__(name=name)self._data_format=
data_format self._orth_linear=orth_linearself.block_config=
block_config def __call__(self, x: Array) ->
Array: torso_out =
x for i, (num_channels, num_blocks) in enumerate(self.block_config): conv = hk.
Conv2D( num_channels,kernel_shape=[3, 3], stride=[1, 1], padding='SAME', data_format=self._data_format, with_bias
=with_bias, ) torso_out =
conv(torso_out) if self._data_format == 'NHWC': torso_out =
hk.max_pool( torso_out, window_shape=[1, 3, 3, 1], strides=[1, 2, 2, 1], padding='SAME', ) elif self._data_format
== 'NCHW': torso_out =
hk.max_pool( torso_out,window_shape=[1, 1, 3,
[Table]
[Table] [Table]
if obs.ndim == 4: # obs in shape [B, ...] torso_out = self._torso_module(obs) #
[Table]
: ignore out= self._head(torso_out) elif obs.ndim =
= 5: # obs in shape [T, B, ...] torso_out =
hk.
BatchApply(self._torso_module)(obs)
out = self._head(torso_out) return out
[Table]
[Table]
[Table]
def ->
: dummy_
= net return params @functools.partial(jax.jit, static_argnums=0) def forward( self,rng_key:
RngKey, params: Params, state: State, **kwargs: Any ) -> Tuple[Action, PVOutput]: policy_logits, value =
self._net.apply(params, state) distribution = rlax.softmax() action = distribution.sample(rng_key, policy_logits) return action, PVOutput(policy_logits=policy_logits, value=value)
Meta Network import jax import rlax import haiku as hk from jax import numpy as jnp from typing import Any, Tuple, Optional, Union from collections import namedtuple from functools import partial
RngKey = jnp.ndarray Array =
Union[np.ndarray, jnp.ndarray]
MetaOutput = collections.namedtuple('MetaOutput',['gamma','lambda
_', 'alpha', 'value_cost', 'entropy_cost', 'policy_cost',])
Action = Any class MetaNetwork (Network): def __init__( self, learning_rate: float, entropy_cost: float, value_cost: float, policy_cost: float, alpha: float, discounting_default: float=0.99, trace_default: float=1.0, return_shape: list=[], loss_shape: list=[], name: Optional[str] =
None ): module = functools.partial( MetaScalarModule_, learning_rate=learning_rate, entropy_cost
[Table]
@functools def forward( self,meta params: Params, state: Any, **kwargs: Any )
[Table]
> StacxOutput: stacx output class MetaScalarModule_ (
[Table]
[Table] [Table]
init
[Table] [Table] [Table]
_rate: float, entropy
[Table]
cost
[Table]
, value [Table] cost:
[Table]
, policy
[Table] cost [Table] [Table]
1.0, discounting_default: float = 0.99, trace_default: float = 1.0, return_shape: list=[], loss_shape: list = [], name: Optional[str] = None ): super(). [Table]
init
[Table]
(name=name)
[Table]
(discounting
[Table]
[Table]
def
[Table]
_call
[Table]
[Table]
, state)
[Table]
> Tuple[
[Table]
[Table]
[Table]
hk.
[Table]
[Table] [Table] , init=hk.initializers. Constant(self.discounting))
lambda_ = hk.get_parameter("lambda", self.return_shape, init=hk.initializers.
Constant(self.lambda
entropy_cost = hk.get_parameter("entropy_cost", self.loss_shape, init=
hk.initializers.
Constant(self.entropy_cost))
value_cost = hk.get_parameter("value_cost", self.loss_shape, init=hk.initializers.
Constant(self.value_cost))
policy_cost = hk.get_parameter("policy_cost", self.loss_shape, init=hk.initializers. Constant(self.policy_cost)) alpha = hk.get_parameter("alpha", self.return_shape, init=hk.initializers. Constant(self.alpha)) return MetaOutput( value_cost=jax.nn.sigmoid(value_cost),entropy_cost =jax.nn.sigmoid(entropy_cost), policy_cost=jax.nn.sigmoid(policy_cost),
[Table]