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 uses the recently developed framework of Taylor expansions for value functions (ref [29] in the paper), expands the results, and shows how they could be applied to Off-Policy Evaluation, and by a non-trivial extension to meta-learning. The main algorthmic contribution of the paper is the application of the ideas to meta-learning, but I feel the main contribution is a conceptual one, further developing the ideas from [29] and drawing theoretical connections between OPE and meta-learning.  The paper includes an experimental section demonstrating marginal improvement over baselines in several domains.