Boosting Trust Region Policy Optimization by Normalizing flows Policy

1 INTRODUCTION . In on-policy optimization , vanilla policy gradient algorithms suffer from occasional updates with large step size , which lead to collecting bad samples that the policy can not recover from ( Schulman et al. , 2015 ) . Motivated to overcome such instability , Trust Region Policy Optimization ( TRPO ) ( Schulman et al. , 2015 ) constraints the KL divergence between consecutive policies to achieve much more stable updates . However , with factorized Gaussian policy , such KL divergence constraint can put a very stringent restriction on the new policy , making it either hard to bypass locally optimal solutions or slow down the learning process . Can we improve the learning process of trust region policy search by using a more expressive policy class ? Intuitively , a more expressive policy class has more capacity to represent complex distributions and the KL constraint may not impose very strict restriction on the sample space . Though prior works ( Haarnoja et al. , 2017 ; 2018b ; a ) have proposed to use implicit generative models as policies , their focus is on off-policy learning . In this work , we show how normalizing flows can be combined with on-policy learning and boost the performance of trust region policy optimization . The structure of our paper is as follows . In Section 2 and 3 , we provide backgrounds on TRPO and related work . In Section 4 , we introduce normalizing flows for control and analyze why KL constraint may not impose a constraint on the sampled action space . On illustrative examples , we show that normalizing flows policy can learn policies with correlated actions and multi-modal policies , which allows for potentially more efficient exploration . In Section 5 , we show by comprehensive experiment results that normalizing flows significantly outperforms baseline policy classes when combined with trust region policy search algorithms . 2 BACKGROUND . 2.1 MARKOV DECISION PROCESS . In the standard formulation of Markov Decision Process ( MDP ) , at time step t ≥ 0 , an agent is in state st ∈ S , takes an action at ∈ A , receives an instant reward rt = r ( st , at ) ∈ R and transitions to a next state st+1 ∼ p ( ·|st , at ) ∈ S. Let π : S 7→ P ( A ) be a policy , where P ( A ) is a set of distribution over the action space A . The discounted cumulative reward under policy π is J ( π ) = Eπ [ ∞∑ t=0 γtrt ] , ( 1 ) where γ ∈ [ 0 , 1 ) is a discount factor . The objective of RL is to search for a policy π that achieves the maximum cumulative reward π∗ = arg maxπ J ( π ) . For convenience , we define action value function Qπ ( s , a ) = Eπ [ J ( π ) |s0 = s , a0 = a ] and value function V π ( s ) = Eπ [ J ( π ) |s0 = s , a0 ∼ π ( ·|s0 ) ] . We also define the advantage function Aπ ( s , a ) = Qπ ( s , a ) − V π ( s ) . 2.2 POLICY OPTIMIZATION . One way to search for π∗ is through direct policy search within a given policy class πθ , θ ∈ Θ where Θ is the parameter space for the policy parameter . We can update the paramter θ with policy gradient ascent , by computing ∇θJ ( πθ ) = Eπθ [ ∑∞ t=0A πθ ( st , at ) ∇θ log πθ ( at|st ) ] , then updating θnew ← θ + α∇θJ ( πθ ) for some learning rate α > 0 . Alternatively , the update can be formulated as trust region optimization problem max θnew Eπθ [ πθnew ( at|st ) πθ ( at|st ) Aπθ ( st , at ) ] , ||θnew − θ||2 ≤ , ( 2 ) for some > 0 . If we do a linear approximation of the objective in ( 2 ) , Eπθ [ πθnew ( at|st ) πθ ( at|st ) A πθ ( st , at ) ] ≈ Eπθ [ Aπθ ( st , at ) ] +∇θJ ( πθ ) T ( θnew − θ ) , we recover the gradient update by properly choosing . 2.3 TRUST REGION POLICY OPTIMIZATION . Trust Region Policy Optimization ( TRPO ) ( Schulman et al. , 2015 ) applies information theoretic constraints instead of Euclidean constraints on θnew and θ to better capture the geometry on the parameter space induced by the policy . In particular , consider the following trust region formulation max θnew Eπθ [ πθnew ( at|st ) πθ ( at|st ) Aπθ ( st , at ) ] , Es [ KL [ πθ ( ·|s ) ||πθnew ( ·|s ) ] ] ≤ , ( 3 ) where Es [ · ] is w.r.t . the state visit distribution induced by πθ . The trust region enforced by the KL divergence entails that the update according to ( 3 ) optimizes a lower bound of J ( πθ ) , so as to avoid accidentally taking large steps that irreversibly degrade the policy performance during training as in vanilla policy gradient ( 2 ) ( Schulman et al. , 2015 ) . For a practical algorithm , the trust region constraint is approximated by a second order expansion Es [ KL [ πθ ( ·|s ) ||πθnew ( ·|s ) ] ] ≈ ( θnew − θ ) T Ĥ ( θnew − θ ) ≤ where Ĥ = ∂ 2 ∂θ2Eπθ [ KL [ πθ ( ·|s ) ||πθnew ( ·|s ) ] ] is the expected Fisher information matrix . If we also linearly approximate the objective , the trust region formulation turns into a quadratic programming max θnew ∇θJ ( πθ ) T ( θnew − θ ) , ( θnew − θ ) T Ĥ ( θnew − θ ) ≤ . ( 4 ) The optimal solution to ( 4 ) is ∝ Ĥ−1∇θJ ( πθ ) . In cases where πθ is parameterized by a neural network with a large number of parameters , Ĥ−1 is formidable to compute . Instead , ( Schulman et al. , 2015 ) proposes to approximate Ĥ−1∇θJ ( πθ ) by conjugate gradient ( CG ) descent ( Wright & Nocedal , 1999 ) since it only requires relatively cheap Hessian-vector products . Given the approximated gradient direction ĝ ≈ Ĥ−1∇θJ ( πθ ) obtained from CG , the KL constraint is enforced by setting ∆θ = √ ĝT Ĥĝ ĝ . Finally a line search is carried out to determine a scaler s by enforcing the exact KL constraint Eπθ [ KL [ πθ+s∆θ||πθ ] ] ≤ and finally θnew ← θ + s∆θ . ACKTR ACKTR ( Wu et al . ) proposes to replace the above CG descent of TRPO by Kroneckerfactored approximation ( Martens & Grosse , 2015 ) when computing the inverse of Fisher information matrix Ĥ−1 . This approximation is more stable than CG descent and yields performance gain over conventional TRPO . 3 RELATED WORK . Most on-policy optimization algorithms are based on policy gradient theorem for function approximation ( Sutton et al. , 2000 ) . Vanilla policy gradient algorithms are typically more stable than off-policy learning due to their optimization based formulation , but can suffer from instability as a result of occasionally large step sizes . In policy based algorithms , updating policies with very large step sizes can be catastrophic to the learning process since the policy will collect bad samples and potentially never recover ( Schulman et al. , 2015 ) . Natural policy gradient ( Kakade , 2002 ) applies natural gradient for the policy updates , which accounts for the information geometry induced by the policy and makes the update more stable . More recently , Trust region policy optimization ( Schulman et al. , 2015 ) derives a tractable trust region policy search algorithm based on the lower bound formulation of ( Kakade & Langford , 2002 ) and achieves promising results on simulated locomotion tasks . The trust region is approximated by the Fisher information matrix , whose inverse is further approximated by conjugate gradient iterations ( Wright & Nocedal , 1999 ) . To further improve the scalability and numerical performance of TRPO , ACKTR ( Wu et al . ) applies Kronecker-factored approximation ( Martens & Grosse , 2015 ) to invert the Fisher information matrix . Orthogonal to prior works , we aim to improve TRPO with a more expressive policy representation , and we show significant improvements on both TRPO and ACKTR . We limit our attention to ( Dinh et al. , 2014 ) while other normalizing flows architectures might provide additional benefits ( Kingma & Dhariwal , 2018 ) . A number of recent prior works have proposed to boost RL algorithms with expressive policy classes . For off-policy learning , Soft Q-learning ( SQL ) ( Haarnoja et al. , 2017 ) takes an implicit generative model as the policy and trains the policy by Stein variational gradient ( Liu & Wang , 2016 ) . Similarly , ( Tang & Agrawal , 2018 ) applies an implicit policy along with a discriminator to compute entropy regularized gradient for the implicit distribution . Latent space policy ( Haarnoja et al. , 2018a ) applies normalizing flows as the policy and displays promising results on hierarchical tasks . Soft Actor Critic ( SAC ) applies a mixture of Gaussian as the policy . So far , expressive policy classes have shown improvement over baselines in the domain of off-policy learning . However , it is not clear whether such benefits come from an enriched policy class or a novel algorithmic procedure . In this work , we fix the trust region search algorithms and study the net effect of expressive policy classes . By definition , normalizing flows stacks layers of invertible transformations to map a source noise into target samples ( Dinh et al. , 2016 ; Rezende & Mohamed , 2015 ) . Through invertible transformations , normalizing flows retains tractable probability densities while being very expressive . Normalizing flows is widely applied in probabilistic generative modeling , such as variational inference ( Rezende & Mohamed , 2015 ) . Previous works have proposed to represent policies using normalizing flows in the context of off-policy learning ( Haarnoja et al. , 2018a ) . Complement to prior works , we show that normalizing flows can significantly boost the performance of on-policy optimization . 4 NORMALIZING FLOWS POLICY FOR ON-POLICY OPTIMIZATION . 4.1 NORMALIZING FLOWS FOR CONTROL . We construct a stochastic policy with normalizing flows . Normalizing flows ( Rezende & Mohamed , 2015 ; Dinh et al. , 2016 ) have been applied in variational inference and probabilistic modeling to represent complex distributions . In general , consider transforming a source noise ∼ ρ0 ( · ) by a series of invertible nonlinear functions gθi ( · ) , 1 ≤ i ≤ K each with parameter θi , to output a target sample x , x = gθK ◦ gθK−1 ◦ ... ◦ gθ2 ◦ gθ1 ( ) . ( 5 ) Let Σi be the inverse of the Jacobian matrix of gθ ( · ) , then the log density of x is computed by change of variables formula , log p ( x ) = log p ( ) + K∑ i=1 log det ( Σi ) . ( 6 ) For a general invertible transformation gθi ( · ) , computing det ( Σi ) is expensive . We follow the architecture of ( Dinh et al. , 2014 ) to ensure that det ( Σi ) is computed in linear time . To combine state information , we embed state s by another neural network Lθs ( · ) with parameter θs and output a state vector Lθs ( s ) with the same dimension as . We can then insert the state vector between any two layers of ( 5 ) to make the distribution conditional on state s. In our implementation , we insert the state vector after the first transformation ( we detail our architecture design in the Appendix B ) . a = gθK ◦ gθK−1 ◦ ... ◦ gθ2 ◦ ( Lθs ( s ) + gθ1 ( ) ) . ( 7 ) Though the additive form of Lθs ( s ) and gθ1 ( ) may in theory limit the capacity of the model , in experiments below we show that the resulting policy is still very expressive . For simplicity , we denote the above transformation ( 7 ) as a = fθ ( s , ) with parameter θ = { θs , θi , 1 ≤ i ≤ K } . It is obvious that the transformation a = fθ ( s , ) is still invertible between a and , which is critical for computing log πθ ( a|s ) according to ( 6 ) . Such representations build complex policy distributions with explicit probability density πθ ( ·|s ) , and hence entail training using score function gradient estimators . In on-policy optimizations , it is necessary to compute gradients of the entropy ∇θH [ πθ ( ·|s ) ] , either for computing Hessian vector product ( Schulman et al. , 2015 ) or for entropy regularization ( Schulman et al. , 2015 ; 2017 ; Mnih et al. , 2016 ) . For normalizing flows there is no analytic form for entropy , we use samples to estimate entropy by re-parameterization , H [ πθ ( ·|s ) ] = Ea∼πθ ( ·|s ) [ − log πθ ( a|s ) ] = E ∼ρ0 ( · ) [ − log πθ ( fθ ( s , ) |s ) ] . The gradient of the entropy can be easily computed by a pathwise gradient and easily implemented using back-propagation ∇θH [ πθ ( ·|s ) ] = E ∼ρ0 ( · ) [ −∇θ log πθ ( fθ ( s , ) |s ) ] .

Summary: The authors in this work present an approach to policy optimization that relies on an alternative policy formulation based on normalizing flows. This is a relatively simple modification (this is no criticism) that essentially uses the same TRPO algorithm as previous approaches, but a different mechanism for generating the distribution over actions. The crux of the authors’ approach is detailed in equations (6) and (7), although it could have been useful to see more of the discussion of the architecture from appendix B in the actual text of the paper.