Return Capping: Sample Efficient CVaR Policy Gradient Optimisation

The main issue suggested in this review is flaws with the theoretical proof of the equivalence of Return Capping and standard CVaR PG optimisation. We would summarise the two points the reviewer raises as:

• Given Eqn (7), when $C$ is selected in the $\min(\cdot)$, this trajectory will have no effect on the optimisation of $\pi_\theta$
• Eqn (7) is not being applied to the actual return capping algorithm

However, the above points are not the case, as we can show below.

Addressing the first point:

Given Equation 7 $J^C (\pi_\theta; C) = E_{\tau \sim \pi_\theta} [\min(R(\tau), C)].$

To compute policy gradients, we need to compute the gradient with respect to $\theta$ $\nabla_\theta J^C (\pi_\theta; C) = \nabla_\theta E_{\tau \sim \pi_\theta} [\min(R(\tau), C)].$ This can be expanded to the integral $\nabla_\theta J^C (\pi_\theta; C) = \nabla_\theta \int_\tau P(\tau|\theta) \min(R(\tau), C).$ This integral highlights the flaw in the first point, as when either $R(\tau)$ or $C$ is selected in the $min(\cdot)$, the probability of that given trajectory is still dependent on $\theta$, and thus will have an effect on the policy optimisation. This addresses the first point made.

Addressing the second point:

From the equation above, we can move the gradient inside the integral and use a log-derivative trick to put the integral in the form $\nabla_\theta J^C (\pi_\theta; C) = \int_\tau P(\tau|\theta) \nabla_\theta \log P(\tau|\theta) \min(R(\tau), C).$

