Paper: ReLU to the Rescue: Improve Your On-Policy Actor-Critic with Positive Advantages (ICLR-2024)

Reviewer comment (Presentation):
The names of metrics shown in Figure 1 (Median, IQM, Mean, OG) should be emphasized above. For example, change 'robust normalized median, interquartile mean, and mean metrics' to 'robust normalized median (Median), interquartile mean (IQM), mean (Mean) and optimal gap (OG) metrics'.

Author response:
> The names of metrics shown in Figure 1(Median, IQM, Mean, OG) should be emphasized above.(e.g. change 'robust normalized median, interquartile mean, and mean metrics' to 'robust normalized median(Median), interquartile mean(IQM), mean(Mean) and optimal gap(OG) metircs'). 

We are adding your suggestions on the metric names.

Reviewer comment (Writing):
The multiple use of some terms in many places made me confused. For example, I think A3C represents two different algorithms respectively in Figure 1 and Figure 2.

Author response:
> The multiple use of some terms in many places made me confused, e.g. I think A3C represent two different algorithms respectively in Figure 1 and 2.

We agree that we can be more clear here and propose making the following changes to the names in Figure 1 in order to help aleviate the confusion:

RMPG -> VSOP "all-actions"
A3C -> VSOP no-ReLU-Adv.
No Spectral -> VSOP no-spectral
No Thompson -> VSOP no-Thompson

Reviewer comment (Theory):
Except in the case of sparse rewards, it is generally not acceptable to assume that the difference between the expected value of the value function over the next states and the value function at the current state is zero.

Author response:
> Except in the case of sparse rewards, it is generally not acceptable to assume that the difference between the expected value of the value function over the next states and the value function at the current state is zero.

Thank you for this comment as your attention has led to a clear improvement. We hope that our comment to all reviewers above has addressed this concern. Specifically, upon a close re-examination of this assumption, we agree it would not generally hold. Fortunately, we can show that it is not necessary and that $C_\pi(\mathbf{s})$ can be smaller than under the zero-expectation assumption, resulting in a smaller additive:

**Lemma.**

$$\text{ReLU}(a) <= |a|$$

