Iterative Amortized Policy Optimization
Year: 2021

Abstract:
Policy networks are a central feature of deep reinforcement learning (RL) algorithms for continuous control, enabling the estimation and sampling of high-value actions. From the variational inference perspective on RL, policy networks, when used with entropy or KL regularization, are a form of amortized optimization, optimizing network parameters rather than the policy distributions directly. However, direct amortized mappings can yield suboptimal policy estimates and restricted distributions, limiting performance and exploration. Given this perspective, we consider the more flexible class of iterative amortized optimizers. We demonstrate that the resulting technique, iterative amortized policy optimization, yields performance improvements over direct amortization on benchmark continuous control tasks. Accompanying code: github.com/joelouismarino/variational_rl.

heading: 1 Introduction
text: 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 cannot 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.

heading: 2 Background

heading: 2.1 Preliminaries
text: 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 γ.

heading: 2.2 KL-Regularized Reinforcement Learning
text: 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].

heading: 2.3 Entropy & KL Regularized Policy Networks Perform Direct Amortization
text: 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.

heading: 2.4 Related Work
text: Previous works have investigated methods for improving policy optimization. QT-Opt [41] uses the cross-entropy method (CEM) [66], an iterative derivative-free optimizer, to optimize a Q-value estimator for robotic grasping. CEM and related methods are also used in model-based RL for performing model-predictive control [60, 14, 62, 33]. Gradient-based policy optimization [36, 71, 10], in contrast, is less common, however, gradient-based optimization can also be combined with CEM
3Additional 2D plots are shown in Figure 19 in the Appendix.
[5]. Most policy-based methods use direct amortization, either using a feedforward [31] or recurrent [28] network. Similar approaches have also been applied to model-based value estimates [13, 16, 4], as well as combining direct amortization with model predictive control [50] and planning [65]. A separate line of work has explored improving the policy distribution, using normalizing flows [29, 74] and latent variables [76]. In principle, iterative amortization can perform policy optimization in each of these settings.
Iterative amortized policy optimization is conceptually similar to negative feedback control [7], using errors to update policy estimates. However, while conventional feedback control methods are often restricted in their applicability, e.g., linear systems and quadratic cost, iterative amortization is generally applicable to any differentiable control objective. This is analogous to the generalization of Kalman filtering [42] to amortized filtering [54] for state estimation.

heading: 3 Iterative Amortized Policy Optimization

heading: 3.1 Formulation
text: Iterative amortization [55] utilizes errors or gradients to update the approximate posterior distribution parameters. While various forms exist, we consider gradient-encoding models [6] due to their generality. Compared with direct amortization (Eq. 10), we use iterative amortized optimizers of the general form
λ← fφ(s,λ,∇λJ ), (iterative amortization) (11) also shown in Figure 1 (Right), where fφ is a deep network and λ are the action distribution parameters. For example, if π = N (a;µ,diag(σ2)), then λ ≡ [µ,σ]. Technically, s is redundant, as the state dependence is already captured in J , but this can empirically improve performance [55]. In practice, the update is carried out using a “highway” gating operation [38, 72]. Denoting ωφ ∈ [0, 1] as the gate and δφ as the update, both of which are output by fφ, the gating operation is expressed as λ← ωφ λ+ (1− ωφ) δφ, (12) where denotes element-wise multiplication. This update is typically run for a fixed number of steps, and, as with a direct policy, the iterative optimizer is trained using stochastic gradient estimates of ∇φJ , obtained through the path-wise derivative estimator [44, 64, 35]. Because the gradients ∇λJ must be estimated online, i.e., during policy optimization, this scheme requires some way of estimating J online through a parameterized Q-value network [58] or a differentiable model [35].

heading: 3.2 Benefits of Iterative Amortization
text: Reduced Amortization Gap. Iterative amortized optimizers are more flexible than their direct counterparts, incorporating feedback from the objective during policy optimization (Algorithm 2), rather than only after optimization (Algorithm 1). Increased flexibility improves the accuracy of optimization, thereby tightening the variational bound [55, 54]. We see this flexibility in Figure 1 (Left), where an iterative amortized policy network iteratively refines the policy estimate (•), quickly arriving near the optimal estimate.
Multiple Estimates. Iterative amortization, by using stochastic gradients and random initialization, can traverse the optimization landscape. As with any iterative optimization scheme, this allows iterative amortization to obtain multiple valid estimates, referred to as “multi-stability” in the generative modeling literature [27]. We illustrate this capability across two action dimensions in Figure 2 for a state in the Ant-v2 MuJoCo environment. Over multiple policy optimization runs, iterative amortization finds multiple modes, sampling from two high-value regions of the action space. This provides increased flexibility in action exploration, despite only using a uni-modal policy distribution.
Generalization Across Objectives. Iterative amortization uses the gradients of the objective during optimization, i.e., feedback, allowing it to potentially generalize to new or updated objectives. We see this in Figure 1 (Left), where iterative amortization, despite being trained with a different value estimator, is capable of generalizing to this new objective. We demonstrate this capability further in Section 4. This opens the possibility of accurately and efficiently performing policy optimization in new settings, e.g., a rapidly changing model or new tasks.

