Paper: Constrained Reinforcement Learning as Wasserstein Variational Inference: Formal Methods for Interpretability (ICLR-2024)

Reviewer comment (Writing):
This paper covers a wide variety of concepts. However, there seems to be issues with key definitions and notations, which remain ambiguous or undefined, thus making several key ideas challenging to comprehend. For instance, in Formula 2), the subscript 'i' is left undefined, and $b_i$ is also undefined.

Author response:
>__Weaknesses:__ * In Formula 2), the subscript 'i' is left undefined, and $b_i$ is also undefined.

This aligns with comments from Reviewer x9MY. We have clarified that $b_i$ is a fixed limit for the i-th constraint (see the red highlight in Page 3).

Reviewer comment (Theory):
In Section 4.1, the definition of $q_\theta$ appears flawed. There seems to be a duplicate probability over '$a$'. The second term should likely be $p(s|a,\theta)$ instead.

Author response:
>* In Section 4.1, the definition of $q_\theta$ appears flawed. There seems to be a duplicate probability over '$a$'. The second term should likely be $p(s|a,\theta)$ instead.

We appreciate your thoughtful review. From a macro perspective, the definition of $q_\theta$ is derived from Equation 1, where the probability $p(s,a|\theta)$ originates directly from the second term on the right-hand side. Please consider our response in Reviewer x9MY (Q1) regarding the correctness of the term $p(s,a|\theta)$ in Equation 1 and in underscoring the formulation as a Markov Decision Process.

On a micro level, we do not perceive any redundancy in probability over 'a' as $q(a)$ is not $p(a)$. The individual $q(a)$ serves as a probability function with respect to 'a,' representing the approximation (variational distribution) for the optimality likelihood $p(\tau|\mathcal{O})$.

We appreciate this may have led to some confusion, and to ensure transparency, we incorporate the use of $p(\tau|\theta)$ instead of $p(s,a|\theta)$ in the revised version (see the red highlight in Equation 1, and the definition of $q_\theta$ and $p(\tau|\mathcal{O})$), consistent with the response to Q1 in Reviewer x9MY.

Reviewer comment (Theory):
In Section 4.1, the definition of $p(\tau|O)$ seems incorrect. The term "equivalence" should likely be replaced with "proportional to" (please refer to formula (1)).

Author response:
>* In Section 4.1, the definition of $p(\tau|\mathcal{O})$ seems incorrect. The term "equivalence" should likely be replaced with "proportional to" (please refer to formula (1)).

Thank you for your feedback. Indeed, in principle, "proportional to" is a more accurate formulation, and it was duly considered in our initial algorithmic framework construction. However, the implementation posed a challenge. Subsequently, we discovered that the direct "equivalence" rather than "proportional to" yields the same convergence performance in real implementation (as detailed in Lemma 2 and Theorem 2 for theoretical understanding). In response to the reviewer's comment, the revised version use "proportional to" instead of "equivalence" along with a clarification that "equivalence" is used in our actual implementation.

Reviewer comment (Writing):
Formula 3) is incomprehensible without the definitions provided in the appendix. The hyperparameter $l$ and the function $\mathcal(G)$ are undefined in the formula and its introductory explanation. It would be helpful to move this content from the appendix into the main body of the paper.

Author response:
>* Formula 3) is incomprehensible without the definitions provided in the appendix. ...

We move the definitions from the appendix into the main body of the paper. Please refer to the red highlights around Equation 3 for details.

Reviewer comment (Writing):
In Formula 3), it's unclear what the omitted term in $\cdot$ is. According to the appendix, it's supposed to be a hyperparameter $l$. If that's the case, it's generally inappropriate to omit the hyperparameter in the formula.

Author response:
>* In Formula 3), it's unclear what the omitted term in $\cdot$ is. According to the appendix, ...

The symbol $\cdot$ signifies the hyperparameter $l$. We would like to retain the use of the symbol $\cdot$ in Equation 3 for consistency with previous papers, such as [1] (Equation 11) and [2] (Equation 15). This choice is made to clearly depict the slicing process in Wasserstein metrics.