*Proof*.
$$ \text{ReLU}(a) = \max(0, a)$$
$$\quad\quad\quad\quad= \frac{1}{2}a + \frac{1}{2}|a|$$
$$\quad\quad\quad\quad= \\{a  \text{ if } a\geq0\text{ else }0\text{ if } a<0$$
$$\quad\quad\quad\quad\leq \\{a  \text{ if } a\geq0\text{ else }-a\text{ if } a<0 (-a > 0)$$
$$\quad\quad\quad\quad= |a| $$

Therefore,

$$C_\pi(\mathbf{s}) = \iint \left(v_\pi(\mathbf{s'}) - v^\pi(\mathbf{s}) \right)^+ dP(\mathbf{s}' \mid \mathbf{S}^\mathrm{t} = \mathbf{s}, \mathbf{A}^\mathrm{t} = \mathbf{a}) d\Pi(\mathbf{a} \mid \mathbf{S}^\mathrm{t} = \mathbf{s})$$
$$\quad\quad\quad\leq \iint \left|v_\pi(\mathbf{s'}) - v_\pi(\mathbf{s}) \right| dP(\mathbf{s}' \mid \mathbf{S}^\mathrm{t} = \mathbf{s}, \mathbf{A}^\mathrm{t} = \mathbf{a}) d\Pi(\mathbf{a} \mid \mathbf{S}^\mathrm{t} = \mathbf{s}) $$
$$\quad\quad\quad\leq \iint K_\pi\left|\left|\mathbf{s'} - \mathbf{s} \right|\right| dP(\mathbf{s}' \mid \mathbf{S}^t = \mathbf{s}, \mathbf{A}^\mathrm{t} = \mathbf{a}) d\Pi(\mathbf{a} \mid \mathbf{S}^\mathrm{t} = \mathbf{s})$$


Please let us know if you have further concerns about this point.

Reviewer comment (Writing):
Typo on (4): $\frac{\beta^\alpha\sqrt{\tau}}{\sqrt{\Gamma(\alpha)2\pi}}\rightarrow \frac{\beta^\alpha \sqrt{\tau}}{\Gamma(\alpha)\sqrt{2\pi}}$, $exp(\beta h_i)\rightarrow exp(-\beta h_i)$.

Author response:
> typo['s] on ..

Nice catches. We have updated the manuscript with, $\frac{\beta^{\alpha}\sqrt{\tau}}{\Gamma(\alpha)\sqrt{2\pi}}$ and 'Letting $\nabla v_\pi^*(\mathbf{s}_0)$'.

Reviewer comment (Theory):
Can you explain in detail how to combine Thompson sampling with Bayesian inference? And why this is a better **state-aware** exploration method?

Author response:
> Neccessary background and discribtion of Thompson sampling is missing. Can you explain in detail how to combine Thompson sampling with Bayesian inference? And why this is a better state-aware exploration method?

Approximate Bayesian inference over the parameters of the policy, $\theta$, yields a distribution over those parameters, $p(\theta \mid \mathcal{D})$. Sampling a policy from this distribution, $\hat{\theta} \sim p(\theta \mid \mathcal{D})$, is as easy as sampling a dropout mask and then running a forward pass of the network, yielding the likelihood, $\pi(\mathbf{a} \mid \mathbf{s}, \hat{\theta})$. Then sampling an action is done by sampling an action from the sampled policy, $\mathbf{a} \sim \pi(\mathbf{a} \mid \mathbf{s}, \hat{\theta})$. This is precisely the procedure descibed by Thompson sampling. We outline this procedure in lines 5-6 of Algorithm 1 and have a section on Thomspon sampling for RL in the related works section, but we agree with your assessment that this can be made more explicit in the main text and we are working on incorporating the above description to the manuscript. We would appreciate any feedback on it should you have time.

We hypothesize that this is a better state-aware exploration method for two reasons. First, for less frequently visited states the diversity of the sampled paramters of the policy will be greater promoting more exploration. As a state is visited more often under actions that yield positive advantages, the diversity of samples will concentrate promoting less exploration. Thus, we get more exploration for states that we have less experience of good actions, and less exploration in states where we know what actions lead to good expected returns. Second, this exploration is done around the mode of the policy distribution, so the model is less likely to explore actions that are far from the mode, which could be more likely to lead to failure.

Reviewer comment (Theory):
Is it necessary to assume it as a gamma-normal distribution?

Author response:
> Is it necessary to assume it as a gamma-normal distribution?

The gamma-normal assumption allows us to interpret adding dropout and weight-decay regularization as sensible approximate Bayesian inference without adding complex computational overhead to the original A3C optimization algorithm. As such, this assumption primarily serves to ground Thompson sampling through approximate Bayesian inference and is not requisite for Theorem 1. As with the original result of the policy gradient theorem, the results in Equations (5-6) do not make any distributional assumptions on $\pi$ and should hold for all policies with differentiable probability densities/distributions.

Reviewer comment (Theory):
In the theoretical component of the algorithm, while there is a theoretical result which shows how the value function improves as a result of the modifications, which is very nice, the work could shed light on the role of these parameters on the overall convergence of the modified A3C.

Author response:
> In the theoretical component of the algorithm, while there is a theoretical result which shows how the value function improves as a result of the modifications, which is very nice, but I think that the work could shed light on the role of these parameters on the overall convergence of the modified A3C?

We are very interested in establishing convergence rates for our algorithm, but believe that such an analysis would consitute a self contained journal paper akin to the works of *Agarwal et al. 2021* or *Shen et al. 2023*. Note that these works were published several years after the acceptance of the original A3C paperto ICML in 2016.

Agarwal, Alekh, et al. "On the theory of policy gradient methods: Optimality, approximation, and distribution shift." The Journal of Machine Learning Research 22.1 (2021): 4431-4506.

Shen, Han, et al. "Towards Understanding Asynchronous Advantage Actor-critic: Convergence and Linear Speedup." IEEE Transactions on Signal Processing (2023).

Reviewer comment (Experiments):
Neither the simulations nor the theoretical results shed light on the exact role of the choice of dropout, ReLU, etc., on the impact of their bounds—both theoretical bounds and empirical bounds—with the exception of the justification of the spectral normalization step.

Author response:
> I also noticed that neither the simulations nor the theoretical results shed light on the exact role of the choice of dropout, relu etc., on the impact of their bounds, both theoretical bounds and empirical bounds, with the exception of the justification of the spectral normalization step.  Or perhaps, a comparison to A3C, since that is the algorithm the current paper is based on?

While dropout does not factor into the bound in Equation (6), we can improve the clarity of Theorem 3.1 to emphasize that the bound is directly dependent on applying the ReLU function to the advantage estimates. In Theorem 3.1 we write, $(x)^+ := \max(0, x)$, which is the definition of the ReLU function.

Reviewer comment (Theory):
A question is whether the constant K is policy dependent, in which case, what would the policy improvement step be optimizing over in (6). The K-Lipschitz assumption is introduced with respect to a particular value function over all \( s \in \mathcal{S} \), which raises the question of whether K is dependent on policies π.

Author response:
> I noticed that the K-Lipschitz assumption is introduced with respect to a particular value function over all which makes me wonder if K is dependent on policies π. A question I have is whether the constant K is policy dependent, in which case, what would the policy improvement step be optimizing over in (6)?

This is an excellent question and we are very thankful for your insight. We hope that the response to all reviewers above serves to answer it. In summary, we agree that K is policy dependent in general, and the dependence has given new insights into the role of spectral normalization as a regularizer of K. Please do not hesitate to follow up if you have furhter clarifying questions regarding this point.

Reviewer comment (Theory):
Another question is whether the assumptions on the action distribution follow the normal-gamma in the bound on the value function improvement in (6), and what other assumptions are incorporated in the bound on the value function improvement.

Author response:
> Another question I have is if the assumptions on the action distribution follow the normal-gamma in the bound on the value function improvement in (6)?

As with the original result of the policy gradient theorem, the result in Equation (6) does not make any distributional assumptions on either $q_\pi$, $v_\pi$, or $\pi$, and should hold for all policy distributions with differentiable densities. 

> What other assumptions are incorporated in the bound on the value function improvement?

As stated in Theorem 3.1, two assumptions are necessary to get the result in Equation (6): 1.) that rewards, $\mathrm{R}_t$, are non-negative (which we make without loss of generality); and 2.) that the gradient of the policy is a conservative vector field. With the additional assumption that the value function is K-Lipschitz, we can bound the constant as in Equation (7). Note that in response to Reviewer kpwW's comment, we can drop the additional assumption that the expected difference in the value function between subsequent states is 0.

Reviewer comment (Writing):
It would be helpful if the authors could improve the motivation of the work. In particular, the Intro and the background do not provide an effective argument as to what problems are really being addressed in this work and why should one care.

Author response:
Dear Reviewer,

Thank you again for sharing your valuable feedback with us. We strongly believe that the motivation, theory, and empirical results make this submission worthy of acceptance. Moreover, we think our response addresses your concerns and have updated the paper accordingly. Specifically, we added commentary on the Lipschitz condition, included your formatting suggestions, and reworked the abstract and introduction to help communicate the motivation better. If you have no further concerns, we kindly ask you to consider raising your score. Otherwise, we would also be happy to discuss further, should you find the time. 

Best regards,
The Authors

Reviewer comment (Presentation):
The reviewer would suggest moving the algorithm pseudocode to an earlier place. It is rather inconvenient that Theorem 3.1 is presented before the algorithm.

Author response:
We are moving the algorithm closer to its first reference on Page 4 in the updated manuscript, which we are working to share shortly.

Reviewer comment (Theory):
Equation (7) seems to assume a Lipschitz condition on v(s)? Please elaborate.

Author response:
We hope that the comment above to all reviewers has addressed this question. Please to not hesitate to follow up if you have further clarifying questions concerning this point.

Reviewer comment (Reproducibility):
It might be helpful to explain how many training steps are implemented. Also, would it be possible to show the training curves?

Author response:
For the MuJoCo-Gym experiments, models interact with each environment for 3 million training steps; for the MuJoCo-Brax experiments, 50 million training steps; and for the ProcGen experiments, 25 million training steps. To save space, we show all training curves in Appendix E. We will make the references to these images in the main text more apparent.

Reviewer comment (Writing):
It would be helpful if the authors could improve the motivation of the work. In particular, the Intro and the background do not provide an effective argument as to what problems are really being addressed in this work and why should one care.

Author response:
Dear Reviewer,

Thank you again for sharing your valuable feedback with us. We strongly believe that the motivation, theory, and empirical results make this submission worthy of acceptance. Moreover, we think our response addresses your concerns and have updated the paper accordingly. Specifically, we added commentary on the Lipschitz condition, included your formatting suggestions, and reworked the abstract and introduction to help communicate the motivation better. If you have no further concerns, we kindly ask you to consider raising your score. Otherwise, we would also be happy to discuss further, should you find the time. 

Best regards,
The Authors

Reviewer comment (Presentation):
The reviewer would suggest moving the algorithm pseudocode to an earlier place. It is rather inconvenient that Theorem 3.1 is presented before the algorithm.

Author response:
We are moving the algorithm closer to its first reference on Page 4 in the updated manuscript, which we are working to share shortly.

Reviewer comment (Theory):
Equation (7) seems to assume a Lipschitz condition on v(s)? Please elaborate.

Author response:
We hope that the comment above to all reviewers has addressed this question. Please to not hesitate to follow up if you have further clarifying questions concerning this point.

Reviewer comment (Reproducibility):
It might be helpful to explain how many training steps are implemented. Also, would it be possible to show the training curves?

Author response:
For the MuJoCo-Gym experiments, models interact with each environment for 3 million training steps; for the MuJoCo-Brax experiments, 50 million training steps; and for the ProcGen experiments, 25 million training steps. To save space, we show all training curves in Appendix E. We will make the references to these images in the main text more apparent.

Reviewer comment (Experiments):
While the results presented in the paper are indeed strong, the paper is still premature in terms of analysis. Throughout the paper, it remained unclear how much each added trick contributes to the overall performance.

Author response:
>Throughout the paper it remained unclear to me how much each added trick contributes to the overall performance. As far as I understand, VSOP was ablated with:
    VSOP (dropout, spectral, relu, thompson)
    A3C (dropout, spectral, thompson)
    No Spectral (dropout, relu, thompson)
    No Thompson (dropout, spectral, relu)
From these ablations it is evident that one cannot determine the contribution of each of the tried tricks. 

Your understanding of the ablation is mostly correct. To help with clarity, we propose to make the following changes to the ablated variant names in Figure 1.:

- RMPG -> VSOP "all-actions"
- A3C -> VSOP no-ReLU-Adv.
- No Spectral -> VSOP no-spectral
- No Thompson -> VSOP no-Thompson

We hyperparameter tune each variant according to the same procedure as for VSOP. The hyperparameter tuning search space includes the dropout rate in the range of 0.0-0.1. As such, we effectivelty ablate dropout rate in the hyperparamter tuning phase to isolate the effect of the added regularization of dropout from the that of Thompson sampling. The estimated optimal dropout rates are 0.025 (VSOP, VSOP no-Thompson, and VSOP no-ReLU), 0.05 (VSOP no-spectral), and 0.0 (VSOP "all-actions").

We respectfully maintain these ablations are sufficient to show that the combination of the four modifications is necessary for the performance of VSOP by demonstrating that the subtraction of any modification results in decreased performance and stability.

We are also interested in the results of seeing the effect size of the remaining nine combinations with respect to A3C, but we have not seen significant evidence that any subset would lead to the performance of VSOP. Lest to justify the substantial computational cost of the hyperparameter tuning and subsequent training runs required for a fair comparison.

Reviewer comment (Experiments):
The presented ablations — VSOP (dropout, spectral, relu, thompson), A3C (dropout, spectral, thompson), No Spectral (dropout, relu, thompson), No Thompson (dropout, spectral, relu) — do not allow one to determine the contribution of each of the tried tricks.

Author response:
>Throughout the paper it remained unclear to me how much each added trick contributes to the overall performance. As far as I understand, VSOP was ablated with:
    VSOP (dropout, spectral, relu, thompson)
    A3C (dropout, spectral, thompson)
    No Spectral (dropout, relu, thompson)
    No Thompson (dropout, spectral, relu)
From these ablations it is evident that one cannot determine the contribution of each of the tried tricks. 

Your understanding of the ablation is mostly correct. To help with clarity, we propose to make the following changes to the ablated variant names in Figure 1.:

- RMPG -> VSOP "all-actions"
- A3C -> VSOP no-ReLU-Adv.
- No Spectral -> VSOP no-spectral
- No Thompson -> VSOP no-Thompson

We hyperparameter tune each variant according to the same procedure as for VSOP. The hyperparameter tuning search space includes the dropout rate in the range of 0.0-0.1. As such, we effectivelty ablate dropout rate in the hyperparamter tuning phase to isolate the effect of the added regularization of dropout from the that of Thompson sampling. The estimated optimal dropout rates are 0.025 (VSOP, VSOP no-Thompson, and VSOP no-ReLU), 0.05 (VSOP no-spectral), and 0.0 (VSOP "all-actions").

We respectfully maintain these ablations are sufficient to show that the combination of the four modifications is necessary for the performance of VSOP by demonstrating that the subtraction of any modification results in decreased performance and stability.

We are also interested in the results of seeing the effect size of the remaining nine combinations with respect to A3C, but we have not seen significant evidence that any subset would lead to the performance of VSOP. Lest to justify the substantial computational cost of the hyperparameter tuning and subsequent training runs required for a fair comparison.

Reviewer comment (Theory):
Advantage clipping is motivated by the fact that regular advantages cannot be Gamma distributed, because Gamma has only positive support. While clipping does fix the advantage estimates' support issue, it is unclear why this operation should make the estimates Gamma distributed as assumed by the theory.

Author response:
>Advantage clipping is motivated by the fact that regular advantages cannot be Gamma distributed, because Gamma has only positive support. While clipping does fix the advantage estimates' support issue, why should this operation make the estimates Gamma distributed as assumed by the theory?

The intuition behind assuming a gamma distribution for the clipped advantages is that advantages ideally have zero mean by construction (we subtract the state-action value by its expected state value over actions), so clipping at zero will result in a heavy-tailed distribution. Gamma distributions are sensible hypotheses for heavy-tailed distributions. In the following linked image, for a training run of Humanoid-v4, we plot the marginal histograms for the advantages (left) and the clipped advantages (right) over each training update:

[Histogram of Clipped Advantages](https://i.imgur.com/wExvlMC.png)

The clipped advantage histogram on the right lends evidence to the gamma assumption (at least for the marginal distribution). We may expect multi-modality at the state-action level, which integration over actions may marginalize out at the state level; however, we would still expect a heavy tail in both cases.