We can then return this back to an expectation $\nabla_\theta J^C (\pi_\theta; C) = E_{\tau \sim \pi_\theta} [\nabla_\theta \log P(\tau|\theta) \min(R(\tau), C)],$ and then we can reformulate the trajectory probability in terms of action probability $\nabla_\theta J^C (\pi_\theta; C) = E_{\tau \sim \pi_\theta} [\sum_{t=0}^T \nabla_\theta \log \pi_\theta (a_t|s_t) \min(R(\tau), C)].$ From here, we can apply commonplace RL techniques such as using return-to-go, rather than using full episode returns, and using a Value function baseline. Note that return-to-go is computed as: $\hat{R}^C_t = \min(\sum_{t'=0}^T R(s_{t'}, a_{t'}, s_{t'+1}), C) - \min(\sum_{t'=0}^t R(s_{t'}, a_{t'}, s_{t'+1}), C).$ When incorporating both of these techniques, the gradient update becomes: $\nabla_\theta J^C (\pi_\theta; C) = E_{\tau \sim \pi_\theta} [\sum_{t=0}^T \nabla_\theta \log \pi_\theta (a_t|s_t) (\hat{R}^C_t - V(s_t)].$ In practice, we used a Generalised Advantage Estimator to compute advantage. This gradient update can then be clipped, using the PPO algorithm, and then this is the exact gradient update used in Return Capping. As such the Return Capping gradient update is derived directly from Eqn (7), addressing the review’s second point.

We are happy to provide any further clarification about the proof, or how the derivations above show that the reviewer’s main issues are unfounded.

Other Issues Raised

• In relation to the scale of the environments, we have discussed this thoroughly in our review to r3Le
• The reason we have only included MIX [1] in Lunar Lander is that this baseline requires an optimal Expected Value policy. Lunar Lander was the only environment where Return Capping required the Expected Value policy to set $C^M$. However we could include this baseline in other environments - it generally just converged to the optimal Expected Value policy, similarly to in Lunar Lander. As explained in the paper, the Return Capping (offset) example was included to demonstrate the relative performance of Return Capping accounting for the environment steps required to train an optimal Expected Value policy.

[1] Luo, Yudong, et al. "A simple mixture policy parameterization for improving sample efficiency of cvar optimization."

In relation to the size of the environments presented, we have discussed this in our rebuttal to reviewer r3Le.

Sample Efficiency

For the specific question on sample efficiency, whilst there are improvements to sample efficiency using Return Capping compared to CVaR PG, it is unlikely to achieve improvements of magnitude $\frac{1}{\alpha}$ as by capping returns, we do lose some information about trajectories. However there is still information to be learnt from capped trajectories.

Primarily, examples of trajectories that reach the cap are present in policy optimisation. In standard CVaR PG, we effectively only have negative examples of trajectories that resulted in low returns. Whilst we lose the relative performance between capped trajectories by capping returns, we still do get a gradient between the uncapped low performing trajectories and these capped trajectories. By training using capped trajectories, the policy is able to learn from positive examples (e.g. actions that result in the trajectory reaching the cap) rather than just having to rely on negative examples.

In terms of speed of convergence observed in empirical results, in both Guarded Maze environments and the Lunar Lander environment, we see much faster convergence for Return Capping compared to the risk-sensitive baselines. The Expected Value policy does converge more quickly in all environments but this is unsurprising given it is learning from all uncapped trajectories. In the betting game, the CVaR PG policy does converge more quickly, albeit to a less optimal policy, than Return Capping. We suggest that this is likely due to the CVaR PG policy being a much more simplistic, conservative policy compared to the more optimal Return Capping policy. We agree that in the AV environment, there is minimal performance improvement over CVaR PPO*, but we suggest that the better performance of Return Capping in all other presented environments suggests it is a better method overall.

*It should be noted that the plot in Figure 5a excludes CVaR PPO outliers that did not converge to the optimal CVaR policy (see Figure 7 in the Appendix for the unmodified plot), so although median performance is comparable to Return Capping, in this environment, CVaR PPO is less consistent at finding the CVaR optimal policy.

Thank you for your time spent reviewing this paper and we appreciate your feedback

Environment Complexity

The main issue raised in this review is the limited complexity of the environments used. Whilst we agree that more complex environments would benefit the paper, as far as we are aware, there is very limited work that has presented promising results for optimising static CVaR for episode return in more complex environments.

Some related work does explore more complex environments. In Safe RL, such as [4], the main benchmarks are modified MuJoCo environments [7]. However, unlike in our work, in Safe RL, the objective is to maximise expected return subject to constraints on cost, where environments have a distinct cost and reward function. Some distributional RL work has been done in MuJoCo [5] and Atari [8]. However these works focus on dynamic CVaR, rather than static CVaR.

The work we are aware of optimising for static CVaR referenced in the paper [1, 2, 6] as well as a similar work [3] suggested by Reviewer iMqv all focus primarily on comparatively similar scale environments to the examples presented in our paper. One of the exceptions are some modified MuJoCo environments in [1]. However, in practice, when optimising policies in these environments we found no distinction between optimal Expected Value and optimal CVaR policies. The other exception is the Atari game Asterix in [6]. This paper presents a modification to distributional DQN to optimise for static CVaR and shows promising results. However, results from [8] show that optimising for dynamic CVaR using distributional DQN in Asterix results in better Expected Value performance that a policy optimising for Expected Value, so it is ambiguous whether the improved CVaR performance in [6] was due to a distinct CVaR optimal policy being found, or due to this aforementioned result shown in [8].

The main issue we have found in scaling to more complex environments is an issue inherent to return CVaR optimisation rather than specifically our method which is that for CVaR optimisation techniques to be relevant, the environment has to have distinct CVaR-optimal and Expected-Value-optimal policies. Generally, what we have observed is that one of these policies will be less complex to learn, and so in more complex environments, either the Expected Value optimal policy converges to the CVaR-optimal policy, or all CVaR optimisation methods converge to the Expected Value optimal policy.

Off Policy Methods

Thank you for pointing out the two papers [4, 5] as both are relevant to risk-sensitive RL and we will include them in reference to Related Work. However, as mentioned above, they are both optimising for different objectives compared to our work.

Other Strengths and Weaknesses

Addressing the two points the review raised in the strengths and weaknesses section:

• The reason we have included (offset) In the Lunar Lander environment, is we found better performance by optimising for Expected Value and then using this policy to set $C^M$ as outlined in Section 4.1, even accounting for the additional environment steps required to train the optimal Expected Value policy (see Return Capping (offset)). In all other environments, we set $C^M$ based on a trivial policy that took no actions.
• Conditioning the cap on the initial state is a sensible suggestion if the desired goal is to optimise for CVaR conditioned on the initial state.

[1] Luo, Yudong, et al. "A simple mixture policy parameterization for improving sample efficiency of cvar optimization."

[2] Greenberg, Ido, et al. "Efficient risk-averse reinforcement learning."

[3] Kim & Min, “Risk-Sensitive Policy Optimization via Predictive CVaR Policy Gradient”

[4] Yang, Qisong, et al. "WCSAC: Worst-case soft actor critic for safety-constrained reinforcement learning."

[5] Ma, Xiaoteng, et al. "Dsac: Distributional soft actor critic for risk-sensitive reinforcement learning."

[6] Lim, Shiau Hong, and Ilyas Malik. "Distributional reinforcement learning for risk-sensitive policies.”

[7] Ji, Jiaming, et al. "Safety gymnasium: A unified safe reinforcement learning benchmark."

[8] Dabney, Will, et al. "Implicit quantile networks for distributional reinforcement learning." International conference on machine learning.

Thank you for referencing [3], we will include it in Related Work

Setting Cap Minimum

Addressing the question on setting $C^M$ in larger environments, it is very possible to set a suitable $C^M$, irrespective of the complexity of the environment.

We show in Proposition 1 that the Return Capping optimisation objective is equivalent to optimising for CVaR if $C$ is set to $VaR_\alpha(\pi^*)$, so we need to ensure that $C^M$, is less than or the equal to $VaR_\alpha(\pi^*)$.

Given Eqn (17) $CVaR_\alpha(\pi) \leq VaR_\alpha(\pi), \quad \forall \pi$ and Eqn (18) $CVaR_\alpha(\pi) \leq CVaR_\alpha(\pi^*), \quad \forall \pi$ we know that the $CVaR_\alpha$ of any policy is necessarily less than the $VaR_\alpha(\pi^*)$. So we know that if we set $C^M$ to $CVaR_\alpha$ of any policy, it will be less than $VaR_\alpha(\pi^*)$. This means we can take any policy and sample a set of trajectories and use this sampled $CVaR_\alpha$ to set $C^M$. This could be a random policy, or it could be a policy that maximises Expected Value.

Even if an environment was sufficiently complex that it was not immediately apparent what an appropriate value of $C^M$ was, it would always be possible to use either the $\text{CVaR}_\alpha$ of a random policy, or a policy trained in any manner to set $C^M$.

Requiring a previously trained policy to set $C^M$ does increase the samples required for training. However, for all but the AV environment, Return Capping would outperform all baselines even accounting for the additional training required to learn an Expected Value optimal policy (shown as Return Capping (offset) in Lunar Lander as this was the only environment that required the cap to be set according to the Expected Value policy, but we could include this in all other environments as well).

Environment Simplicity

We have discussed issues with more complex environments further in our rebuttal to reviewer r3Le, but these issues arise more from challenges with optimising for CVaR in general, rather than any specific characteristics of Return Capping.

Additional Baselines

Addressing the lack of inclusion of [1, 2] as baselines, the reason we have not included [2] is because it requires the environment to be formulated as a Context-MDP where the context encapsulates the environment randomness. This limits the applications of the baseline as not all problems can be formulated as such. The reason we have only included [1] in Lunar Lander is that this baseline requires an optimal Expected Value policy. Lunar Lander was the only environment where Return Capping required the Expected Value policy to set $C^M$. However we could include this baseline in other environments - it generally just converged to the optimal Expected Value policy, similarly to in Lunar Lander.

[1] Luo, Yudong, et al. "A simple mixture policy parameterization for improving sample efficiency of cvar optimization."

[2] Greenberg, Ido, et al. "Efficient risk-averse reinforcement learning."

[3] Kim & Min, “Risk-Sensitive Policy Optimization via Predictive CVaR Policy Gradient”