Variance Reduction for Reinforcement Learning in Input-Driven Environments

1 INTRODUCTION . Deep reinforcement learning ( RL ) has emerged as a powerful approach for sequential decision-making problems , achieving impressive results in domains such as game playing ( Mnih et al. , 2015 ; Silver et al. , 2017 ) and robotics ( Levine et al. , 2016 ; Schulman et al. , 2015a ; Lillicrap et al. , 2015 ) . This paper concerns RL in input-driven environments . Informally , input-driven environments have dynamics that are partially dictated by an exogenous , stochastic input process . Queuing systems ( Kleinrock , 1976 ; Kelly , 2011 ) are an example ; their dynamics are governed by not only the decisions made within the system ( e.g. , scheduling , load balancing ) but also the arrival process that brings work ( e.g. , jobs , customers , packets ) into the system . Input-driven environments also arise naturally in many other domains : network control and optimization ( Winstein & Balakrishnan , 2013 ; Mao et al. , 2017 ) , robotics control with stochastic disturbances ( Pinto et al. , 2017 ) , locomotion in environments with complex terrains and obstacles ( Heess et al. , 2017 ) , vehicular traffic control ( Belletti et al. , 2018 ; Wu et al. , 2017 ) , tracking moving targets , and more ( see Figure 1 ) . We focus on model-free policy gradient RL algorithms ( Williams , 1992 ; Mnih et al. , 2016 ; Schulman et al. , 2015a ) , which have been widely adopted and benchmarked for a variety of RL tasks ( Duan et al. , 2016 ; Wu & Tian , 2017 ) . A key challenge for these methods is the high variance in the gradient estimates , as such variance increases sample complexity and can impede effective learning ( Schulman et al. , 2015b ; Mnih et al. , 2016 ) . A standard approach to reduce variance is to subtract a “ baseline ” from the total reward ( or “ return ” ) to estimate the policy gradient ( Weaver & Tao , 2001 ) . The most common choice of a baseline is the value function — the expected return starting from the state . Our main insight is that a state-dependent baseline — such as the value function — is a poor choice in input-driven environments , whose state dynamics and rewards are partially dictated by the input process . In such environments , comparing the return to the value function baseline may provide limited information about the quality of actions . The return obtained after taking a good action may be poor ( lower than the baseline ) if the input sequence following the action drives the system to unfavorable states ; similarly , a bad action might end up with a high return with an advantageous input sequence . Intuitively , a good baseline for estimating the policy gradient should take the specific instance of the input process — the sequence of input values — into account . We call such a baseline an input-dependent baseline ; it is a function of both the state and the entire future input sequence . We formally define input-driven Markov decision processes , and we prove that an input-dependent baseline does not introduce bias in standard policy gradient algorithms such as Advantage Actor Critic ( A2C ) ( Mnih et al. , 2016 ) and Trust Region Policy Optimization ( TRPO ) ( Schulman et al. , 2015a ) , provided that the input process is independent of the states and actions . We derive the optimal input-independent baseline and a simpler one to work with in practice ; this takes the form of a conditional value function — the expected return given the state and the future input sequence . Input-dependent baselines are harder to learn than their state-dependent counterparts ; they are highdimensional functions of the sequence of input values . To learn input-dependent baselines efficiently , we propose a simple approach based on meta-learning ( Finn et al. , 2017 ; Vilalta & Drissi , 2002 ) . The idea is to learn a “ meta baseline ” that can be specialized to a baseline for a specific input instantiation using a small number of training episodes with that input . This approach applies to applications in which an input sequence can be repeated during training , e.g. , applications that use simulations or experiments with previously-collected input traces for training ( McGough et al. , 2017 ) . We compare our input-dependent baseline to the standard value function baseline for the five tasks illustrated in Figure 1 . These tasks are derived from queuing systems ( load balancing heterogeneous servers ( Harchol-Balter & Vesilo , 2010 ) ) , computer networks ( bitrate adaptation for video streaming ( Mao et al. , 2017 ) ) , and variants of standard continuous control RL benchmarks in the MuJoCo physics simulator ( Todorov et al. , 2012 ) . We adapted three widely-used MuJoCo benchmarks ( Duan et al. , 2016 ; Clavera et al. , 2018a ; Heess et al. , 2017 ) to add a stochastic input element that makes these tasks significantly more challenging . For example , we replaced the static target in a 7-DoF robotic arm target-reaching task with a randomly-moving target that the robot aims to track over time . Our results show that input-dependent baselines consistently provide improved training stability and better eventual policies . Input-dependent baselines are applicable to a variety of policy gradient methods , including A2C , TRPO , PPO , robust adversarial RL methods such as RARL ( Pinto et al. , 2017 ) , and meta-policy optimization such as MB-MPO ( Clavera et al. , 2018b ) . Video demonstrations of our experiments are available at https : //sites.google.com/view/input-dependent-baseline/ . 2 PRELIMINARIES . Notation . We consider a discrete-time Markov decision process ( MDP ) , defined by ( S , A , P , ρ0 , r , γ ) , where S ⊆ Rn is a set of n-dimensional states , A ⊆ Rm is a set ofm-dimensional actions , P : S ×A× S → [ 0 , 1 ] is the state transition probability distribution , ρ0 : S → [ 0 , 1 ] is the distribution over initial states , r : S ×A → R is the reward function , and γ ∈ ( 0 , 1 ) is the discount factor . We denote a stochastic policy as π : S × A → [ 0 , 1 ] , which aims to optimize the expected return η ( π ) = Eτ [ ∑∞ t=0 γ tr ( st , at ) ] , where τ = ( s0 , a0 , ... ) is the trajectory following s0 ∼ ρ0 , at ∼ π ( at|st ) , st+1 ∼ P ( st+1|st , at ) . We use Vπ ( st ) = Eat , st+1 , at+1 , ... [ ∑∞ l=0 γ lr ( st+l , ar+l ) |st ] to define the value function , and Qπ ( st , at ) = Est+1 , at+1 , ... [ ∑∞ l=0 γ lr ( st+l , ar+l ) |st , at ] to define the state-action value function . For any sequence ( x0 , x1 , ... ) , we use x to denote the entire sequence and xi : j to denote ( xi , xi+1 , ... , xj ) . Policy gradient methods . Policy gradient methods estimate the gradient of expected return with respect to the policy parameters ( Sutton et al. , 2000 ; Kakade , 2002 ; Gu et al. , 2017 ) . To train a policy πθ parameterized by θ , the Policy Gradient Theorem ( Sutton et al. , 2000 ) states that ∇θη ( πθ ) = Es∼ρπ a∼πθ [ ∇θ log πθ ( a|s ) Qπθ ( s , a ) ] , ( 1 ) where ρπ ( s ) = ∑∞ t=0 [ γ t Pr ( st = s ) ] denotes the discounted state visitation frequency . Practical algorithms often use the undiscounted state visitation frequency ( i.e. , γ = 1 in ρπ ) , which can make the estimation slightly biased ( Thomas , 2014 ) . Estimating the policy gradient using Monte Carlo estimation for the Q function suffers from high variance ( Mnih et al. , 2016 ) . To reduce variance , an appropriately chosen baseline b ( st ) can be subtracted from the Q-estimate without introducing bias ( Greensmith et al. , 2004 ) . The policy gradient estimation with a baseline in Equation ( 1 ) becomes Eρπ , πθ [ ∇θ log πθ ( a|s ) ( Qπθ ( s , a ) − b ( s ) ) ] . While an optimal baseline exists ( Greensmith et al. , 2004 ; Wu et al. , 2018 ) , it is hard to estimate and often replaced by the value function b ( st ) = Vπ ( st ) ( Sutton & Barto , 2017 ; Mnih et al. , 2016 ) . 3 MOTIVATING EXAMPLE . We use a simple load balancing example to illustrate the variance introduced by an exogenous input process . As shown in Figure 2a , jobs arrive over time and a load balancing agent sends them to one of two servers . The jobs arrive according to a Poisson process , and the job sizes follow a Pareto distribution . The two servers process jobs from their queues at identical rates . On each job arrival , the load balancer observes state st = ( q1 , q2 ) , denoting the queue length at the two servers . It then takes an action at ∈ { 1 , 2 } , sending the job to one of the servers . The goal of the load balancer is to minimize the average job completion time . The reward corresponding to this goal is rt = −τ × j , where τ is the time elapsed since the last action and j is total number of enqueued jobs . In this example , the optimal policy is to send the job to the server with the shortest queue ( Daley , 1987 ) . However , we find that a standard policy gradient algorithm , A2C ( Mnih et al. , 2016 ) , trained using a value function baseline struggles to learn this policy . The reason is that the stochastic sequence of job arrivals creates huge variance in the reward signal , making it difficult to distinguish between good and bad actions . Consider , for example , an action at the state shown in Figure 2a . If the arrival sequence following this action consists of a burst of large jobs ( e.g. , input sequence 1 in Figure 2a ) , the queues will build up , and the return will be poor compared to the value function baseline ( average return from the state ) . On the other hand , a light stream of jobs ( e.g. , input sequence 2 in Figure 2a ) will lead to short queues and a better-than-average return . Importantly , this difference in return has little to do with the action ; it is a consequence of the random job arrival process . We train two A2C agents ( Mnih et al. , 2016 ) , one with the standard value function baseline and the other with an input-dependent baseline tailored for each specific instantiation of the job arrival process ( details of this baseline in §4 ) . Since the the input-dependent baseline takes each input sequence into account explicitly , it reduces the variance of the policy gradient estimation much more effectively ( Figure 2b , left ) . As a result , even in this simple example , only the policy learned with the input-dependent baseline comes close to the optimal ( Figure 2b , right ) . Figure 2c visualizes the policies learned using the two baselines . The optimal policy ( pick-shortest-queue ) corresponds to a clear divide between the chosen servers at the diagonal . In fact , the variance of the standard baseline can be arbitrarily worse than an input-dependent baseline : we refer the reader to Appendix A for an analytical example on a 1D grid world .

Summary: This work considers the problem of learning in input-driven environments -- which are characterized by an addition stochastic variable z that can affect the dynamics of the environment and the associated reward the agent might see. The authors show how the PG theorem still applied for a input-aware critic and then they show that the best baseline one can use in conjecture with this critic is a input-dependent one. My main concerns are highlighted in points (3) and (4) in the detailed comments below.