heading: 3.3 Consideration: Mitigating Value Overestimation
text: Why are more powerful policy optimizers typically not used in practice? Part of the issue stems from value overestimation. Model-free approaches generally estimate Qπ using function approximation and temporal difference learning. However, this has the pitfall of value overestimation, i.e., positive bias in the estimate, Q̂π [75]. This issue is tied to uncertainty in the value estimate, though it is distinct from optimism under uncertainty. If the policy can exploit regions of high uncertainty, the resulting target values will introduce positive bias into the estimate. More flexible policy optimizers exacerbate the problem, exploiting this uncertainty to a greater degree. Further, a rapidly changing policy increases the difficulty of value estimation [63].
Various techniques have been proposed for mitigating value overestimation in deep RL. The most prominent technique, double deep Q-network [81] maintains two Q-value estimates [80], attempting to decouple policy optimization from value estimation. Fujimoto et al. [25] apply and improve upon this technique for actor-critic settings, estimating the target Q-value as the minimum of two Q-networks, Qψ1 and Qψ2 :
Q̂π(s,a) = min i=1,2 Qψ′i(s,a),
where ψ′i denotes the target network parameters. As noted by Fujimoto et al. [25], this not only counteracts value overestimation, but also penalizes high-variance value estimates, because the minimum decreases with the variance of the estimate. Ciosek et al. [15] noted that, for a bootstrapped ensemble of two Q-networks, the minimum operation can be interpreted as estimating
Q̂π(s,a) = µQ(s,a)− βσQ(s,a), (13) with mean µQ(s,a) ≡ 12 ∑ i=1,2Qψ′i(s,a), standard deviation σQ(s,a) ≡ ( 1 2 ∑ i=1,2(Qψ′i(s,a)− µQ(s,a)) 2)1/2, and β = 1. Thus, to further penalize high-variance value estimates, preventing value overestimation, we can increase β. For large β, however, value estimates become overly pessimistic, negatively impacting training. Thus, β reduces target value variance at the cost of increased bias.
Due to the flexibility of iterative amortization, the default β = 1 results in increased value bias and a more rapidly changing policy as compared with direct amortization (Figure 3). Further penalizing high-variance target values (β = 2.5) reduces value overestimation and improves stability. For details, see Appendix A.2. Recent techniques for mitigating overestimation have been proposed, such as adjusting α [22]. In offline RL, this issue has been tackled through the action prior [24, 48, 84] or by altering Qnetwork training [2, 49]. While such techniques could be used here, increasing β provides a simple solution with no additional computational overhead. This is a meaningful insight toward applying more powerful policy optimizers.

heading: 4 Experiments

heading: 4.1 Setup
text: To focus on policy optimization, we implement iterative amortized policy optimization using the soft actor-critic (SAC) setup described by Haarnoja et al. [32]. This uses two Q-networks, uniform action prior, pθ(a|s) = U(−1, 1), and a tuning scheme for the temperature, α. In our experiments, “direct” refers to direct amortization employed in SAC, i.e., a direct policy network, and “iterative” refers to iterative amortization. Both approaches use the same network architecture, adjusting only the number of inputs and outputs to accommodate gradients, current policy estimates, and gated updates (Sec. 3.1). Unless otherwise stated, we use 5 iterations per time step for iterative amortization, following [55]. For details, refer to Appendix A and Haarnoja et al. [31, 32].

heading: 4.2 Analysis

heading: 4.2.1 Visualizing Policy Optimization
text: We provide 2D visualizations of iterative amortized policy optimization in Figures 1 & 2, with additional 2D plots in Appendix B.5. In Figure 4, we visualize iterative refinement using a single action dimension from Ant-v2 across time steps. The refinements in Figure 4a give rise to the objective improvements in Figure 4b. We compare with Adam [45] (gradient-based) and CEM [66] (gradient-free) in Figure 4c, where iterative amortization is an order of magnitude more efficient. This trend is consistent across environments, as shown in Appendix B.4.

