Numerical Pitfalls in Policy Gradient Updates

Thanks for the detailed feedback, we hope that the following response can address your concerns.

For the weaknesses:

The issue of numerical instability due to importance sampling is well recognized, and several of the potential solutions explored here have also been discussed in prior work.

We agree that the issue of numerical instability due to importance sampling has been studied in prior work, such as [1], as discussed in the Related Work section. Our work differs in that, while previous studies highlighted how the product of density ratios over a long horizon can grow exponentially, leading to exploding variance in off-policy estimation, we demonstrate that even a single density ratio can take excessively large values and cause numerical overflow.

To the best of our knowledge, no existing work explicitly relates arithmetic overflow in deep policy gradients to the numerical instability of importance sampling. Currently, this numerical issue is generally viewed as a hyperparameter tuning problem, as suggested in the Stable-Baselines3 documentation (https://stable-baselines3.readthedocs.io/en/master/guide/checking_nan.html). Our work contributes to this line of research by directly identifying the ill-conditioned nature of importance sampling as the root cause of the problem.

Please also refer to the global response, where we summarize the contributions of our work.

W1: Our analysis focuses on Gaussian policies because they are the default in policy gradient methods and most theoretical frameworks and algorithms are designed for this case. However, the discussion on the ill-conditionedness of importance sampling is not limited to Gaussian policies and can be applied to various probability distributions.

W2: We agree that using more environments could further strengthen the analysis. However, we believe that the three environments (Hopper, Walker, Humanoid) considered in the paper are sufficient to support the theoretical analysis, as they were adopted from two well-known prior works [2, 3].

W3: In Figure 1(a), it shows how the condition number $\kappa(\theta)$ changes with respect to the mean of $\pi$ ($\mu_0$), and the difference between the mean of $\pi_\theta$ ($\mu$) and $\mu_0$. The condition number is related to the probability ratio $p_\theta$. These plots indicate that the condition number of the importance sampling ratio grows exponentially as $\pi_\theta$ deviates from $\pi$, potentially resulting in arithmetic overflow. We thank the reviewer for pointing it out and now we have fixed it.

References:

[1] "Breaking the curse of horizon: Infinite-horizon off-policy estimation" by Liu et al., NeurIPS, 2018;

[2] "Implementation Matters in Deep Policy Gradients: A Case Study on PPO and TRPO" by Engstrom et al., ICLR, 2019;

[3] "A closer look at deep policy gradients" by Ilyas et al., ICLR, 2020;

For the questions:

Q1: As we mentioned, the importance sampling ratio can become ill-conditioned when the probability density function takes excessively small values. This can occur with beta distributions, which have zero or infinite density at the boundaries. Therefore, importance sampling remains the root cause of exploding gradients, as it may involve dividing by extremely small values.

Q2: They are related as follows: the importance sampling ratio can grow exponentially, leading to arithmetic overflow. When this occurs, the objective value to be differentiated becomes NaN, which corrupts the entire computational graph of auto-differentiation, thereby causing exploding gradient issues. This is why techniques like gradient clipping cannot prevent the issue—because the overflow happens before the gradient computation.

Q3: We agree that it would ensure the numerical stability for PPO. One potential problem is that when both $\pi_\theta$ and $\pi$ are very small, the probability ratio will always take very small values and hence affect the clipping law in PPO (which assumes that the ratio stays near 1). We proposed a method in the revision that replaces the importance sampling ratio with its logarithm, and it works out.

Q4: Yes, that is why we said "assume that the mean $\mu(s)$ is SUFFICIENTLY large".

Q5: The network architecture used for the value function is described in Appendix A. While using ReLU leads to non-smooth functions, they are at least piecewise linear. However, the true value function, as shown there, can have fractal structures and be much more non-smooth than ReLU functions.

Q6: In Figure 4 (now in the appendix), the y-axis in (a) represents the cumulative reward. The constant lines in (b) are due to the convergence of the algorithm.

We also thank the reviewer for pointing out the notation issues, which have been fixed in the revised draft.

Dear Reviewer uwTa,

We are delighted to hear that the revised draft has addressed your concerns. Thank you for raising the score and for supporting our work!

Thanks for the positive review. Please see the following responses to your concerns.

For the weaknesses:

FINAL RETURN is not a fully informative comparison measure. The hyperparameter tuning process in your experiments is not clearly explained, and five runs seem insufficient for reliable results. Increasing this to ten runs would provide more robust data. Additionally, were these runs conducted with different random seeds?

We agree that using more random seeds would yield more robust data. Our experimental setup is primarily based on [1], which used five random seeds for each environment. Similarly, [2] also employed five seeds, balancing the computation budget with the need for robustness.

Regarding the performance metric, we reported both FINAL RETURN and MAXIMAL RETURN when they differ significantly, and only FINAL RETURN when it is equal to (or nearly equal to) the MAXIMAL RETURN.

For the hyperparameter tuning process, we followed the procedure outlined in Appendix A, which was adopted from previous works.

In all experiments presented in Section 5 and 6, we fixed the random seeds (42–46) to ensure consistency and comparability of the results. This approach was also applied to the additional experiments detailed in Appendix B (now in Section 4).

I would have given a strong recommendation if the paper had introduced a rigorous algorithm to bypass the need for importance sampling in TRPO/PPO.

That’s a great point! We just proposed a new solution called "vanilla policy gradient with clipping," which removes the importance sampling ratio from the PPO objective and applies the same clipping rule to the vanilla policy gradient objective. This algorithm combines the effectiveness of PPO with the robustness of VPG. The paper you suggested is highly relevant to the topic and we have discussed it in Related Work. Please refer to the global response and the revised draft for further details.

For the question:

Why did you use FINAL RETURN to report? why the final is important? not previous ones?

In response to your question, both [1] and [2] used the final return as the performance metric. We reported both maximal return and final return when they differed, and only the final return when they were nearly the same.

References:

[1] "Implementation Matters in Deep Policy Gradients: A Case Study on PPO and TRPO" by Engstrom et al., ICLR, 2019;

[2] "Deep Reinforcement Learning that Matters" by Henderson et al., AAAI, 2017.

Thank you for the feedback and please see the following responses to your questions.

For the weaknesses:

The authors claim that the success of policy gradient algorithms is due to numerical instability -- this seems like an odd thing to infer from their ablations.

In the revised draft, we have adjusted the wording to better emphasize the point that the importance sampling term may not be necessary in policy gradient methods.

While the point made in L137-140 about sampling from $\pi$ notwithstanding, clipping them to a bounded range leading to an unintentionally small value for pi, this doesn't seem like a strong enough case to justify the broad claim in the title in the section, "importance sampling considered harmful".

Thank you for pointing that out. We have rewritten the section title accordingly.

Most of the insights contained in the paper are either not entirely novel or they are hard to be convinced of (e.g. claims about importance sampling being harmful, or numerical instability benefitting deep pg algorithms), from the empirical observations. There isn't much theory to speak of, and the main theorem is just a simple statement about the standard gaussian normal density.

We have added more empirical analysis to the revised draft and proposed an algorithm that combines the merits of vanilla policy gradient and PPO. This algorithm effectively overcomes the numerical issues while achieving performance comparable to PPO. We believe that the contributions of this work are now clearly presented. Controversial statements have also been removed.

For the questions:

Figure 2 unclear what is the experiment ablation being studied across the three variants (a), (b) and (c).

In Figure 2 (now Figure 3), we empirically demonstrate that optimization constraints, such as clipping and KL penalties, are not sufficient to prevent the importance sampling ratio from exploding. The difference between these plots lies in the constraints applied to the policy training: (a) probability ratio clipping, (b) KL-penalty, and (c) no constraints. We have clarified this point in the revised draft.

In Table 2, can you describe "explosion rate" a bit more and how this impacts the learning? Is this the fraction of samples for which the gradients are NaN and in which case the sample is effectively discarded from the mini batch that computes the model update? Or is it the fraction of minibatches where the update is not computed, i.e. effectively a full training step wasted.

In Table 2, the explosion rate represents the percentage of individual runs that encounter arithmetic overflow and return NaN. For example, if 2 out of 5 random seeds encounter gradient explosion, the explosion rate would be 40%.

We are happy to answer any follow-up questions.

Thank you for the detailed feedback. Please see the following response for the concerns.

Regarding the weaknesses:

My main concern stems from the contrived analysis.

It is really a good question! Let us clarify these points:

The authors state that 1) can be controlled (eg, entropy bonus).

