Unifying Gradient Estimators for Meta-Reinforcement Learning  via Off-Policy Evaluation

1 Introduction . Recent years have witnessed the success of reinforcement learning ( RL ) in challenging domains such as game playing [ 1 ] , board games [ 2 ] , and robotics control [ 3 ] . However , despite such breakthroughs , state-of-the-art RL algorithms are still plagued by sample inefficiency , as training agents requires orders of magnitude more samples than humans would experience to reach similar levels of performance [ 1 , 2 ] . One hypothesis on the source of such inefficiencies is that standard RL algorithms are not good at leveraging prior knowledge , which implies that whenever presented with a new task , the algorithms must learn from scratch . On the contrary , humans are much better at transferring prior skills to new scenarios , an innate ability arguably obtained through evolution over thousand of years . Meta-reinforcement learning ( meta-RL ) formalizes the learning and transfer of prior knowledge in RL [ 4 ] . The high-level idea is to have an RL agent that interacts with a distribution of environments at meta-training time . The objective is that at meta-testing time , when the agent interacts with previously unseen environments , it can learn much faster than meta-training time . Here , faster learning is measured by the number of new samples needed to achieve a good level of performance . If an agent can achieve good performance at meta-testing time , it embodies the ability to transfer knowledge from prior experiences at meta-training time . meta-RL algorithms can be constructed in many ways , such as those based on recurrent memory [ 5 , 6 ] , gradient-based adaptations [ 4 ] , learning loss functions [ 7–9 ] , probabilistic inference of context variables [ 10–12 ] , online adaptation of hyper-parameters within a single lifetime [ 13 , 14 ] and so on . Some of these formulations have mathematical connections ; see e.g . [ 15 ] for an in-depth discussion . We focus on gradient-based adaptations [ 4 ] , where the agent carries out policy gradient updates [ 16 ] at both meta-training and meta-testing time . Conceptually , since meta-RL seeks to optimize the 1https : //github.com/robintyh1/neurips2021-meta-gradient-offpolicy-evaluation 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . way in which the agent adapts itself in face of new environments , it needs to differentiate through the policy gradient update itself . This effectively reduces the meta-RL problem into estimations of Hessian matrices of value functions . Challenges for computing Hessian matrix of value functions . To calculate Hessian matrices ( or unbiased estimates thereof ) for supervised learning objectives , it suffices to differentiate through gradient estimates , which can be easily implemented with auto-differentiation packages [ 17–19 ] . However , it is not the case for value functions in RL . Intuitively , this is because value functions are defined via expectations with respect to distributions that themselves depend on the policy parameters of interest , whereas in supervised learning the expectations are defined with respect to a fixed data distribution . As a result , implementations that do not take this into account may lead to estimates [ 4 ] whose bias is not properly characterized and might have a negative impact on downstream applications [ 20 ] . Motivated by this observation , a number of prior works suggest implementation alternatives that lead to unbiased Hessian estimates [ 21 ] , with potential variance reduction [ 22 , 23 , 20 ] ; or biased estimates with small variance [ 24 ] . However , different algorithms in this space are motivated and derived in seemingly unrelated ways : for example , [ 21–23 ] derive code-level implementations within the general context of stochastic computation graphs [ 25 ] . On the other hand , [ 24 ] derives the estimates by explicitly analyzing certain terms with potentially high variance , which naturally produces bias in the final estimate . Due to apparently distinct ways of deriving estimates , it is not immediately clear how all such methods are related , and whether there could be other alternatives . Central contribution . We present a unified framework for estimating higher-order derivatives of value functions , based on the concept of off-policy evaluation . The main insights are summarized in Table 1 , where most aforementioned prior work can be interpreted as special cases of our framework . Our framework has a few advantages : ( 1 ) it conceptually unifies a few seemingly unrelated prior methods ; ( 2 ) it elucidates the bias and variance trade-off of the estimates ; ( 3 ) it naturally produces new methods based on Taylor expansions of value functions [ 29 ] . After a brief background introduction on meta-RL in Section 2 , we will discuss the above aspects in detail in Section 3 . From an implementation perspective , we will show in Section 4.1 that both the general framework and the new method can be conveniently implemented in auto-differentiation libraries [ 30 , 19 ] , making it amenable in a practical setup . Finally in Section 5 , we validate important claims based on this framework with experimental insights . 2 Background . We begin with the notation and background on RL and meta-RL . 2.1 Task-based reinforcement learning . Consider a Markov decision process ( MDP ) with state space X and action space A . At time t ≥ 0 , the agent takes action at ∈ A in state xt ∈ X , receives a reward rt and transitions to a next state xt+1 ∼ p ( ·|xt , at ) . Without loss of generality , we assume a single starting state x0 . Here , we assume the reward rt = r ( xt , at , g ) to be a deterministic function of state-action pair ( xt , at ) and the task variable g ∈ G. The task variable g ∼ pG is resampled for every episode . For example , xt ∈ X is the sensory space of a running robot , at ∈ A is the control , g is the episodic target direction in which to run and r ( xt , at , g ) is the speed in direction g. A policy π : X → P ( A ) specifies a distribution over actions at each state . For convenience , we define the value function V π ( x , g ) : = Eπ [ ∑∞ t=0 γ tr ( xt , at , g ) | x0 = x ] and Q-function Qπ ( x , a , g ) : = Eπ [ ∑∞ t=0 γ tr ( xt , at , g ) | x0 = x , a0 = a ] . Without loss of generality , we assume the policy to be smoothly parameterized as πθ with parameter θ ∈ RD . We further assume that the MDPs terminate within a finite horizon of H < ∞ under any policy . 2.2 Gradient-based meta-reinforcement learning . The motivation of meta-RL is to identify a policy πθ such that given a task g , after updating the parameter from θ to θ′ with a parameter update computed under rewards r ( x , a , g ) , the resulting policy πθ′ performs well . Formally , define U ( θ , g ) ∈ RD as a parameter update to θ , for example , one policy gradient ascent step . Model-agnostic meta-learning ( MAML ) [ 4 ] formulates meta-RL as optimizing the value function at the updated policy Eg∼pG [ V πθ′ ( x , g ) ] where θ′ = θ+U ( θ , g ) is the updated policy . The optimization is with respect to the initial policy parameter θ . The aim is to find θ such that it entails fast learning ( or adaptation ) to the environment , through the inner loop update operator U ( θ , g ) , that leads to high-performing updated policy θ′ . Consider the following problem , max θ F ( θ ) : = Eg∼pG [ V πθ′ ( x0 , g ) ] , θ′ = θ + U ( θ , g ) . ( 1 ) We can decompose the meta-gradient into two terms based on the chain rule , ∇θF ( θ ) = Eg∼pG [ ∇θF ( θ , g ) ] : = Eg∼pG [ ∇θθ′∇θ′V πθ′ ( x0 , g ) ] ( 2 ) where∇θθ′ = I+∇θU ( θ , g ) ∈ RD×D is a matrix , and∇θ′V πθ′ ( x0 , g ) is the vanilla policy gradient evaluated at the updated parameter θ′ [ 31 ] . A straightforward way to optimizing Eqn ( 1 ) is to carry out the outer loop update θ ← θ + α∇θF ( θ ) with learning rate α > 0 . Policy gradient update . Following MAML [ 4 ] , we focus on policy gradient update where U ( θ , g ) = η∇θV πθ ( x0 , g ) for a fixed step size η > 0 . The matrix ∇θU ( θ , g ) = η∇2θV πθ ( x0 , g ) is the Hessian of the value function with respect to policy parameters ; henceforth , we define Hθ ( x0 , g ) : = ∇2θV πθ ( x0 , g ) ∈ RD×D . Estimating meta-gradients . Practical algorithms construct stochastic estimates of the metagradients using a finite number of samples . In Eqn ( 2 ) , we decompose the meta-gradients into the multiplication of a Hessian matrix with a policy gradient vector . Given a fixed task variable g , a common practice of prior work is to construct the gradient estimate as the product of the Hessian estimate and the policy gradient estimate ∇̂F ( θ , g ) = ( I + ηĤθ ( x0 , g ) ) ∇̂θ′V πθ′ ( x0 , g ) ; see , e.g. , [ 4 , 32 , 33 , 24 , 21–23 , 20 ] . See Algorithm 1 for the full pseudocode for estimating meta-gradients by sampling multiple tasks . Since there is a large literature on constructing accurate estimates to policy gradient ( e.g. , actor-critic algorithms use baselines for variance reduction [ 34 ] ) , the main challenge consists in estimating the Hessian matrix accurately . Bias of common plug-in estimators . Common practices in meta-RL algorithms rely on the premise that if both Hessian and policy gradient estimates are unbiased , then the meta-gradient estimate is unbiased too . E [ Ĥθ ( x0 , g ) ] = Hθ ( x0 , g ) , E [ ∇̂θ′V πθ′ ( x0 , g ) ] = ∇θ′V πθ′ ( x0 , g ) ⇒ E [ ∇̂θF ( θ , g ) ] = ∇θF ( θ , g ) . Unfortunately , this is not true . This is because the two estimates are in general correlated when the sample size is finite , leading to the bias of the overall estimate . We provide further discussions in Appendix A . For the rest of the paper , we follow practices of prior work and focus on the properties of Hessian estimates , leaving a more proper treatment of this bias to future work . 3 Deriving Hessian estimates with off-policy evaluation . Since the meta-gradient estimates are computed by averaging over task variables g ∼ pG , in the following , we focus on Hessian estimates at a single state and task variable Hθ ( x , g ) with a fixed g. In this section , we also drop the dependency of the value function on g , such that , e.g. , V πθ ( x0 ) ≡ V πθ ( x0 , g ) and Q ( x , a , g ) ≡ Q ( x , a ) .

Summary: The paper is concerned with meta reinforcement learning. The author(s) introduced a framework for learning high-order derivatives of value functions using off-policy evaluation. Specifically, the author(s) proposed to adopt existing off-policy evaluation methods to derive the value estimator and then take the derivative with respect to the target parameters to construct the estimate. This technique is further adopted to implement meta reinforcement learning.