Paper: Guaranteed Trust Region Optimization via Two-Phase KL Penalization (ICLR-2024)

Reviewer comment (Experiments):
The authors propose to add a new fixup phase to make the updated policy strictly satisfy trust region constraint, which can reduce the instability during training. However, we can easily address the instability issue by increasing the coefficient of KL regularization $\beta$, which is not only easier to implement, but also achieves better performances according to the ablation study (fig.5, $\beta=10$).

Author response:
It is true that setting $\beta = 10$ performs much better than expected, especially on the HalfCheetah benchmark. However, this does not guarantee the trust region is enforced without also using the fixup phase, which that ablation does use, since even $\beta = 10$ may not be sufficient to enforce the trust region. In some runs of other environments, $\beta$ reached much higher values than 10.

Reviewer comment (Evaluation):
The advantage of FixPO over previous methods is not consistent. In mujoco domain, FixPo does not exhibit better performance than baselines on halfcheetah, hopper and reacher tasks. There is also no significant different between FixPO and APPO in DMLab tasks. Meanwhile, FixPO cannot even exceed No Fixup phase baseline in fig.5, which makes the effectiveness of fixup phase very questionable.

Author response:
We do not claim that FixPO consistently achieves higher rewards than prior works. Our experiments are intended to demonstrate that FixPO is able to approximately match the rewards of proximal methods while providing stronger guarantees, and is able to match the guarantees of prior trust region methods while being more computationally efficient than them.

Reviewer comment (Writing):
The line below eq.(1), $L_{\pi_i}$ should be "policy loss" instead of "policy gradient".

Author response:
$L_{\pi_i}$ is the “importance sampled policy gradient loss”, that is described in Section 5 of [1]. We have edited the description to make this more clear.

Reviewer comment (Writing):
What is the exact definition of $D_{KL}^{max}[\pi_1,\pi_2]$?

Author response:
The exact definition of $D_{KL}^{max}[\pi_1, \pi_2]$ is $max_{s \in D} D_{KL}[\pi_1(a|s) || \pi_2(a|s)] = max_{s \in D} \int_a \pi_1(a|s) \frac {\pi_1(a|s)} {\pi_2(a|s)} \mathrm{d}a$

Reviewer comment (Experiments):
TRPO is a relevant baseline but is absent from experiments.

Author response:
Thank you for your kind and thoughtful comments.

We have revised the paper to include a TRPO baseline. It performs similarly to other baselines, so it does not significantly affect the interpretation of the results.

Reviewer comment (Theory):
Shouldn’t $D_{KL}$ be multiplied by $C_\beta$ in line 11 of Algorithm 1 (according to eq. 4)?

Author response:
> Shouldn’t $D_{KL}$ be multiplied by $C_\beta$ in line 11 of Algorithm 1?

This is a subtle but important aspect of the algorithm. By using $C_\beta$ in equation 4 and not using it on line 11, we effectively optimize a more conservative trust region in the primary phase minibatch steps than in the fixup phase minibatch steps.
This is important for two reasons.
First, because if the primary phase optimizes so that $D_{KL}^{max}$ is approximately at the trust-region boundary, it will often be slightly outside the trust-region boundary, requiring fixup gradient steps to correct it.
Second, because the minibatch is likely to produce a smaller $D_{KL}^{max}$ than the entire batch, we use $C_\beta$ in the primary phase minibatches so the KL computed there approximately matches the trust region constraint checked in the fixup loop.

If we were to use $C_\beta$ on line 11, it would be mathematically equivalent to changing the hyper parameters so that $\epsilon_{KL} \leftarrow \epsilon_{KL} / C_\beta$, $\gamma_\beta \leftarrow \gamma_\beta / C_\beta$, and $C_\beta \leftarrow 1$.
However, setting $C_\beta = 1$ performs very badly, as shown in the ablation.
In comparison, the extra conservativeness of the minibatch trust region does not seem to have a negative effect on reward, probably because the exact size of the trust region was largely chosen arbitrarily.

We have revised the paragraph after equation 4 to make this more clear.

Reviewer comment (Novelty):
Is the fixup phase analogous to line search in TRPO?

Author response:
> Is the fixup phase analogous to line search in TRPO?

