On the Convergence Theory of Debiased Model-Agnostic Meta-Reinforcement Learning

1 Introduction . Meta-learning has recently attracted much attention as a learning to learn approach that enables quick adaptation to new tasks using past experience and data . This is a particularly promising approach for Reinforcement Learning ( RL ) where in several applications , such as robotics , a group of agents encounter new tasks and need to learn new behaviors or policies through a few interactions with the environment building on previous experience [ 1–9 ] . Among various forms of Meta-learning , gradient-based Model-Agnostic Meta-Learning ( MAML ) formulation [ 1 ] is a particularly effective approach which , as its name suggests , can be applied to any learning problem that is trained with gradient-based updates . In MAML , we exploit observed tasks at training time to find an initial model that is trained in a way that rapidly adapts to a new unseen task at test time , after running a few steps of a gradient-based update with respect to the loss of the new task . The MAML formulation can be extended to RL problems if we represent each task as a Markov Decision Process ( MDP ) . In this setting , we assume that we are given a set of MDPs corresponding to the tasks that we observe during the training phase and assume that the new task at test time is drawn from an underlying probability distribution . The goal in Model-Agnostic Meta-Reinforcement Learning ( MAMRL ) is to exploit this data to come up with an initial policy that adapts to a new task ( drawn from the same distribution ) at test time by taking a few stochastic policy gradient steps [ 1 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Several algorithms have been proposed in the context of MAMRL [ 1 , 9–12 ] which demonstrate the advantage of this framework in practice . None of these methods , however , are supported by theoretical guarantees for their convergence rate or overall sample complexity . Moreover , these methods aim to solve a specific form of MAMRL that does not fully take into account the stochasticity aspect of RL problems . To be more specific , the original MAMRL formulation proposed in [ 1 ] assumes performing one step of policy gradient to update the initial model at test time . However , as mentioned in the experimental evaluation section in [ 1 ] , it is more common in practice to use stochastic policy gradient , computed over a batch of trajectories , to update the initial model at test time . This is mainly due to the fact that computing the exact gradient of the expected reward is not computationally tractable due to the massive number of possible state-action trajectories . As a result , the algorithm developed in [ 1 ] is designed for finding a proper initial policy that performs well after one step of policy gradient , while in practice it is implemented with stochastic policy gradient steps . Due to this difference between the formulation and what is used in practice , the ascent step used in MAML takes a gradient estimate which suffers from a non-diminishing bias . As the variance of gradient estimation is also non-diminishing , the resulting algorithm would not achieve exact first-order optimality . To be precise , in stochastic nonconvex optimization , if we use an unbiased gradient estimator , along with a small stepsize or a large batch size to control the variance , the iterates converge to a stationary point . However , if we use a biased estimator with non-vanishing bias and variance , exact convergence to a stationary point is not achievable , even if the variance is small . Contributions . The goal of this paper is to solve the modified formulation of model-agnostic metareinforcement learning problem in which we perform a stochastic policy gradient update at test time instead of ( deterministic ) policy gradient . To do so , we propose a novel stochastic gradientbased method for Meta-Reinforcement Learning ( SG-MRL ) , which is designed for stochastic policy gradient steps at test time . We show that SG-MRL implements an unbiased estimate of its objective function gradient which allows achieving first-order optimality in non-concave settings . Moreover , we characterize the relation between batch sizes and other problem parameters and the best accuracy that SG-MRL can achieve in terms of gradient norm . We show that , for any > 0 , SG-MRL can find an -first-order stationary point if the learning rate is sufficiently small or the batch of tasks is large enough . To the best of our knowledge , this is the first result on the convergence of MAMRL methods . Moreover , we show that our analysis can be extended to the case where more than one step of stochastic policy gradient is taken during test time . For simplicity , we state all the results in the body of the paper for the single-step case and include the derivations of the general multiple steps case in the appendices . We also empirically validate the proposed SG-MRL algorithm in larger-scale environments standard in modern reinforcement learning applications , including a 2D-navigation problem , and a more challenging locomotion problem simulated with the MuJoCo library . Related work . Although this paper provides the first theoretical study of MAML for RL , several recent papers have studied the complexity analysis of MAML in other contexts . In particular , the iMAML algorithm which performs an approximation of one step of proximal point method ( instead of a few steps of gradient descent ) in the inner loop was proposed in [ 13 ] . The authors focus on the deterministic case , and show that , assuming the inner loop loss function is sufficiently smooth , i.e. , the regularized inner loop function is strongly convex , iMAML converges to a first-order stationary point . Another recent work [ 14 ] establishes convergence guarantees of the MAML method to first-order stationarity for non-convex settings . Also , [ 15 ] extends the theoretical framework in [ 14 ] to the multiple-step case . However , the results in [ 14 , 15 ] can not be applied to the reinforcement learning setting . This is mainly due to the fact that the probability distribution over possible trajectories of states and actions varies with the policy parameter , leading to a different algorithm that has an additional term which makes the analysis , such as deriving an upper bound on the smoothness parameter , more challenging . We will discuss this point in subsequent sections . The online meta-learning setting has also been studied in a number of recent works [ 16–18 ] . In particular , [ 17 ] studies this problem for convex objective functions by casting it in the online convex optimization framework . Also , [ 16 ] extends the model-agnostic setup to the online learning case by considering a competitor which adapts to new tasks , and propose the follow the meta leader method which obtains a sublinear regret for strongly convex loss functions . It is also worth noting that another notion of bias that has been studied in the MAMRL literature [ 10 , 19 ] differs from what we consider in our paper . More specifically , as we will show later , the derivative of the MAML objective function requires access to the second-order information , i.e. , Hessian . In [ 1 ] , the authors suggest a first-order approximation which ignores this second-order term . This leads to a biased estimate of the derivative of the MAML objective function , and a number of recent works [ 10 , 19 ] focus on providing unbiased estimates for the second-order term . In contrast , here we focus on biased gradient estimates where the bias stems from the fact that in most real settings we do not have access to all possible trajectories and we only have access to a mini-batch of possible trajectories . In this case , even if one has access to the second-order term required in the update of MAML , the bias issue we discuss here will remain . 2 Problem formulation . Let { Mi } i be the set of Markov Decision Processes ( MDPs ) representing different tasks1 . We assume these MDPs are drawn from a distribution p ( which we can only draw samples from ) , and also the time horizon is fixed and is equal to { 0 , 1 , ... , H } for all tasks . For the i-th MDP denoted by Mi , which corresponds to task i , we denote the set of states and actions by Si and Ai , respectively . We also assume the initial distribution over states in Si is given by µi ( · ) and the transition kernel is denoted by Pi , i.e. , the probability of going from state s ∈ Si to s′ ∈ Si given taking action a ∈ Ai is Pi ( s′|s , a ) . Finally , we assume at state s and by taking action a , the agent receives reward ri ( s , a ) . To summarize , an MDPMi is defined by the tuple ( Si , Ai , µi , Pi , ri ) . For MDPMi , the actions are chosen according to a random policy which is a mixed strategy over the set of actions and depends on the current state , i.e. , if the system is in state s ∈ Si , the agent chooses action a ∈ Ai with probability πi ( a|s ) . To search over the space of all policies , we assume these policies are parametrized with θ ∈ Rd , and denote the policy corresponding to parameter θ by πi ( ·|· ; θ ) . A realization of states and actions in this setting is called a trajectory , i.e. , a trajectory of MDPMi can be written as τ = ( s0 , a0 , ... , sH , aH ) where ah ∈ Ai and sh ∈ Si for any 0 ≤ h ≤ H . Note that , given the above assumptions , the probability of this particular trajectory is given by qi ( τ ; θ ) : = µi ( s0 ) H∏ h=0 πi ( ah|sh ; θ ) H−1∏ h=0 Pi ( sh+1|sh , ah ) . ( 1 ) Also , the total reward received over this trajectory isRi ( τ ) : = ∑H h=0 γ hri ( sh , ah ) , where 0 ≤ γ ≤ 1 is the discount factor . As a result , for MDPMi , the expected reward obtained by choosing policy π ( ·|· ; θ ) is given by Ji ( θ ) : = Eτ∼qi ( · ; θ ) [ Ri ( τ ) ] . ( 2 ) It is worth noting that the gradient∇Ji ( θ ) admits the following characterization [ 20–22 ] ∇Ji ( θ ) = Eτ∼qi ( · ; θ ) [ gi ( τ ; θ ) ] , ( 3 ) where gi ( τ ; θ ) is defined as gi ( τ ; θ ) : = H∑ h=0 ∇θ log πi ( ah|sh ; θ ) Rhi ( τ ) , ( 4 ) if we define Rhi ( τ ) as Rhi ( τ ) : = ∑H t=h γ tri ( st , at ) . In practice , evaluating the exact value of ( 3 ) is not computationally tractable . Instead , one could first acquire a batch Di , θ of trajectories drawn independently from distribution qi ( · ; θ ) , and then , estimate∇Ji ( θ ) by ∇̃Ji ( θ , Di , θ ) : = 1 |Di , θ| ∑ τ∈Di , θ gi ( τ ; θ ) . ( 5 ) Also , we denote the probability of choosing ( with replacement ) an independent batch of trajectories Di , θ by qi ( Di , θ ; θ ) ( see Appendix A.1 for a remark on this ) . In this setting , the goal of Model-Agnostic Meta-Reinforcement Learning problem introduced in [ 1 ] is to find a good initial policy that performs well in expectation when it is updated using one or a few steps of stochastic policy gradient with respect to a new task . In particular , for the case of performing one step of stochastic policy gradient , the problem can be written as2 max θ∈Rd V1 ( θ ) : = Ei∼p [ EDitest [ Ji ( θ + α∇̃Ji ( θ , Ditest ) ) ] ] . ( 6 ) 1To simplify the analysis , we assume the number of tasks is finite 2From now on , we suppress the θ dependence of batches to simplify the notation . Note that by solving this problem we find an initial policy ( Meta-policy ) that in expectation performs well if we evaluate the output of our procedure after running one step of stochastic policy gradient on this initial policy for a new task . This formulation can be extended to the setting with more than one step of stochastic policy gradient as well . To state the problem formulation in this case , let us first define Ψi which is an operator that takes model θ and batch Di as input and performs one step of stochastic gradient policy at point θ and with respect to function Ji and batch Di , i.e. , Ψi ( θ , Di ) : = θ + α∇̃Ji ( θ , Di ) . Now , we extend problem ( 6 ) to the case where we are looking for an initial point which performs well on expectation after it is updated with ζ steps of stochastic policy gradient with respect to a new MDP drawn from distribution p. This problem can be written as max θ∈Rd Vζ ( θ ) : = Ei∼p [ E { Ditest , t } ζt=1 [ Ji ( Ψi ( . . . ( Ψi ( θ , Ditest,1 ) . . . ) , Ditest , ζ ) ) ] ] , ( 7 ) where the operator Ψi is applied ζ times inside the expectation . In this paper , we establish convergence properties of policy gradient methods for both single step and multiple steps of stochastic gradient cases , but for simplicity in the main text we focus on the single step case .

Summary: This paper deals with the problem of few-shot meta learning, i.e. where one explicitly optimizes for the performance of doing one (or more) steps of SGD on a new task. This was intially introduced through the MAML method. However, the authors argue that MAML only discusses gradient descent while in practise, approximations lead to doing stochastic gradient desecent instead, thus making the original formulation of the problem biased. To that end, this paper proposes an SGD based meta learning objective, which ensures simpler convergence analysis as well as better performance in two benchmark problems from the RL setting.