Policy Gradient, Policy Iteration, and in between

Policy gradient (PG) and actor critic are the most trendy RL methods today. As a theorist, it often gives me headache when people bring up all sorts of variants of PG and ask me for opinions, and all I can say is …

Part of the reason is that many algorithms adopt changes to the original PG that look fairly heuristic, and I cannot find a conceptually simple and satisfactory framework for understanding and deriving them from first principles. This situation changed a little bit when Lihong recently asked me about PPO and shared some interesting empirical phenomenon, which got me thinking and organizing together some ideas that have been in the back of my mind for a long time. So this post will share some perspectives on how I view these algorithms.

---

Policy Gradient, Implemented (In)Correctly

Let’s start by reviewing the good-old PG, which is $\nabla J(\pi) = \frac{1}{1-\gamma}\mathbb{E}_{(s,a)\sim d^\pi}[Q^\pi(s,a) \nabla \log \pi(s,a)],$ and $Q^\pi(s,a)$ can be replaced by the advantage function $A^\pi(s,a) = Q^\pi(s,a) - V^\pi(s)$, or the Monte-Carlo return starting from $(s,a)$ under policy $\pi$.

In reality, however, a number of approximations is often deployed that departs from the equation above, with my comments at the end in parentheses:

1. The true $A^\pi$ is often unknown and we may use instead an approximate version $\hat{A}$ often learned via TD, which we refer to as the critic.¹ (Totally reasonable.)
2. In principle, after every update to the actor $\pi$, we should train a new critic for the new policy, ideally to convergence. But that’s obviously computationally too intensive, so we often just constantly update the critic, knowing that it likely has not converged to (the best approximation of) the advantage for the current policy. Theoretically, an updating scheme that reproduces the same issue (and is easier to study) is to freeze the critic and update the actor for a while, and then update the critic.
3. The equation requires on-policy sampling $(s,a)\sim d^\pi$. While most (but not all) practical variants still respect the action on-policyness, to my knowledge nobody cares about $s \sim d^\pi$.² Instead of $\mathbb{E}_{d^\pi}$ people often just write $\mathbb{E}_t$ to refer to a recent batch (RAGE METER BUILDING UP…)
4. Things get worse in off-policy setting, where an importance-weighting correction is needed. In practice, the policy on the denominator of importance weight does not always match the data used for updating the actor, which breaks importance sampling. (At this point man, what can I say.)
5. … and so on.

---

Policy Iteration

Despite my complaints about all these “wrong” practices, there is an angle that makes a lot of them justifiable. Consider the following: let’s push the alternating update in #2 to its extreme: in each round, we freeze $A^\pi$ and keep updating the actor $\pi$ until convergence. If the actor parameterization is sufficiently expressive, we would get

$$\pi_{\textrm{new}} \to \textrm{greedy}(Q^{\pi_{\textrm{old}}}), \qquad (1)$$

which is policy iteration (PI).³ Now, the textbook material on PI, viewed as a planning algorithm, just lifts so many restrictions of PG, including

• You can optimize the policy very far (all the way to greedy w.r.t. current critic) before updating critic.
• There is no need to be on-policy. In fact, the policy improvement step does not even talk about any notion of data distribution.

All these relaxations are nice. But they of course come at a price, otherwise we can simply use fitted PI and do not need PG at all. As I will explain below, my view is that PI relaxes the requirement on data distributions at the cost of strong (and sometimes unrealistic) assumptions on (1) policy structure, and (2) critic quality. We need strict PG if none of these assumptions hold. However, when some of them are reasonable, we may relax PG towards PI and enjoy the resulting flexibility.

---

On Policy Structure

In its original form, PI does not come with a separate actor and the policy is implicitly induced as the greedy policy of the value functions. For PI to be properly written in the form of Eq.(1), a key assumption is that

$$\textrm{greedy}(A^{\pi}) \in \Pi := \{\pi_\theta: \theta \in \Theta\},$$

that is, the greedy policy w.r.t. the advantage function (or its approximation $\hat{A}$ in the learning setting) can always be well-approximated by our policy parameterization $\Pi$. This is because the PD lemma tells us that

$$J(\pi) - J(\pi_{\textrm{old}}) = \frac{1}{1-\gamma}\mathbb{E}_{d^\pi}[A^{\pi_{\textrm{old}}}(s,\pi)],$$

where $A^{\pi_{\textrm{old}}}(s,\pi) = Q^{\pi_{\textrm{old}}}(s,\pi) - Q^{\pi_{\textrm{old}}}(s,\pi_{\textrm{old}})$. PI finds a new policy that is guaranteed to improve because the new policy has non-negative advantage over the old one in all states simultaneously, which is the key property that makes PI insensitive to data distributions. When we mimic PI using “critic-frozen-PG”, we can guarantee this property by choosing the tabular softmax policy, which has a factorized structure in the sense that the policy can choose the action distributions on each state independently.

In contrast, when we use a policy parameterization that entangles the action distributions across different states, we may run into the problem that the gradients of $A^\pi(s, \pi)$ for different $s$ may not be aligned with each other, and making improvements in some states may cause performance degradation in other states. In this case, PG theorem is the final safe net that guarantees overall improvement in $J(\pi)$: it tells us that the correct trade-off between the conflicting objectives at different states is precisely given by the distribution $d^\pi$.

---

On Critic Quality

Another major assumption behind PI is that $A^\pi$ is either given or its approximation $\hat{A}$ is universally good across the state-action space, which is rarely the case in practical learning settings. A much more reasonable assumption is that critic is only accurate under the distribution of states and actions induced by the current policy (see the seminal CPI paper, though there is more nuance to this assumption⁴), or the historical data, since this is where the critic is trained.⁵ Under this assumption, we cannot freeze the critic and optimize the actor indefinitely, as it will run into the regime where critic is poor due to the lack of data. Instead, we need to stop after $\pi$ is a bit far from $\pi_{\textrm{old}}$, collect data, and re-train the critic to make sure the actor never gets into the completely uncharted regime and stays within the trust region.

---

Understanding the Practical Relaxations

The takeaway message of the above discussion is that PI makes a number of implicit assumptions, which allow it to have much less restriction than PG in several corresponding aspects. Judging from how people tweak PG in practice (as mentioned at the beginning), we can “reverse-engineer” and have some reasonable guess of why those tweaks are successful: given that people care less about on-policyness but enforce trust region on actor updates pretty carefully, we can conjecture that (1) modern neural-net architectures are rich enough that conflicting objectives across states are less of an issue (especially when actor and critic share layers),⁶ and (2) critic is still only good under existing data, which necessitates trust region. (Update 09/30) After I show this post to Wei Xiong, he told me that there are recent works that start to take distribution issues in PG more seriously.⁷ I guess we will just have to wait and see how things play out.

Footnote Remarks