Variance Reduced Domain Randomization for Policy Gradient

1 INTRODUCTION . Deep reinforcement learning ( DRL ) has achieved an impressive success on complex sequential decision-making tasks , such as playing video games at top human-level ( Ye et al. , 2020 ) , continuous robot manipulation ( Agarwal et al. , 2021 ) and traffic signal control ( Egeaa et al. , 2021 ) . For an extensive exploration , DRL requires a massive amount of samples to train the policy , which is usually done in a simulator constructed for the task . Due to the lack of diversity , however , policy trained by DRL tends to overfit to a specific training environment ( Cobbe et al. , 2019 ) , causing possibly a severe performance degradation ( i.e. , reality gap ) when the policy learned in simulator is transferred to the real deployment ( Peng et al. , 2018 ; Kang et al. , 2019 ) . To close this gap , domain randomization ( DR ) has been proposed to randomize the environment parameters that may vary the dynamics and observation in simulator , imposing diversity on the trained policy ( Tobin et al. , 2017 ; Jiang et al. , 2021 ) . Recent success in robot control has shown that DR enables the trained policy to generalize in real deployment , even with presence of fundamental modeling errors in the simulator ( Muratore et al. , 2021 ; Xie et al. , 2020 ) . DR has also demonstrated its capability of robustness to contend with extremely rare environments ( Muratore et al. , 2021 ) . Direct applications of DR on policy gradient methods can be thought as policy optimization on multiple randomly generated environments , which incurs additional variability of the observed data . Therefore , a critical challenge is the low sample efficiency due to the extremely high variance of the gradient estimator , where the variability not only stems from the gradient approximation of expected return from a batch of sampled trajectories ( Greensmith et al. , 2004 ) , but is also imposed by the randomization of environment parameters . In standard DRL , a commonly used method to reduce variance is to construct a bias-free and action-independent baseline that is subtracted from the expected return ( Sutton & Barto , 2018 ) . For example , a typical choice of baseline is the state value function of a policy . When directly applied to DR , as proposed by Andrychowicz et al . ( 2020 ) , it is equivalent to learning a state value function that predicts the expected return over all possible environments . Though remaining unbiased , such a state-dependent baseline may be a poor choice in DR , since the additional variability of randomized environments is not taken into account . In this paper , we aim to address the high variance issue of policy gradient methods for domain randomization , with a particular focus on reducing the additional variance imposed by randomization of environments . Our key insight is that the additional information on varying environments can be incorporated into the baseline to further reduce this variance . Theoretically , we derive the optimal state/environment-dependent baseline , and demonstrate that it improves consistently the variance reduction over baselines that is constant or uses state information only . In order for the practical implementation to strike the tradeoff between the variance reduction performance and computational complexity for maintaining the state/environment-dependent baselines , we propose a variance reduced domain randomization ( VRDR ) approach for policy gradient methods , which can improve the sample efficiency while maintaining a reasonable number of baselines associated with a set of specifically designed environment subspaces . Our main contributions can be summarized as follows . • Theoretically optimal state/environment-dependent baseline for DR. For the policy gradient in a variety of environments that differ in dynamics , we derive theoretically an optimal baseline that depends on both the states and environments . We further quantify the variance reduction improvement achieved by the proposed state/environment-dependent baseline over two common choices of baselines for DR , i.e. , the constant and state-dependent baselines . • Criterion for constructing practical state/subspace-dependent baseline . Since the accurate estimation of state/environment-dependent baseline for each possible environment is infeasible during practical implementation of RL , we propose to alternatively divide the entire space of environment parameters into a limited number of subspaces , and estimate instead the optimal baseline for every pair of state and environment subspace . We further show that the clustering of environments into subspaces should follow the policy ’ s expected returns on these environments , which can guarantee an improvement of variance reduction over the state-dependent baseline . • Variance reduced domain randomization ( VRDR ) with empirical evaluation . To strike the tradeoff between the variance reduction performance and computational complexity for maintaining optimal baselines , we develop a variance reduced domain randomization ( VRDR ) approach for policy gradient methods . Specifically , VRDR learns an acceptable number of baselines for each pair of state and environment subspace , where the environment subspaces are determined based on the above criterion . We then conduct experiments on six robot control tasks with their fundamental physical parameters randomized , demonstrating that VRDR can accelerate the convergence of policy training in DR settings , as compared to the standard state-dependent baselines . In some specific tasks , VRDR can even achieve a higher reward . 2 BACKGROUND . Notation . Under the standard reinforcement learning ( RL ) setting , the environment is modeled as a Markov decision process ( MDP ) defined by a tuple < S , A , Tp , Φ > , where S and A denote the state and action spaces , respectively . For the convenience of derivation , we assume that they are finite . Tp : S × A × S → [ 0 , 1 ] is the environment transition model that is essentially determined by environment parameter p ∈ P , with P denoting the space of environment parameters . In robot control , for example , environment parameter p can be a vector containing the rolling friction of each joint and the mass of torso and arms . Throughout the rest of this paper , by environment p we mean that the environment has dynamics determined by parameter p. Φ : S × A → R is the reward function . At each time step t , the agent observes its state st ∈ S and takes an action at ∈ A under the guidance of policy πθ ( at|st ) parameterized by θ . It will then receive a reward rt = Φ ( st , at ) , while the environment shifts to the next state st+1 with probability T ( st+1|st , at , p ) . The goal of standard RL is to search for a policy π that maximizes the expected discounted return η ( π , p ) = Eτ [ R ( τ ) ] over all possible trajectories τ = { st , at , rt , st+1 } ∞t=0 of states and actions , where R ( τ ) = ∑∞ t=0 γ trt and γ ∈ [ 0 , 1 ] is the discount factor . We can then define the state value function as Vπ ( s , p ) = E [ ∑∞ k=0 γ krt+k|st = s ] , the action value function asQπ ( s , a , p ) = E [ ∑∞ k=0 γ krt+k|st = s , at = a ] , and the advantage function as Aπ ( s , a , p ) = Qπ ( s , a , p ) − Vπ ( s , p ) . Policy gradient methods for DR . In DR , the environment parameter p is a random variable that follows the probability distribution P over P . For the convenience of derivation , we assume a finite en- vironment parameter space P with cardinality |P| . By introducing DR , the goal of policy optimization is to maximize the expected return over all possible environment parameters : Ep∼P [ η ( π , p ) ] . The policy gradient with action-independent baseline ( Sutton & Barto , 2018 ) can be formulated as : ∇θEp∼P [ η ( π , p ) ] = EP [ Eµpπ , π [ ∇θ log πθ ( a|s ) [ Qπ ( s , a , p ) − b ] ] ] , ( 1 ) where we define µpπ ( s ) = ∑∞ t=0 γ tPπ ( st = s|p ) as the discounted state visitation frequency , with Pπ ( st = s|p ) denoting the probability of shifting from the initial state s0 to state s after t steps under policy π in environment p. For convenience , we further denote g ( θ , s , a , p ) , ∇θ log πθ ( a|s ) [ Qπ ( s , a , p ) − b ] , which is the gradient estimator for the state-action pair under environment p. As long as the baseline b is action-independent , we have Ea [ ∇θ log πθ ( a|s ) b ] = ∇θEa [ b ] = 0 and thus EP , µpπ , π [ g ] = EP , µpπ , π [ ∇θ log πθ ( a|s ) Qπ ( s , a , p ) ] , E [ g ] , where the subscript is dropped for convenience . Therefore , by subtracting this action-independent baseline b , the variance of gradient estimator V ar ( g ) = E [ gTg ] − E [ g ] TE [ g ] can be reduced without introducing any bias ( Greensmith et al. , 2004 ) . 3 OPTIMAL BASELINES FOR DOMAIN RANDOMIZATION . To derive the optimal baselines for DR , we formulate the following optimization problem that aims to minimize the variance of gradient estimate w.r.t . the baseline b ( Greensmith et al. , 2004 ) : min b E [ G ( a , s ) [ Qπ ( s , a , p ) − b ] 2︸ ︷︷ ︸ gTg ] − E [ g ] TE [ g ] ⇔ min b EP , µpπ , π [ G ( a , s ) [ Qπ ( s , a , p ) − b ] 2 ] ( 2 ) where we denote G ( a , s ) , ∇θ log πθ ( a|s ) T∇θ log πθ ( a|s ) . Due to its independence of b , the second term E [ g ] TE [ g ] does not affect the minimizer , and is thus omitted on RHS of Eq . ( 2 ) . In the following , we first derive for DR two common choices of baselines that are constant or depend on the state only . We then propose the optimal state/environment-dependent baseline , and show its ability to further reduce the variance incurred by the randominzation of environment parameters . 3.1 TWO COMMON CHOICES OF ACTION-INDEPENDENT BASELINES . Optimal constant baseline . We first consider a constant baseline bc that depends neither on the action nor on the state . The optimization problem in Eq . ( 2 ) then becomes minimizing the expectation of a quadratic function , which can be proved to be convex . Referring to the detailed derivation in Appendix A.1 , the optimal constant baseline b∗c for DR-based policy gradient methods is : b∗c = EP , µpπ , π [ G ( a , s ) Qπ ( s , a , p ) ] EP , µpπ , π [ G ( a , s ) ] = EP , µpπ [ V ′ ( s , p ) ] . ( 3 ) This optimal baseline b∗c can be understood as the expected state value function V ′ ( s , p ) over all states and possible environments , where the state value function V ′ ( s , p ) is computed as the weighted average of the action value function : V ′ ( s , p ) = Eπ [ G ( a , s ) EP , µpπ , π [ G ( a , s ) ] Qπ ( s , a , p ) ] . Optimal state-dependent baseline . As in the standard RL , the utilization of state-dependent baseline in the DR setting can be considered as an expected value function b ( s ) = EP [ Eπ [ Qπ ( s , a , p ) ] ] for each state over all possible environments , which predicts the expected return over the distribution of dynamics caused by the variation of environment parameters . Still referring to Appendix A.1 , the optimal state-dependent baseline for DR is given as follow : b∗ ( s ) = EP ( p|s ) Eπ [ G ( a , s ) [ Qπ ( s , a , p ) ] ] Eπ [ G ( a , s ) ] . ( 4 )

Summary: This paper tackles the variance of policy gradient due to the domain randomization used in RL in simulations. The authors prove that the policy gradient variance can be further reduced by learning a state-dependent baseline for each environment parameter compared to state-dependent baselines. The authors then develop a practical algorithm based on the analysis and analyze the properties of the algorithm. The algorithm is implemented and tested on six robot control tasks. It consistently accelerates policy training.