Shrinkage-based Bias-Variance Trade-off for Deep Reinforcement Learning

1 INTRODUCTION . Deep reinforcement learning ( RL ) has achieved remarkable successes recently , as witnessed by the human/super-human level performance achieved in various challenging artificial intelligence tasks ; notable examples include AlphaGo ( Silver et al. , 2017 ; 2016 ) , Atari games ( Mnih et al. , 2013 ) , robotic learning ( Levine et al. , 2016 ; Andrychowicz et al. , 2017 ) , among many others . A remarkable nature of RL is that multiple types of algorithms have been developed , based on different principles , highlighting different features and trade-offs . This includes model-based approaches which derive optimal policies from estimated dynamic models of the unknown environment ( Deisenroth & Rasmussen , 2011 ; Nagabandi et al. , 2017 ) , model-free off-policy methods such as Q-learning ( Mnih et al. , 2013 ) , as well as model-free on policy methods such as Proximal Policy Optimization ( PPO ) ( Schulman et al. , 2017 ) and Trust Region Policy Optimization ( TRPO ) ( Schulman et al. , 2015 ) . These methods exploit different properties of the Markov decision process , and have different pros and cons . Model-based and off-policy methods are known to be more sample efficient when implemented successfully , while on-policy methods are less sample efficient , but are simpler and easier to train , especially for problems with continuous or high dimensional action spaces . Although there have been emerging advances in improving each individual algorithms , more substantial improvements are likely achieved by adaptively integrating multiple strategies into an overall intelligence system . This is evident from human brain , which has been discovered to contain multiple distinct or even competing pathways for learning from reward ( Daw et al. , 2005 ; Rangel et al. , 2008 ) , including both model-free strategies consisting of error-and-trial learning to repeat rewarded actions , and model-based learning that learns models of the environment , and use it to evaluate candidate actions by mental simulations of their consequences . We would expect integrating different methods in RL will improve the performance , especially when integrated into an adaptive fashion . Algorithmically , a key challenge for an adaptive combination is to evaluate the reliability of each individual strategy , without knowing the ground truth answer while performing the task . In this work , we focus on the problem of adaptively combining policy gradient estimators from various sources . We leverage the fact that different strategies yield different gradient estimators with different bias-variance trade-off , which can be exploited by statistical shrinkage estimators to obtain optimal combinations that have better mean square error than any individual estimators . In particular , policy gradient estimators provide ( nearly ) unbiased estimators of the true gradient of the expected reward , but often suffers from high variance and is sample inefficient . Model-based and on-policy methods , on the other hand , provide more biased estimates but can have much smaller variance by taking the advantage of the large amount of imaginary data efficiently simulated from the model , or off-policy data stored in the replay buffer . In particular , we propose to use James-Stein ( JS ) shrinkage estimator ( Stein , 1956 ; James & Stein , 1961 ; Efron & Morris , 1975 ) to combine the unbiased on-policy gradient estimator with biased but low variance estimators constructed based on models or off-policy data . The basic idea of JS estimator is to access the quality of the biased estimator by inspecting its consistency with unbiased estimators , and use it as a clue for constructing an optimally weighted combination to trade-off bias and variance . JS estimator admits strong theoretical guarantees for standard Gaussian cases , and more importantly , also provides a simple practical heuristic for more complex practical non-Gaussian cases . It can potentially provide an ideal tool for solving the very challenge of bias-variance trade-off in various components of deep reinforcement learning , but seems to be still largely under-explored in the literature . In this work , we explore the practical application of JS estimators in deep RL , illustrated by two examples : Shrinking Towards Model-based Gradient We develop a hybrid algorithm which adaptively integrates on-policy policy gradient with Dyna-style model-based learning agent ( Sutton , 1991 ) . We train a transition model using the existing rollout data in parallel to the on-policy gradient descent ; imaginary rollouts are simulated from the model , providing a biased , lower variance gradient estimator . This model-based gradient estimator compared against the unbiased , high variance on-policy gradient , and a weighted combination is constructed using JS estimator based on the consistency of the two estimators . Our method improves the sample efficiency by exploiting the data more efficiently using model learning while avoiding significant bias introduced by the model using adaptive JS estimator . Shrinking Towards Temporally Smoothed Gradient It has been observed heuristically that proper bias is helpful for training ( Sutskever et al. , 2013 ; Nesterov , 2013 ) . We present a simple method that explicitly constructs an optimal bias-variance trade-off by shrinking the on-policy gradient estimator towards a smoothed gradient estimator obtained by a weighted sum over previous iterations . The smoothed gradient provides a reasonable guess of the true gradient of the current estimation and forms a biased , low variance estimator . By shrinking towards the smoothed gradient , we significantly decrease the variance of the on-policy gradient , at the cost of introducing a small bias . This effectively creates a statistical momentum , which shares similarity with the traditional momentum methods derived from optimization perspectives , but with an adaptive coefficient determined by the JS estimator . It can also be combined with traditional momentum to yield gains on both statistical and optimization sides . This simple algorithm effectively leverages the historic ( off-policy ) data via the temporally smoothed gradient , yielding higher sample efficiency than standard policy gradient . We perform empirical studies to test the performance of both above examples , showing our adaptive combination strategies can yield better results than any fixed-weight combination , significantly improving the sample efficiency on standard on-policy methods . This illustrates the practical power of JS estimators , which we expect to be broadly applicable to more settings in deep RL in future works . 2 BACKGROUND . We first introduce the basics of reinforcement learning and model-free policy gradient , and set up the notation that we will use in the paper . We then introduce a model-based policy gradient method . 2.1 REINFORCEMENT LEARNING . Reinforcement learning considers the problem of finding an optimal policy for an agent which interacts with an uncertain environment and collects reward per action . The goal of the agent is to maximize the long-term cumulative reward . Formally , this problem can be formulated as a Markov decision process over the environment states s ∈ S and agent actions a ∈ A , under an unknown environmental dynamic defined by a transition probability T ( s′|s , a ) and a reward signal r ( s , a ) immediately following the action a performed at state s. The agent ’ s action a is selected by a conditional probability distribution π ( a|s ) called policy . In policy optimization , we consider a set of candidate policies πθ ( a|s ) parameterized by θ and obtain the optimal policy by maximizing the expected cumulative reward or return J ( θ ) = Es∼ρπ , a∼π ( a|s ) [ r ( s , a ) ] , where ρπ ( s ) = ∑∞ t=1 γ t−1Pr ( st = s ) is the normalized discounted state visitation distribution with discount factor γ ∈ [ 0 , 1 ) . To simplify the notation , we denote Es∼ρπ , a∼π ( a|s ) [ · ] by simply Eπ [ · ] in the rest of paper . 2.2 ON-POLICY POLICY GRADIENT . According to the policy gradient theorem ( Sutton & Barto , 1998 ) , the gradient of J ( θ ) equals ∇θJ ( θ ) = Eπ [ ∇θ log π ( a|s ) Qπ ( s , a ) ] , ( 1 ) whereQπ ( s , a ) = Eπ [ ∑∞ t=1 γ t−1r ( st , at ) |s1 = s , a1 = a ] denotes the expected return under policy π starting from state s and action a . Different policy gradient methods are based on different stochastic estimations of the expected gradient in Eq ( 1 ) . Perhaps the most straightforward way is to simulate the environment with the current policy π to obtain a trajectory { ( st , at , rt ) } nt=1 and estimate∇θJ ( θ ) using the Monte Carlo estimation : ∇̂θJ ( θ ) = 1 n n∑ t=1 γt−1∇θ log π ( at|st ) Q̂π ( st , at ) , ( 2 ) where Q̂π ( st , at ) is an empirical estimate of Qπ ( st , at ) , e.g. , Q̂π ( st , at ) = ∑ j≥t γ j−trj . This is known as an on-policy gradient estimator because it only uses the data from the current policy . Proximal Policy Optimization ( PPO ) ( Schulman et al. , 2017 ; Heess et al. , 2017 ) is one of the state-theart model-free policy gradient methods for policy optimization . It uses a proximal Kullback-Leibler ( KL ) divergence penalty to regularize and stabilize the policy gradient update . Given an existing policy πold , PPO obtains a new policy by maximizing the following surrogate loss function Jppo ( θ ) = Eπold [ πθ ( a|s ) πold ( a|s ) Qπ ( s , a ) − λKL [ πold ( ·|s ) || πθ ( ·|s ) ] ] , where the first term is an approximation of the expected reward , and the second term enforces the the updated policy to be close to the previous policy under KL divergence . The gradient of Jppo ( θ ) can be rewritten as ∇θJppo ( θ ) = Eπold [ wπ ( s , a ) ∇θ log π ( a|s ) Qπλ ( s , a ) ] where wπ ( s , a ) : = πθ ( a|s ) /πold ( a|s ) is the density ratio of the two polices , and Qπλ ( s , a ) : = Qπ ( s , a ) + λwπ ( s , a ) −1 where the second term comes from the KL penalty . The advantage of using on-policy estimation is that it provides ( nearly ) unbiased estimation of the gradient of expected reward . A well-known challenge , however , is that it has large variance , given that it is only possible to use a small number of data from the current policy at each iteration because the cost of simulation from the real environment is high . In addition , it does not take advantage of the off-policy data collected previously . Using off-policy data allows us to obtain low variance but potentially highly biased estimates , which can be combined with on-policy gradient with James-Stein estimator as we propose in this work . In the sequel , we introduce model-based policy gradient as an important way for leveraging off-policy data . 2.3 MODEL-BASED POLICY GRADIENT . The model-based methods derive optimal polices by learning a transition model of the underlying dynamic T ( s′|s , a ) of the Markov Decision Process . Assume the states s is continuous , it is common to assume s′ = s+fφ ( s , a ) +ξ where fφ is a parametric function ( e.g. , neural network ) with learnable parameter φ , and ξ is a zero-mean noise . The model parameter φ is often trained by minimizing the L2 one-step prediction loss , min φ n∑ t=1 ‖st+1 − st − fφ ( st , at ) ‖22 , ( 3 ) where D : = ( st , at , st+1 ) nt=1 is a set of transition pairs collected from previous rollouts . See Deisenroth & Rasmussen ( 2011 ) ; Fu et al . ( 2016 ) ; Nagabandi et al . ( 2017 ) ; Thanard Kurutach ( 2018 ) ; Depeweg et al . ( 2017 ) for variety of the recent variants of model-based learning methods . With a learned model , the problem of estimating the optimal policy becomes a planning problem , which can be solved either using various classical planning and control methods , typically based on variants of dynamic programming ( see e.g. , Sutton & Barto , 1998 ) , or sample-based methods that generate fictitious data by simulating the model forward and obtain optimal policies by applying standard RL algorithm on the fictitious data . In this work , we will mainly consider a simple sample-based policy gradient method whose main idea is rooted from Dyna ( Sutton , 1991 ) . The idea is simply generating a fictitious data ( s̃t , ãt , r̃t , s̃t+1 ) by simulating the model , and plug it into the policy gradient formula ( 2 ) , ∇̂modelθ J ( θ ) = n∑ t=1 γt−1∇θ log π ( ãt|s̃t ) Q̂π ( s̃t , ãt ) . ( 4 ) By simulating a large number of fictitious data from the model , which is much less expensive than simulating from the real environment , we can make the variance of ∇̂modelθ J ( θ ) significantly smaller than the variance of the on-policy gradient in ( 2 ) . The idea of using ∇̂modelθ J ( θ ) can be found in ( Gu et al. , 2016 ; Thanard Kurutach , 2018 ) . The problem , however , is that the transition model can deviate significantly from the true model , introducing a potentially large bias that makes the training problem unstable . Thanard Kurutach ( 2018 ) found that by training an ensemble of models can help stabilize the training . As we discuss in the sequel , one of our key idea is to calculate on-policy and model-based gradient simultaneously , and use the unbiased information of the on-policy gradient to evaluate the biasness of model-based gradient , and weigh its importance accordingly to ensure the stability while exploiting the model-based information .

Summary: This paper presents two algorithms that improve on PPO by using James-Stein (JS) shrinkage estimator. The first algorithm, PPO-MBS, combines low bias of on-policy methods with low variance of model-based RL algorithms. The second, PPO-STS, uses JS to create a statistical momentum and reduce variance of the PPO algorithm. Both algorithms are evaluated on Mujoco environments aiming to show improvements in average cumulative reward, reduced bias and variance.