The fixup phase performs a similar function to line search in TRPO, and in that sense they are definitely analogous. There are also several important differences. The most significant difference is that the fixup phase continues to use minibatches of gradient steps (SGD), while TRPO line search uses a projected gradient of the entire batch. For some environments, such as DMLab-30, the entire batch is very large (>50k samples), and consequently it is impractical to compute the full batch projected gradient.
The fixup phase also differs from the line search in TRPO through how the initial value step size is calculated. TRPO calculates a starting step size using the Hessian vector product, then shrinks that step using an exponential line search until the trust region is satisfied. FixPO checks to see if any gradient steps are necessary, and takes stochastic gradient steps to “decrease the step size” until the constraint is satisfied.

Reviewer comment (Experiments):
Can an entropy term be added to encourage exploration? I understand the point made in Fig. 4 but still wonder if additional exploration could be beneficial.

Author response:
> Can an entropy term be added to encourage exploration?

An entropy term can definitely be added to $L_\theta$. We briefly experimented with doing so, but did not use it in the final version so that we could focus more on other details in Section 3.

Reviewer comment (Experiments):
I find it interesting that in HalfCheetah, FixPO underperforms both PPO (fig. 3) and the constant $\beta$ ablation (fig. 5). I wonder if using a high constant $\beta$ (maybe without fixup) results in approximate trust region similar to what clipping does in PPO.

Author response:
> Question about HalfCheetah

We were surprised by the high performance of the $\beta = 10$ ablation on the HalfCheetah environment. When using $\beta = 10$ on HalfCheetah, the KL divergence remains very low and almost no fixup steps are required. We suspect that HalfCheetah benefits from having a smaller trust region than the other environments, but did not want to tune hyper parameters separately for different environments.