Reviewer comment (Writing):
The definition of $R_\mu(l,\tilde{\theta})$ in the appendix is also confusing. Please verify whether the $x$ in the formula should be $\mu$ and provide clarification.

Author response:
>* The definition of $\mathcal{R}_{\mu}(l,\widetilde{\theta})$ in the appendix is also confusing. Please verify ...  

We thoroughly verify that $x$ in the formula is correct as $\mu$ represents probability measures and $x$, defined as $\{x\in \mathbb{R}^d|\langle x,\widetilde{\theta} \rangle=l\}$, denotes hyperplanes. It is essential to integrate over $\mu$ rather than $x$. For further clarification, please refer to both [2] (Equation 6/7) and [3]. We appreciate your observation.

Reviewer comment (Reproducibility):
The inference step of $p(\theta|D)$ in Algorithm 2 needs clarification. Its reference cannot be found in Section 3. Please specify the exact location or formula. Please also distinguish between $\theta$, $\theta^\mu$, and $\theta^Q$. The rationale behind employing different updating methods for $\theta$ and the others also needs to be elucidated.

Author response:
>* The inference step of $p(\theta|D)$ in Algorithm 2 needs clarification. Its reference cannot be found in Section 3. Please specify the exact location or formula. ...  

Thanks for your comment.   

Regarding the reference of $p(\theta|D)$: this information can be found in the second paragraph of Section 3 (where we state, "the prior probability $p_{D}(\theta)$ is derived from the posterior probability $p(\theta|D)$, where the parameter $\theta$ is inferred from the training dataset $D$"). The term 'inference' here signifies the process of obtaining parameters $\theta$ from the dataset $D$, which, in our specific implementation, corresponds to the training process to determine the parameter values of ${\theta}=[\theta^\mu,\theta^Q]$ in actor and critic networks.  

Distinguishing $\theta$, $\theta^\mu$ and $\theta^Q$: The term $\theta$ serves as a general representation of the parameters. Specifically, in our implementation, actor-critic networks are employed, and the parameters are denoted as $\theta^\mu$ and $\theta^Q$ (refer to 'Initialize:' in Algorithm 2).  

Rationale of different updating methods: The specific updating methods are elucidated in Section 4.2 (refer to 'Policy Updating'), where the detailed gradients of the actor and critic networks, $\delta_{\theta^\mu}$ and $\delta_{\theta^Q}$, are defined in Equation 10 in Appendix B.

Reviewer comment (Writing):
There are typos in the paper. For example, the "hupersurface" before the formula (3) should be "hypersurface", and in Section 4.2, the definition of $\tilde{g}_i(\tau)$ lacks a square bracket on the right.

Author response:
>* Typos:  

We have fixed in the revised version. Thank you!

Reviewer comment (Theory):
The theoretical results appear to lack key assumptions. Many conclusions are questionable and would benefit from further contemplation. For example, the proof of Proposition 1 claims that A-GSWD is a pseudometric rather than a metric. More significantly, it's questionable whether the proof can leverage the fact that "Wasserstein distance is a metric," given that A-GSWD is defined between $\mu$ and $\nu$, not $\mathcal{G}\mu$ and $\mathcal{G}\nu$. If you swap $\mu$ and $\nu$, their corresponding hyperparameters in $\mathcal{G}$ (as output by the neural network) will not be interchanged. Thus, symmetry might not generally hold.

Author response:
>* The proof of Proposition 1 claims that A-GSWD is a pseudometric rather than a metric. More significantly, ...  
  
Thanks for this detailed and thoughtful comment. For clarity, it can be proven that A-GSWD is still symmetric, given that it is defined between $\mu$ and $\nu$. To prove symmetry, since the k-Wasserstein distance is a metric [1]:  
  
$$W_k\left(\mathcal{G}\_\mu(\cdot,\widetilde{\theta};g_{rl}),\mathcal{G}\_\nu(\cdot,\widetilde{\theta};g_{rl})\right)=W_k\left(\mathcal{G}\_\nu(\cdot,\widetilde{\theta};g_{rl}),\mathcal{G}\_\mu(\cdot,\widetilde{\theta};g_{rl})\right)$$  
  
Thus, there exists:  
  