heading: 4.2.2 Performance Comparison
text: We evaluate iterative amortized policy optimization on the suite of MuJoCo [78] continuous control tasks from OpenAI gym [12]. In Figure 5, we compare the cumulative reward of direct and iterative amortized policy optimization across environments. Each curve shows the mean and ± standard deviation of 5 random seeds. In all cases, iterative amortized policy optimization matches or outperforms the baseline direct amortized method, both in sample efficiency and final performance. Iterative amortization also yields more consistent, lower variance performance.

heading: 4.2.3 Improved Exploration: Multiple Policy Modes
text: As described in Section 3.2, iterative amortization is capable of obtaining multiple estimates, i.e., multiple modes of the optimization objective. To confirm that iterative amortization has captured multiple modes, at the end of training, we take an iterative agent trained on Walker2d-v2 and histogram the distances between policy means across separate runs of policy optimization per state (Fig. 7a). For the state with the largest distance, we plot 2D projections of the optimization objective, J , across action dimensions in Figure 7b, as well as the policy density across 10 optimization runs (Fig. 7c). The multi-modal policy optimization surface shown in Figure 7b results in the multi-modal policy in Figure 7c. Additional results on other environments are presented in Appendix B.7.
To better understand whether the performance benefits of iterative amortization are coming purely from improved exploration via multiple modes, we also compare with direct amortization with a
multi-modal policy distribution. This is formed using inverse autoregressive flows [46], a type of normalizing flow (NF). Results are presented in Appendix B.2. Using a multi-modal policy reduces the performance deficiencies on Hopper-v2 and Walker2d-v2, indicating that much of the benefit of iterative amortization is due to lifting direct amortization’s restriction to a single, uni-modal policy estimate. Yet, direct + NF still struggles on HalfCheetah-v2 compared with iterative amortization, suggesting that more complex, multi-modal distributions are not the only consideration.

heading: 4.2.4 Improved Optimization: Amortization Gap
text: To evaluate policy optimization accuracy, we estimate per-step amortization gaps, performing additional iterations of gradient ascent on J w.r.t. the policy parameters, λ ≡ [µ,σ] (see Appendix A.3). To analyze generalization, we also evaluate the iterative agents trained with 5 iterations for an additional 5 amortized iterations. Results are shown in Figure 6. We emphasize that it is challenging to directly compare amortization gaps across optimization schemes, as these involve different value functions, and therefore different objectives. Likewise, we estimate the amortization gap using the learned Q-networks, which may be biased (Figure 3). Nevertheless, we find that iterative amortized policy optimization achieves, on average, lower amortization gaps than direct amortization across all environments. Additional amortized iterations at evaluation yield further estimated improvement, demonstrating generalization beyond the optimization horizon used during training.
The amortization gaps are small relative to the objective, playing a negligible role in evaluation performance (see Appendix B.6). Rather, improved policy optimization is helpful for training, allowing the agent to explore states where value estimates are highest. To probe this further, we train iterative amortized policy optimization while varying the number of iterations per step in {1, 2, 5}, yielding optimizers with varying degrees of accuracy. Note that each optimizer is, in theory, capable of finding multiple modes. In Figure 8, we see that training with additional iterations improves performance and optimization accuracy. We stress that the exact form of this relationship depends on the Q-value estimator and other factors. We present additional results in Appendix B.6.

heading: 4.2.5 Generalizing to Model-Based Value Estimates
text: Direct amortization is a purely feedforward process and is therefore incapable of generalizing to new objectives. In contrast, because iterative amortization is formulated through gradient-based feedback, such optimizers may be capable of generalizing to new objective estimators, as shown in Figure 1. To demonstrate this capability further, we apply iterative amortization with model-based value estimators, using a learned deterministic model on HalfCheetah-v2 (see Appendix A.5). We evaluate the generalizing capabilities in Figure 9 by transferring the policy optimizer from a modelfree agent to a model-based agent. Iterative amortization generalizes to these new value estimates, instantly recovering the performance of the model-based agent. This highlights the opportunity for instantly incorporating new tasks, goals, or model estimates into policy optimization.