When the covariance of the policy has small eigenvalues, the importance sampling ratio becomes ill-conditioned. This issue can be controlled by setting lower bounds on the standard deviations.

They argue the main issue stems from 2), but I think it is unlikely unless the policy’s initialization is pathological.

As mentioned in (3), the policy distribution can shift outside the action space during training due to mini-batch updates, even if the policy network is well-initialized. Therefore, the initialization itself is less important compared to the shift caused during policy training. This is not a contrived scenario, but rather the general case during PPO training. This phenomenon has been reported in previous work ([1]), though it was not discussed in relation to numerical issues. We have revised the wording to clarify this point.

I think this is directly targetted by the KL constrained.

In Section 4 (now Section 5), we showed that the KL divergence is not able to fully control the mean from leaving the action space, since the KL divergence measures the averaged distance between two distributions over the entire space, whereas the probability ratio can take large values at specific points. In Figure 2 (now Figure 3), we observe that while the KL distance remains small throughout, the maximum of the probability ratios explodes.

A second concern is that the authors study potential fixes but none works.

Since we have observed that vanilla policy gradient is numerically more robust compared to PPO, we propose a new solution called "vanilla policy gradient with clipping." This method removes the importance sampling ratio from the PPO objective and applies the same clipping rule to the vanilla policy gradient objective. The algorithm combines the effectiveness of PPO with the robustness of VPG. Please refer to the global response and the revised draft for further details.