$$\bf{A-GSWD}\_{k}(\mu,\nu)=\left(\int\_{\mathcal{R}\_{\widetilde{\theta}}} W\_k\left(\mathcal{G}\_\mu(\cdot,\widetilde{\theta};g_{rl}),\mathcal{G}\_\nu(\cdot,\widetilde{\theta};g_{rl})\right)\mathrm{d}\widetilde{\theta}\right)\^{\frac{1}{k}}=\left(\int\_{\mathcal{R}\_{\widetilde{\theta}}} W\_k\left(\mathcal{G}\_\nu(\cdot,\widetilde{\theta};g_{rl}),\mathcal{G}\_\mu(\cdot,\widetilde{\theta};g_{rl})\right)\mathrm{d}\widetilde{\theta}\right)\^{\frac{1}{k}}=\bf{A-GSWD}\_{k}(\nu,\mu)$$   
  
Therefore, symmetry holds. We emphasize this in the proof of proposition 1 in Appendix C.3 (please see the red highlight in the revised manuscript).

Reviewer comment (Theory):
In the proof of Proposition 3, the expectation following the max operator should be eliminated. If the action is selected based on an existing policy, the max operator becomes irrelevant.

Author response:
>* In the proof of Proposition 3, the expectation following the max operator should be eliminated. ...  
  
To clarify, $\pi$ represents a policy in the Policy Improvement process (to get the optimized policy that achieves the highest value), not a stationary policy. Consequently, we believe that the expectation should not be eliminated.

Reviewer comment (Theory):
In Equation (16), the first inequality does not generally hold. It often holds when $\pi_{new}$ corresponds to the argmax Q in the Bellman optimality function. This point needs clarification.

Author response:
>* In Equation (16), the first inequality does not generally hold. It often holds ...  
  
The first inequality in Equation 16 relies on Equation 15 (please note that 'According to Equation 15 and Equation 8, it yields: ...' ). We highlight (red) this relationship in the proof in Appendix C.3. Thank you for highlighting this.

Reviewer comment (Theory):
The Wasserstein distance cannot be related to the discrepancy between quantiles in (18) unless $p$ (parameter) and $d$ (dimension) are both 1 in Wasserstein distance.

Author response:
>* The Wasserstein distance cannot be related to the discrepancy between quantiles in (18) unless $p$ (parameter) and $d$ (dimension) are both 1 in Wasserstein distance.  
  
Yes, we agree with this comment. Our proof in Equation 18 is consistent with this statement, and uses the Wasserstein distance based on $W_1$ ($p=1$) and one-dimensional quantiles $q_t^{i,*}$ ($d=1$). We appreciate the detailed feedback, and in the revised version, we highlight (red) this alignment in the proof.

Reviewer comment (Theory):
In Formula (19), Lemma 5.1 in Cai et al. (2019) relies on the linearity of the action-value function, which does not generally apply in value networks.

Author response:
>* In Formula (19), Lemma 5.1 in Cai et al. (2019) relies on the linearity of the action-value function, ...  
  
We agree with this comment. In Equation 19, both the terms $f_i^{(H)}$ and $f_{0,i}^{(H)}$ represent quantiles, serving as local linearization that can be applied with Lemma 5.1 in Cai et al. (2019). In the revised version, we emphasize this point (red highlight) in the proof.

Reviewer comment (Evaluation):
The empirical results lack certain key definitions, and their credibility is questionable. The constraint limit in the plots is undefined. Is this intended to represent the tolerance?

Author response:
>* The constraint limit in the plots is undefined. Is this intended to represent the tolerance?  
  
The constraint limit relies on the benchmark setting [2] (Acrobot) and the specific tasks, with the intention of facilitating a clear comparison of constraint convergence. We provide one additional sentence in Section 6 to clarify this in the revised version.

Reviewer comment (Theory):
The implementation of distributional Reinforcement Learning (RL) in this paper lacks a clear rationale. Despite its application, it remains unclear why it's necessary and what benefits it brings to the study.

Author response:
__Motivation.__ The implementation of distributional Reinforcement Learning (RL) in this paper lacks a clear rationale. Despite its application, it remains unclear why it's necessary and what benefits it brings to the study.

