Iterative Amortized Policy Optimization

1 Introduction . Reinforcement learning ( RL ) algorithms involve policy evaluation and policy optimization [ 73 ] . Given a policy , one can estimate the value for each state or state-action pair following that policy , and given a value estimate , one can improve the policy to maximize the value . This latter procedure , policy optimization , can be challenging in continuous control due to instability and poor asymptotic performance . In deep RL , where policies over continuous actions are often parameterized by deep networks , such issues are typically tackled using regularization from previous policies [ 67 , 68 ] or by maximizing policy entropy [ 57 , 23 ] . These techniques can be interpreted as variational inference [ 51 ] , using optimization to infer a policy that yields high expected return while satisfying prior policy constraints . This smooths the optimization landscape , improving stability and performance [ 3 ] . However , one subtlety arises : when used with entropy or KL regularization , policy networks perform amortized optimization [ 26 ] . That is , rather than optimizing the action distribution , e.g. , mean and variance , many deep RL algorithms , such as soft actor-critic ( SAC ) [ 31 , 32 ] , instead optimize a network to output these parameters , learning to optimize the policy . Typically , this is implemented as a direct mapping from states to action distribution parameters . While such direct amortization schemes have improved the efficiency of variational inference as “ encoder ” networks [ 44 , 64 , 56 ] , they also suffer from several drawbacks : 1 ) they tend to provide suboptimal estimates [ 20 , 43 , 55 ] , yielding a so-called “ amortization gap ” in performance [ 20 ] , 2 ) they are restricted to a single estimate [ 27 ] , thereby limiting exploration , and 3 ) they can not generalize to new objectives , unlike , e.g. , gradient-based [ 36 ] or gradient-free optimizers [ 66 ] . Inspired by techniques and improvements from variational inference , we investigate iterative amortized policy optimization . Iterative amortization [ 55 ] uses gradients or errors to iteratively update the parameters of a distribution . Unlike direct amortization , which receives gradients only after ∗Now at DeepMind , London , UK . Correspondence to josephmarino @ deepmind.com . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . outputting the distribution , iterative amortization uses these gradients online , thereby learning to iteratively optimize . In generative modeling settings , iterative amortization empirically outperforms direct amortization [ 55 , 54 ] and can find multiple modes of the optimization landscape [ 27 ] . The contributions of this paper are as follows : • We propose iterative amortized policy optimization , exploiting a new , fruitful connection between amortized variational inference and policy optimization . • Using the suite of MuJoCo environments [ 78 , 12 ] , we demonstrate performance improvements over direct amortized policies , as well as more complex flow-based policies . • We demonstrate novel benefits of this amortization technique : improved accuracy , providing multiple policy estimates , and generalizing to new objectives . 2 Background . 2.1 Preliminaries . We consider Markov decision processes ( MDPs ) , where st ∈ S and at ∈ A are the state and action at time t , resulting in reward rt = r ( st , at ) . Environment state transitions are given by st+1 ∼ penv ( st+1|st , at ) , and the agent is defined by a parametric distribution , pθ ( at|st ) , with parameters θ.2 The discounted sum of rewards is denoted as R ( τ ) = ∑t γtrt , where γ ∈ ( 0 , 1 ] is the discount factor , and τ = ( s1 , a1 , . . . ) is a trajectory . The distribution over trajectories is : p ( τ ) = ρ ( s1 ) T−1∏ t=1 penv ( st+1|st , at ) pθ ( at|st ) , ( 1 ) where the initial state is drawn from the distribution ρ ( s1 ) . The standard RL objective consists of maximizing the expected discounted return , Ep ( τ ) [ R ( τ ) ] . For convenience of presentation , we use the undiscounted setting ( γ = 1 ) , though the formulation can be applied with any valid γ . 2.2 KL-Regularized Reinforcement Learning . Various works have formulated RL , planning , and control problems in terms of probabilistic inference [ 21 , 8 , 79 , 77 , 11 , 51 ] . These approaches consider the agent-environment interaction as a graphical model , then convert reward maximization into maximum marginal likelihood estimation , learning and inferring a policy that results in maximal reward . This conversion is accomplished by introducing one or more binary observed variables [ 19 ] , denoted as O for “ optimality ” [ 51 ] , with p ( O = 1|τ ) ∝ exp ( R ( τ ) /α ) , where α is a temperature hyper-parameter . We would like to infer latent variables , τ , and learn parameters , θ , that yield the maximum log-likelihood of optimality , i.e. , log p ( O = 1 ) . Evaluating this likelihood requires marginalizing the joint distribution , p ( O = 1 ) = ∫ p ( τ , O = 1 ) dτ . This involves averaging over all trajectories , which is intractable in high-dimensional spaces . Instead , we can use variational inference to lower bound this objective , introducing a structured approximate posterior distribution : π ( τ |O ) = ρ ( s1 ) T−1∏ t=1 penv ( st+1|st , at ) π ( at|st , O ) . ( 2 ) This provides the following lower bound on the objective : log p ( O = 1 ) = log ∫ p ( O = 1|τ ) p ( τ ) dτ ( 3 ) ≥ ∫ π ( τ |O ) [ log p ( O = 1|τ ) p ( τ ) π ( τ |O ) ] dτ ( 4 ) = Eπ [ R ( τ ) /α ] −DKL ( π ( τ |O ) ‖p ( τ ) ) . ( 5 ) 2In this paper , we consider the entropy-regularized case , where pθ ( at|st ) = U ( −1 , 1 ) , i.e. , uniform . However , we present the derivation for the KL-regularized case for full generality . 21 Equivalently , we can multiply by α , defining the variational RL objective as : J ( π , θ ) ≡ Eπ [ R ( τ ) ] − αDKL ( π ( τ |O ) ‖p ( τ ) ) . ( 6 ) This objective consists of the expected return ( i.e. , the standard RL objective ) and a KL divergence between π ( τ |O ) and p ( τ ) . In terms of states and actions , this objective is : J ( π , θ ) = Est , rt∼penv at∼π [ T∑ t=1 rt − α log π ( at|st , O ) pθ ( at|st ) ] . ( 7 ) At a given timestep , t , one can optimize this objective by estimating the future terms in the sum using a “ soft ” action-value ( Qπ ) network [ 30 ] or model [ 62 ] . For instance , sampling st ∼ penv , slightly abusing notation , we can write the objective at time t as : J ( π , θ ) = Eπ [ Qπ ( st , at ) ] − αDKL ( π ( at|st , O ) ||pθ ( at|st ) ) . ( 8 ) Policy optimization in the KL-regularized setting corresponds to maximizing J w.r.t . π . We often consider parametric policies , in which π is defined by distribution parameters , λ , e.g. , Gaussian mean , µ , and variance , σ2 . In this case , policy optimization corresponds to maximizing : λ← arg max λ J ( π , θ ) . ( 9 ) Optionally , we can then also learn the policy prior parameters , θ [ 1 ] . 2.3 Entropy & KL Regularized Policy Networks Perform Direct Amortization . Policy-based approaches to RL typically do not directly optimize the action distribution parameters , e.g. , through gradient-based optimization . Instead , the action distribution parameters are output by a function approximator ( deep network ) , fφ , which is trained using deterministic [ 70 , 52 ] or stochastic gradients [ 83 , 35 ] . When combined with entropy or KL regularization , this policy network is a form of amortized optimization [ 26 ] , learning to estimate policies . Again , denoting the action distribution parameters , e.g. , mean and variance , as λ , for a given state , s , we can express this direct mapping as λ← fφ ( s ) , ( direct amortization ) ( 10 ) denoting the corresponding policy as πφ ( a|s , O ; λ ) . Thus , fφ attempts to learn to optimize Eq . 9 . This setup is shown in Figure 1 ( Right ) . Without entropy or KL regularization , i.e . πφ ( a|s ) = pθ ( a|s ) , we can instead interpret the network as directly integrating the LHS of Eq . 4 , which is less efficient and more challenging . Regularization smooths the optimization landscape , yielding more stable improvement and higher asymptotic performance [ 3 ] . Viewing policy networks as a form of direct amortized variational optimizer ( Eq . 10 ) allows us to see that they are similar to “ encoder ” networks in variational autoencoders ( VAEs ) [ 44 , 64 ] . However , there are several drawbacks to direct amortization . Algorithm 1 Direct Amortization Initialize φ for each environment step do λ← fφ ( st ) at ∼ πφ ( at|st , O ; λ ) st+1 ∼ penv ( st+1|st , at ) end for for each training step do φ← φ+ η∇φJ end for Algorithm 2 Iterative Amortization Initialize φ for each environment step do Initialize λ for each policy optimization iteration do λ← fφ ( st , λ , ∇λJ ) end for at ∼ πφ ( at|st , O ; λ ) st+1 ∼ penv ( st+1|st , at ) end for for each training step do φ← φ+ η∇φJ end for Amortization Gap . Direct amortization results in suboptimal approximate posterior estimates , with the resulting gap in the variational bound referred to as the amortization gap [ 20 ] . Thus , in the RL setting , an amortized policy , πφ , results in worse performance than the optimal policy within the parametric policy class , denoted as π̂ . The amortization gap is the gap in following inequality : J ( πφ , θ ) ≤ J ( π̂ , θ ) . Because J is a variational bound on the RL objective , i.e. , expected return , a looser bound , due to amortization , prevents one from more completely optimizing this objective . This is shown in Figure 1 ( Left ) ,3 where J is plotted over two dimensions of the policy mean at a particular state in the MuJoCo environment Hopper-v2 . The estimate of a direct amortized policy ( ) is suboptimal , far from the optimal estimate ( F ) . While the relative difference in the objective is relatively small , suboptimal estimates prevent sampling and exploring high-value regions of the action-space . That is , suboptimal estimates have only a minor impact on evaluation performance ( see Appendix B.6 ) but hinder effective data collection . Single Estimate . Direct amortization is limited to a single , static estimate . In other words , if there are multiple high-value regions of the action-space , a uni-modal ( e.g. , Gaussian ) direct amortized policy is restricted to only one region , thereby limiting exploration . Note that this is an additional restriction beyond simply considering uni-modal distributions , as a generic optimization procedure may arrive at multiple uni-modal estimates depending on initialization and stochastic sampling ( see Section 3.2 ) . While multi-modal distributions reduce the severity of this restriction [ 74 , 29 ] , the other limitations of direct amortization still persist . Inability to Generalize Across Objectives . Direct amortization is a feedforward procedure , receiving gradients from the objective only after estimation . This is contrast to other forms of optimization , which receive gradients ( feedback ) during estimation . Thus , unlike other optimizers , direct amortization is incapable of generalizing to new objectives , e.g. , if Qπ ( s , a ) or pθ ( a|s ) change , which is a desirable capability for adapting to new tasks or environments . To improve upon this scheme and overcome these drawbacks , in Section 3 , we turn to a technique developed in generative modeling , iterative amortization [ 55 ] , retaining the efficiency of amortization while employing a more flexible iterative estimation procedure .

Summary: This paper draws an interesting connection to variational inference, categorizing current policy optimization methods with KL regularization as direct amortized optimizers. The paper shows how direct amortized policy optimization can be suboptimal and proposes a new class of method called iterative amortized policy optimization. This ‘iterative’ amortized policy optimization performs iterative optimization before the regular optimization step, offering several advantages: better at reaching the optima, better at sampling multiple modes, and support more rapidly changing policy. The paper shows nice visualization to support these claims and also run benchmark experiments to show its improvement.