heading: 5 Discussion
text: We have introduced iterative amortized policy optimization, a flexible and powerful policy optimization technique. In so doing, we have identified KL-regularized policy networks as a form of direct amortization, highlighting several limitations: 1) limited accuracy, as quantified by the amortization gap, 2) restriction to a single estimate, limiting exploration, and 3) inability to generalize to new objectives, limiting the transfer of these policy optimizers. As shown through our empirical analysis, iterative amortization provides a step toward improving each of these restrictions, with accompanying improvements in performance over direct amortization. Thus, iterative amortization can serve as a drop-in replacement and improvement over direct policy networks in deep RL.
This improvement, however, is accompanied by added challenges. As highlighted in Section 3.3, improved policy optimization can exacerbate issues in Q-value estimation stemming from positive bias. Note that this is not unique to iterative amortization, but applies broadly to any improved optimizer. We have provided a simple solution that involves adjusting a factor, β, to counteract this bias. Yet, we see this as an area for further investigation, perhaps drawing on insights from the offline RL community [49]. In addition to value estimation issues, iterative amortized policy optimization incurs computational costs that scale linearly with the number of iterations. This is in comparison with direct amortization, which has constant computational cost. Fortunately, unlike standard optimizers, iterative amortization adaptively tunes step sizes. Thus, relative improvements
rapidly diminish with each additional iteration, enabling accurate optimization with exceedingly few iterations. In practice, even a single iteration per time step can work surprisingly well.
Although we have discussed three separate limitations of direct amortization, these factors are highly interconnected. By broadening policy optimization to an iterative procedure, we automatically obtain a potentially more accurate and general policy optimizer, with the capability of obtaining multiple modes. While our analysis suggests that improved exploration resulting from multiple modes is the primary factor affecting performance, future work could tease out these effects further and assess the relative contributions of these improvements in additional environments. We are hopeful that iterative amortized policy optimization, by providing a more powerful, exploratory, and general optimizer, will enable a range of improved RL algorithms.
Acknowledgments and Disclosure of Funding
JM acknowledges Scott Fujimoto for helpful discussions. This work was funded in part by NSF #1918839 and Beyond Limits. JM is currently employed by Google DeepMind. The authors declare no other competing interests related to this work.

Review summaries:
The paper proposes to replace the actor/policy network with an iterative version to encode the action distribution parameters, which is inspired by prior work on iterative amortized optimization. This scheme generates the action distribution parameters for each state at the end of an inner loop which takes the objective gradient wrt to the parameters also as input (as opposed to a regular actor network which does one forward pass). This seems similar to prior work on learning to learn/meta learning but the presentation can be improved (more detailed comments below).  Experimental evidence is given for benefits over SAC in mujoco environments.

**Contributions**: The authors propose to use iterative amortization for policy optimization to help reduce suboptimality in policy optimization. They find that like in variational inference for generative modeling, iterative amortization is able cover multiple modes of distributions and lower the amortization gap, and show slightly improved performance on mujoco gym benchmarks. Additionally, they find that the iteratively amortized policy is able to better exploit the learned Q function and lead to more overestimation bias, and address this by tuning a parameter to make the Q-function backups more pessimistic to compensate.

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.

This work analyzes policy optimization from an inference perspective. It first shows that policy optimization with entropy or KL regularization in DRL with a policy network is a form of amortized optimization where instead of directly optimizing the action distribution parameters, we optimize a network which outputs the action distribution parameters. This results in an amortization gap which negatively impacts performance and exploration. Next the authors propose an alternative iterative amortization method for training the policy network. This new approach is shown to have various benefits which the authors demonstrated through a series of experiments on the MuJoCo environments.

The authors, via a standard reinterpretation of reinforcement learning as an inference problem, discusses the resulting “amortization gap” that results from using policy networks to optimize the variational lower bound. In order to close this gap, the authors propose to train an iterative procedure parameterized by $\phi$, which uses estimated gradients of the return to iteratively refine $\pi(\cdot | s_t)$ for the current state $s_t$. They then provide an extensive experimental evaluation demonstrating their proposed method and several desirable properties, including adaptability to gradients estimated by different methods.

This paper introduces a novel approach to policy optimization which leverages the connection between amortized variational inference and policy-based RL with function approximation. Specifically, rather than rely on a "direct" parameterization of the state-conditioned distribution over actions, in which an input state is mapped directly to distribution parameters (e.g., mean and variance), the approach adds an additional optimization procedure over the distribution parameters themselves, rather than those of the policy network. Experimental results are then used to demonstrate the benefits of this approach for performance, accuracy, and flexibility in the optimization process.