In principle, distributional representation is an intuitive choice that may offer a more comprehensive and richer value-distribution information, enhancing the stability of the learning process. This is particularly important for Bayesian learning processes, which often face challenges in achieving stable convergence. Additionally, policy optimization tends to be biased toward actions with high-variance value estimates, as some of these values may be overestimated by random chance [2].

Moreover, the distributional representation significantly enhances interpretability. As demonstrated in Equation 18 of the proof, it uses quantiles derived from the distributional representation to formally establish the transparency (i.e., quantitive convergence rate) of convergence process outlined in Theorem 2. Ultimately, our experiments, both simulated and real-world, verify a more stable and interpretable learning process, as illustrated in Figure 3 and 6.

We appreciate this thoughtful comment, and to provide further clarity we emphasize this in Appendix B.1.

Reviewer comment (Novelty):
While the introduction of AGSWD stands as a significant contribution, given that GSWD has already been investigated, simply integrating GSWD with neural network outputs doesn't constitute a major advancement in the field.

Author response:
__Novelty.__ While the introduction of AGSWD stands as a significant contribution, given that GSWD has already been investigated, simply integrating GSWD with neural network outputs doesn't constitute a major advancement in the field.

We appreciate the viewpoint and thank the reviewer for recognising the significance of the contribution.

Reviewer comment (Writing):
How should we interpret $g_{rl}$? Its role in Formula (3) isn't clear. Is it intended to be a hyperparameter that defines the plane on which $\Tilde{\theta}$ lies? Or is it a term returned by the RL algorithm?

Author response:
>* How should we interpret $g\_{rl}$? Its role in Formula (3) isn't clear. Is it intended to be a hyperparameter ...

Thanks for the comment. $g_{rl}$ is a measurable mapping $g_{rl}: \mathcal{X}\rightarrow\mathcal{R}_{\widetilde{\theta}}$, where $\mathcal{R}\_{\widetilde{\theta}}$ is a space containing a compact set of all feasible parameters $\widetilde{\theta}$ for $g\_{rl}$ (see Page 4 for details). $g\_{rl}$ can be interpreted as a parameter for the spatial Radon transform $\mathcal{G}$, which is $\mathcal{G}\_{\mu}(l,\widetilde{\theta}) = \int\_{\mathbb{G}^d} \delta(l-\langle g\_{rl},\widetilde{\theta} \rangle) \mathrm{d}\mu$.

Reviewer comment (Writing):
In Formula (18), what does the symbol $H$ denote?

Author response:
>* In Formula (18), what does the symbol $H$ denote?  

$H$ is the layers of the neural network, defined in the statement of Theorem 2 (Section 5) and Table 1. We highlight this in red in Equation 18.

Reviewer comment (Theory):
On page 3, last line: "where, upon satisfying the constraints, the agent enters a state considered as safe". In my understanding, in the expected constraint formulation (Equation 2), if the constraints are satisfied, then the policy is considered "safe". "Safe/unsafe" states typically refer to a CMDP formulation with constraint sets, i.e. a part of the state-action space is safe, and the rest is unsafe. While it is useful to obtain a graphical model framework equivalent to the control as inference framework [1], a more appropriate justification of Figure 1 should have been that the optimality variables are influenced by the constraint, rather than saying that the agent enters a safe state. This means that \( p(\tau|O) \) becomes 0 as soon as the constraint is not satisfied in expectation. The authors do not model the probability in this way, but rather learn a policy by maximizing reward when within the feasible region, and minimizing constraint when outside the feasible region (to get within the feasible region). This makes sense intuitively, but equivalence to constrained RL (equation 2) is not formally established.

Author response:
We are grateful for your assessment of our research. We agree with many of your statements like '"Safe/unsafe" states typically refer to a Constrained Markov Decision Process formulation with constraint sets' and 'learn a policy by maximizing reward when within the feasible region, and minimizing constraint when outside the feasible region'. However, we would like to address some specific statements to improve our shared understanding:

> 1. a more appropriate justification of Figure 1 should have been that the optimality variables are influenced by the constraint,