Reviewer comment (Novelty):
Just in case the authors would appreciate more related work on the Lagrangian optimization (where the Lagrange multiplier is learned), similar approaches are used in Adversarial (Inverse RL / Imitation Learning / Generative Networks) (https://arxiv.org/abs/1810.00821), Differentiable Economics (https://arxiv.org/abs/1706.03459, https://arxiv.org/abs/2202.13110), and Multi-Agent RL (https://arxiv.org/abs/2306.08419).

Author response:
> Additional citations for Lagrangian Methods

Thank you very much for providing us with these citations, we have incorporated them into the paper.

Reviewer comment (Theory):
I do not totally understand the intuition behind equation (4), which is perhaps the most important step in the proposed algorithm. Why it works and controls the distance between pi and pi'? It would be great if this result could be written as a proposition with proof.

Author response:
As you mention, equation 4 is one of the most important steps in the proposed algorithm. This equation results in $\beta$ acting as a Lagrange multiplier, which is a methodology also employed in prior works mentioned in the “Lagrangian Methods” paragraph of the related work section. An alternative explanation is to see $L_\beta$ as applying a linear feedback controller to the loss that approximately enforces the trust region constraint on a per minibatch basis. The use of momentum optimization for $\beta$ then causes this controller to be a PI (Proportional & Integral) controller. From this perspective, $C_\beta$ translates between a per-policy-step target $\epsilon_{KL}$ to a per-minibatch-step target $\epsilon_{KL} / C_\beta$. We describe why we expect this to result in convergence to the per-minibatch-step target in Section 3.2, but unfortunately we were unable to find a proof of this convergence in all cases.

Reviewer comment (Theory):
Could you mathematically describe the reasoning behind equation (4)?

Author response:
As you mention, equation 4 is one of the most important steps in the proposed algorithm. This equation results in $\beta$ acting as a Lagrange multiplier, which is a methodology also employed in prior works mentioned in the “Lagrangian Methods” paragraph of the related work section. An alternative explanation is to see $L_\beta$ as applying a linear feedback controller to the loss that approximately enforces the trust region constraint on a per minibatch basis. The use of momentum optimization for $\beta$ then causes this controller to be a PI (Proportional & Integral) controller. From this perspective, $C_\beta$ translates between a per-policy-step target $\epsilon_{KL}$ to a per-minibatch-step target $\epsilon_{KL} / C_\beta$. We describe why we expect this to result in convergence to the per-minibatch-step target in Section 3.2, but unfortunately we were unable to find a proof of this convergence in all cases.

Reviewer comment (Theory):
How do you compute the gradients of L_theta and L_beta?

Author response:
We compute gradients of our losses using the back propagation algorithm, as implemented in PyTorch.

Reviewer comment (Theory):
What is the complexity of the fixed up phase?

Author response:
Unfortunately, proving an upper bound on the per-iteration complexity of the proposed method would involve proving a finite-time convergence proof for SGD applied to deep neural networks, which is a major open problem in deep learning theory.
We cite the most recent progress we are aware of on this open problem, [1], in Section 3.2.
In practice, the average time spent in the fixup phase is quite low, as shown in Figure 1, and appears to be roughly constant throughout training for a specific choice of hyper parameters.
We do observe that occasionally a specific policy update out of several hundred can take significantly longer than others, however this never added more than 10 minutes to a 3 hour experiment run. As an optimization, it may be worth backtracking to a prior $\theta$ value if the fixup phase exceeds a time limit, but we found this to be unnecessary in practice.

Reviewer comment (Writing):
The paper does not provide context on the setting considered. Preliminaries are completely missing, the general tone is too colloquial, it is not clear which RL setting is considered (finite-horizon, infinite-horizon, average reward etc.).

Author response:
We have added a clarification to the description of each benchmark describing the RL setting for that benchmark. The Mujoco and DMLab-30 benchmarks are finite-horizon with discounted rewards and the Meta-World benchmark is infinite–horizon and uses discounted rewards with a success rate metric. These are the recommended settings for these benchmarks as used in prior work.

We appreciate the feedback that the tone is not formal enough. We have significantly revised the introduction and ablation sections to use more formal language and include more details.

Reviewer comment (Writing):
The advantage function $\hat{A}$ is never defined.

Author response:
We compute $\hat{A}$ using GAE[1]. We have added a citation to the second paragraph in Section 3.1 and to the Subroutines paragraph to clarify how $\hat{A}$ is computed.

Reviewer comment (Writing):
$L_\beta$ is introduced and not used later.

Author response:
We apologize, but we believe you are mistaken about $L_\beta$ not being used. $L_\beta$ is used in two important places (line 7 and line 14) in Algorithm 1. We also describe in detail how $L_\beta$ is used in Section 3.2.

Reviewer comment (Writing):
What is $L_{VF}$?

Author response:
$L_{VF}$ is described in the last paragraph of Section 3.1. We have not given a formal description of value function clipping because it is described in detail in prior works such as [2].

Reviewer comment (Writing):
How is it computed $L_\pi$?

Author response:
$L_\pi$ is described in Equation 3 to be the importance sampled policy gradient, $-\frac {\pi_{\theta}(a|s)} {\pi_{\theta'}(a|s)} \hat{A}$. We have added a note to clarify that this description is the definition.

Reviewer comment (Evaluation):
Figure 1: How can we extrapolate from the figure that the performances are reduced?

Author response:
Figure 1 is not intended to show the rewards the policy achieves for different values of $C_\beta$, and instead is intended to show stability in the number of gradient steps required throughout training. Figure 5 compares $C_\beta = 1$ to the default we chose ($C_\beta = 3$), which shows that $C_\beta = 1$ results in significantly decreased rewards. We have added a note to the description of Figure 1 referring to the $C_\beta = 1$ ablation. Figure 1 does show that $C_\beta = 1$ results in taking significantly more gradient steps. Gradient steps must be performed in sequence, and consequently taking a large number of them decreases the computational performance of Algorithm 1.

Reviewer comment (Evaluation):
Why are your results on Inverted Pendulum PPO-clip ~5000 when in [1] they achieved ~8000?

Author response:
We assume you mean Inverted Double Pendulum. The original PPO paper only reports rewards for InvertedDoublePendulum up to 1 million timesteps. We show PPO achieving similar performance (~8000) at that time in Figure 3. Continuing to train PPO-clip can lead to decreased rewards when not combined with KL penalization, as described in [5].

Reviewer comment (Experiments):
Why are your results on Swimmer PPO-clip ~50 when in [1] they achieved ~100?

Author response:
We were not able to reproduce the Swimmer results reported in the original PPO paper. Results on Swimmer are missing from most papers using PPO, including [3] and [4]. Where rewards on Simmer are reported in [5], they match our results of ~50.

Reviewer comment (Experiments):
The comparison with TRPO is missing (as with other policy optimization methods), although the two methods are quite similar.

Author response:
We have added a comparison to TRPO for OpenAI Gym. It is computationally difficult to add results for TRPO on Meta-World within the time limit of the revisions period, but we currently compare FixPO to [4] on this benchmark, which was published at ICLR and is reported to perform better than TRPO.

Reviewer comment (Experiments):
In general, the experimental evaluation is not convincing since the proposed method does not provide better results compared to PPO and the comparison with other policy optimization algorithms is missing.

Author response:
The primary objective of this work is not to achieve higher rewards than PPO, our experimental results are intended to show that we can maintain the same peak reward at similar computational cost while providing stronger guarantees.