In fact, I believe there is a well-known solution: for bounded action spaces, use the so called “Tanh-Gaussian” distribution as done in the Soft Actor-Critic and derivative works (Haarnoja et al, 2018).

This approach still suffers from the ill-conditioned nature of importance sampling. For instance, consider the modified probability density function in the one-dimensional case, given by $\pi'(a|s) = \pi(u|s) (1 - \tanh^2(u))^{-1}$ where $u \sim \pi$ is a Gaussian random variable and $a = \tanh(u)$. Note that when the mean of $\pi$ is too large, it is likely that the random variable $u$ will take on extreme values, resulting in an excessively small value for the term $1 - \tanh^2(u)$ and causing numerical overflow in subsequent steps. Thanks for the comment and we have added this discussion to the revised draft.

In fact, I’m doubtful of some of the experimental results. SAC with the Tanh-Gaussian gets among the best results on the Mujoco tasks studied in this paper but the authors report negative results when trying this distribution with TRPO / PPO.

In fact, the approach we attempted in the "Mean transformation" section is to directly bound the mean of the Gaussian distribution, rather than its output. Therefore, these are different approaches.

I don’t know what’s the difference between the plots in Fig. 2.

In Figure 2 (now Figure 3), we empirically show that optimization constraints such as clipping and KL penalties are not sufficient to prevent the importance sampling ratio from exploding. The difference between these plots lies in the constraints applied to the policy training: (a) probability ratio clipping, (b) KL-penalty, and (c) no constraints. We have clarified this point in the revision, and thank you for pointing it out.

is max computed over all optimization steps or just the current one

Just the current step.

for Table 1 it would be good to include the reference results without any clipping or tanh.

Now we have included the reference results without any clipping or tanh in Table 1.

For the questions:

Why would the policy distribution drift significantly outside the box defined by the action space?

It is very likely that the policy distribution shifts significantly outside the action space, as PPO performs multiple policy updates in each epoch due to mini-batching. It has been reported that PPO training often leads to policies that generate boundary actions ([1]).

l. 357: Why would $\mu$ be outside the action space at initialization?

In line 357, we mentioned that when the policy network is large, the variance of its output is also large at initialization, which implies that some actions may already be outside the action space. While this is not necessarily always the case, it is more likely to occur with larger networks.

Do the same issues appear on discrete action spaces?

That’s a good question. Although it is somewhat beyond the scope of this work, we believe that numerical instability could also arise in discrete action spaces. The root causes of instability—policy distribution shifts and singular covariance matrices—are not restricted to the continuous case.

We hope that our response can address your concerns, and we are happy to answer any further questions.

References:

[1] Is Bang-Bang Control All You Need? Solving Continuous Control with Bernoulli Policies, Seyde et al. NeurIPS, 2021.

Dear Reviewer Wwkj,

Thank you once again for your constructive comments. As the discussion period is closing soon, we hope our response has addressed your main concerns. Please let us know if you have any remaining questions.

Thank you for your important feedback. We hope that the following responses may address the questions and concerns.

Regarding the weaknesses:

The paper does not provide a detailed empirical analysis of the numerical instabilities to support the theoretical claims.

We have conducted an empirical analysis of the influence of optimization constraints and the impact of large means on numerical stability. Some experimental results were initially placed in the Appendix due to the page limit. We have now added a new section in the main text dedicated to the empirical analysis. Please refer to the revised draft for details.

The propoesed solutions are either difficult to implement or hurt the performance of the algorithms. There is no clear and effective solution presented in this paper.

Since we observed that the vanilla policy gradient (VPG) is numerically more robust compared to PPO, we propose a new solution called "vanilla policy gradient with clipping." This approach simplifies the PPO objective by removing the importance sampling ratio while applying the same clipping rule to the VPG objective. The resulting algorithm combines the effectiveness of PPO with the robustness of VPG. Please refer to the global response and the revised draft for further details.

Although the paper provides some experimental results for the proposed results, they are not comprehensive and do not provide a clear comparison with existing methods.

We have expanded the experimental results in the main text; please refer to the revised draft for details.

Regarding the questions:

Can you provide more insights into the numerical instabilities in policy gradient methods and how they manifest during training? How do these instabilities affect the convergence and performance of deep reinforcement learning algorithms?

Based on our analysis, the numerical instability of TRPO/PPO arises from large means and/or small standard deviations in the stochastic policy during training. When such instabilities occur, the algorithm may encounter arithmetic overflow (e.g., division by zero), leading to NaN or Inf values. While this issue has been recognized as a fundamental challenge in policy gradient methods (as discussed in https://stable-baselines3.readthedocs.io/en/master/guide/checking_nan.html), it is often regarded as a hyperparameter tuning problem. However, the root cause has not been well understood.

Can you provide a comprehensive empirical study on the proposed solutions to mitigate numerical instabilities in policy gradient methods? Can you explain more about the limitations of these solutions?

We have proposed a solution that effectively addresses the numerical issues caused by importance sampling. Please refer to the revised draft for further details.

Please let us know if you have any further questions.