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 high variance problem caused by the randomization of environments for estimating the policy gradients. The idea is to derive a bias-free and state/environment-dependent optimal baseline for domain randomization. The authors further develop a VRDR method by dividing the entire environment space into subspaces and estimating the state/subspace-dependent baseline.