Yes - we concur that 'the optimality variables are influenced by the constraint'. To elaborate, when $\mathcal{O}=1$, it signifies that the trajectory $\tau$ is __not only optimized but also satisfied the constraints__. In other words, it denotes an ideally optimized trajectory that meets the prescribed constraints. A clear explanation has been highlighted (red) at the outset of Section 3 on page 3. Thank you for highlighting this.

> 2. rather than saying that the agent enters a safe state. ... This makes sense intuitively, but equivalence to constrained RL (equation 2) is not formally established.

It is important to clarify that, in principle, constrained Reinforcement Learning (RL) cannot guarantee full adherence to constraints. Throughout the learning process, an agent is awarded or penalized, pushing it towards a 'safe' state. Consequently, there exists a certain probability of violating constraints in states or actions. This challenge poses a significant barrier for applying constrained RL in safety-critical applications.

More importantly, while AWaVO does not ensure absolute safety, we offer a clear theoretical insight in Theorem 2. This theorem quantitatively interprets the violation: the convergence rate will follow a sublinear $\Theta(1/\sqrt{T})$ if constraints are violated with an error of $\Theta(1/{m^{\frac{H}{4}}})$ (which decreases as the width and layer count of the neural network, denoted as $m$ and $H$ respectively, increase), with a probability of at least $1-\delta$. For detailed information, we refer to Equations 26-28 in the proof of Theorem 2.  

Consequently, we believe that AWaVO is not 'equivalent' to Equation 2, but surpasses constrained RL in terms of interpretability.

Reviewer comment (Theory):
In Equation 1, isn't it more appropriate to use \( p(\tau|\theta) \) instead of \( p(s,a|\theta) \) for the second term of the right hand side?

Author response:
Yes - we appreciate this insightful interpretation of Equation 1. Using $p(\tau|\theta)$ in the second term would bring Equation 1 closer in form to Bayes' theorem, making it more intuitively understandable. Essentially, the formation process of the second term outlines a Markov Decision Process, wherein $p(a_t|s_t,\theta)$ and $p({s}_{t+1}|{s}_t,{a}_t)$ represent the policy state $\pi$ and transitions, respectively.   

The rationale behind our use of $p(s, a|\theta)$ in this context lies in the Markov property, indicating that the present state encompasses all the information necessary to predict future states, irrespective of historical states. Consequently, we adopt $p(\tau|\theta)$ in Equation 1 (see the red highlight in our revision) and explicitly clarify the Markov property for better comprehension.

Reviewer comment (Writing):
In equation 2, what is \( b_i \)? Maybe I missed it, but why are there two thresholds?

Author response:
$b_i$ is a fixed limit for the i-th constraint, while $\tau_c$ is the tolerance. The reason we use the tolerance $\tau_c$ here is for the convergence process when the approximated constraints $J_{g,i}(\pi)$ violate the fixed constraint limit $b_i$, necessitating the bounding of errors (cf. Equation 27 in the proof of Theorem 2).

Reviewer comment (Writing):
(Suggestion) Page 4, after equation 3: "\( \tilde \theta \) are the returns from the actor networks" could be re-worded. Returns have a specific meaning in RL literature, and in this context, I think return just means the output of the actor network, and not the usual return. (please ignore if my understanding is incorrect)

Author response:
Thanks for highlighting this. We use 'outputs' in the revised paper when conveying this concept.

Reviewer comment (Reproducibility):
Maybe I missed this, but how are the individual \( p(L_i|D) \) obtained (in the flight task setup)?

Author response:
The individual $p(L_i|D)$ represents the agent's capability to detect the factor $L_i$ within the environments $D$, typically acquired from sensors or perception modules. In the context of our flight task, $L_0$ represents external forces, and $D$ represents depth images sampled from a depth camera, namely the Intel Realsense D435i. Therefore, $p(L_i|D)$ individually represents the estimation of external forces $L_i$ within the environment $D,' and this estimation is obtained through VID-Fusion [1].  

Reference: 

[1] Ding, Ziming, et al. "Vid-fusion: Robust visual-inertial-dynamics odometry for accurate external force estimation." 2021 IEEE International Conference on Robotics and Automation (ICRA). IEEE, 2021.