Paper: S$2$AC: Energy-Based Reinforcement Learning with Stein Soft Actor Critic (ICLR-2024)

Reviewer comment (Novelty):
I believe that the authors missed an important related work [1]. There are lots of similarities between the proposed method and [1], e.g. using SVGD to approximate the energy-based policy parameterized by the $Q$-function. In fact, if [1] does not amortize the policy, I believe it is nearly identical to the proposed methods.

Author response:
**Reference [1]:**
  Thank you for bringing this to our attention. We actually intended to cite [1] in the SVGD-augmented RL paragraph (last paragraph in related work on page 9) but we ended up citing a different paper by referencing liu2017stein (identical first author, year, and first word of title). Lines 2-4 are actually comparing [1] to our work. Our work is fundamentally different. In [1], the authors learn a multi-modal distribution over the policy parameters ($\theta$) space, while we learn a single policy with multi-modal action ($a$) spaces. We argue that learning a multi-modal distribution over $\theta$ following [1] doesn't necessarily result in learning diverse policies as it's possible to learn the same mapping from states to actions using different parameterizations due to local-minima symmetries in the deepnet loss function. In contrast, learning multi-modal distributions over action spaces results in diverse actions by design. Besides, in real-world setups, we would like to follow the optimal policy and have the possibility to pick a suboptimal action in case of test-time perturbation, without completely switching to a different policy (with a different parametrization $\theta$). We fixed the citation.

Reviewer comment (Theory):
I feel there are several modifications that make the proposed method alleviating from the concept of SVGD, e.g. the truncation of the particle update. Furthermore, as there will be some discretization error from the gradient flow of KL on Wasserstein space (which is the motivation of SVGD), I'm thinking if the entropy of the discretized SVGD is a proper estimate of the KL on energy-based policy. I'm also wondering that given the $Q$-function, is it possible to directly estimate the entropy of the energy-based policy, which potentially gets rid of the other issues.

Author response:
**Particles truncation:**
  We did not implement the constraint in Eq.9 by truncating the particles to be within 3 std of the mean. Since we use the change of variable formula, we can only apply invertible transformations on the particles (truncation is not invertible). Instead, we sample more particles than we need to select the ones that stay within the range. Since the Q-landscapes become smoother as the training progresses (entropy maximization effect), there are always enough particles to satisfy the constraint. We added this clarification on page 9 under Eq 9.

**Estimating the entropy of EBMs:**
  Estimating the entropy of EBMs ($p(x) = \exp{E(x)}/Z$, $E(x) \in \mathbb{R}$ is an energy function) would require evaluating the intractable partition function $Z$ (integration over an infinite number of configurations of $x$). Hence, directly computing the entropy of EBMs is intractable. In Sec B.2 of the appendix, we review previous approaches approximating this quantity and elaborate on how they compare to our method.

Reviewer comment (Evaluation):
The experimental results in fact do not beat the baseline a lot.

Author response:
**MuJoCo experiments:**
  We updated our experiments. Please check the common reply.

Reviewer comment (Novelty):
The relationship between the proposed methods and the references I mentioned in the weakness part.

Author response:
**Reference [1]:**
  Thank you for bringing this to our attention. We actually intended to cite [1] in the SVGD-augmented RL paragraph (last paragraph in related work on page 9) but we ended up citing a different paper by referencing liu2017stein (identical first author, year, and first word of title). Lines 2-4 are actually comparing [1] to our work. Our work is fundamentally different. In [1], the authors learn a multi-modal distribution over the policy parameters ($\theta$) space, while we learn a single policy with multi-modal action ($a$) spaces. We argue that learning a multi-modal distribution over $\theta$ following [1] doesn't necessarily result in learning diverse policies as it's possible to learn the same mapping from states to actions using different parameterizations due to local-minima symmetries in the deepnet loss function. In contrast, learning multi-modal distributions over action spaces results in diverse actions by design. Besides, in real-world setups, we would like to follow the optimal policy and have the possibility to pick a suboptimal action in case of test-time perturbation, without completely switching to a different policy (with a different parametrization $\theta$). We fixed the citation.

Reviewer comment (Theory):
Is it possible to directly estimate the entropy of the energy-based policy?

Author response:
**Estimating the entropy of EBMs:**
  Estimating the entropy of EBMs ($p(x) = \exp{E(x)}/Z$, $E(x) \in \mathbb{R}$ is an energy function) would require evaluating the intractable partition function $Z$ (integration over an infinite number of configurations of $x$). Hence, directly computing the entropy of EBMs is intractable. In Sec B.2 of the appendix, we review previous approaches approximating this quantity and elaborate on how they compare to our method.

Reviewer comment (Theory):
Equation (2) seems to be inconsistent with the setting in SQL, as there is no discount factor in Equation (2). Is this a typo (as well as in Equation (3))?

Author response:
**Discount factor in Eq. 2:**  
   This is correct; there should be a discount factor. We added it. We originally followed the notation in the SAC paper which omits the discount factor. We agree that it’s more rigorous this way.

Reviewer comment (Theory):
Theorem 3.1 is not rigorous. My main concern is around the condition $\epsilon \lVert \nabla h \rVert_{\infty} \ll 1$. First, this is not a precise statement. Second, under this condition, Equation (10) cannot hold with equality (as already displayed in the appendix, there is an approximately equal argument). My suggestion is to either make a precise statement articulating the relationship between the condition and the final assertion to respect the rigor of a theorem, or make it casual by changing it to be a claim with approximate equality. After all, this theorem serves as a benign consequence of preserving the invertibility and as a motivation to properly choose $\sigma$.

Author response:
**Theorem 3.1:**  
   We agree. We changed the equality sign to an approximately equal sign and added the order of the approximation error $\mathcal{O}(\epsilon^2 Ld)$.

Reviewer comment (Writing):
In Proposition 3.2, it is better to recall the definition of $\sigma$, as there is an ambiguity of $\sigma$ being the variance of Gaussian or kernel function.

Author response:
**Definition of $\sigma$:**  
   We updated Proposition 3.2 to include the definition $\sigma$.

Reviewer comment (Theory):
I am not confident in the claim that "This empirically confirms that SGLD, DLD, and HMC update rules are not invertible" (beginning of page 7) can be obtained from the entropy is not accurately estimated. Can authors provide more context around it?

Author response:
**Sec 4.1:**  
   We replaced the sentence "This empirically confirms that SGLD, DLD, and HMC update rules are not invertible" with "This empirically **supports** Proposition 3.4 that SGLD, DLD, and HMC update rules are not invertible".

Reviewer comment (Writing):
Page 3. The notation \rho_\pi. Is it the stationary distribution of the state under \pi, or the state distribution s_t, which depends on the initial distribution? If it is the former, please clarify. If it is the latter, then \rho_\pi should depend on t.

Author response:
$\rho_\pi(s_t)$ is the state marginal induced by $\pi$ which does depend on the initial distribution. By convention (e.g., in the Sutton-Barto RL book and the SQL paper), we omitted the transition probability for simplicity.

Reviewer comment (Theory):
Page 3 basic setting. It seems that the authors are considering an RL problem with infinite horizon and without discount. With such a setting, it is not a trivial issue to guarantee the total reward is finite. However, the authors make no comments on this issue.

Author response:
Thanks. We added a discount factor. We originally followed the notation in the SAC paper which omits the discount factor. We agree that it's more rigorous this way.

Reviewer comment (Theory):
Page 3 eqn (4). -\alpha should be +\alpha?

Author response:
That is correct. We revised it.

Reviewer comment (Theory):
Page 4 bottom. In the change of variable formula, there should be an epsilon after I+?

Author response:
Thanks. We added it.

Reviewer comment (Theory):
Page 5 equation (9). The motivation for the actor update is to minimize the KL divergence between the policy distribution and EBM of Q-values, which is clear to me. But it is not obvious that minimizing this KL divergence is equivalent to (9). Can you also add a short proof in the appendix. Also, I think \mathcal{D} is used without definition. Q(a|s) looks like a typo, is it Q(a,s) or q(a|s)?

Author response:
We added the proof in Appendix E. We added the definition of $\mathcal{D}$ (replay buffer) and fixed the typo ($Q(s_t,a_t)$ instead of $Q(s_t|a_t)$).

Reviewer comment (Theory):
Page 5 theorem 3.1. Eqn (10) should not be equal, but an approximation. Maybe it is better to add the order of the approximation error. I think it is O(\epsilon^2 * L), provided sufficient regularity.

Author response:
Great point. We changed to using the approximation sign instead and added the order of the approximation error $\mathcal{O}(\epsilon^2 Ld)$. $L$ is the number of SVGD steps and $d$ is the action dimension.

Reviewer comment (Theory):
Page 6. Prop 3.5 says HMC is not invertible, but the following paragraph says, “While the HMC update is invertible”. Is it a contradiction?

Author response:
Thanks. We clarified. The HMC update rule is only invertible with respect to the $(a,v)$ (Neal et al. (2011)), i.e., when conditioning on $v$. Since $v$ is sampled from a random distribution, it has the effect of the noise variable in SGLD. Hence, HMC is not invertible w.r.t $a$.

Reviewer comment (Novelty):
Related work. I think the paper “Single Timescale Actor-Critic Method to Solve the Linear Quadratic Regulator with Convergence Guarantees” (JMLR 2023) could also be added to the related work. The LQR setting also has continuous state and action space, and the actor is a soft policy.

Author response:
Yes, we added the citation.

Reviewer comment (Writing):
Page 17. The first \epsilon should not appear in the second last line.

Author response:
Thanks. Corrected.

Reviewer comment (Experiments):
There are also variants of SAC for discrete settings. Do you think the S2AC would benefit in those scenarios?

Author response:
We appreciate the reviewer's highly positive comments.

**Generalizing $S^2AC$ to discrete settings**

Generalizing $S^2AC$ to discrete settings is non-trivial as SVGD is by default a sampler for continuous probabilities. Since the actions in discrete setups are **finite** and pre-defined, the policy is a softmax distribution with dimensions corresponding to the number of actions. Hence, the entropy can be obtained by computing the empirical mean of the negative log-likelihood of this distribution over all actions.

Reviewer comment (Experiments):
My biggest reservation for acceptance is significance. How much does the increased expressiveness of the stochastic policy (at the cost of computation) matter? The authors could only show significant differences to existing methods on hand-crafted toy problems. For the traditional benchmarks, i.e., Mujoco, the method does not appear to perform better than standard approaches such as PPO. The problem here is not the fact that it does not perform better, but that it appears better experiments could be designed to show how large/if there is a significant advantage.

Author response:
**Significance on complex environment**: Humanoid-v2 (Fig. 8-e) is a multi-modal environment (the action space is 17 dimensional with a large number of continuous control factors including combinations of legs and arms movements, posture, $\cdots$). Note that our approach (violet) is more sample efficient and performs better than SAC (red) due to the improved ability to explore different local optima. We thought about testing robustness by modifying the humanoid environment, for instance by freezing some variables in the action space to force the robot into a different forward movement mode. This is, however, non-trivial as it impacts the robot's balance. We agree, as RL algorithms are being deployed in real-world setups, ensuring the robustness of these agents is becoming increasingly important. We added a sentence in the conclusion to encourage the community to invest more in this direction.

Reviewer comment (Experiments):
If that is the argumentation, I would expect the authors to design experiments under these scenarios that are not toy. This could have been achieved, for instance, by augmenting the MuJoCo tasks with perturbation events that prevent execution of certain actions in certain situations. Then a more expressive stochastic policy could perform better. For now it makes it hard to assess if the increased time- and algorithmic-complexity is worth the cost of obtaining better performing policies.

Author response:
**Significance on complex environment**: Humanoid-v2 (Fig. 8-e) is a multi-modal environment (the action space is 17 dimensional with a large number of continuous control factors including combinations of legs and arms movements, posture, $\cdots$). Note that our approach (violet) is more sample efficient and performs better than SAC (red) due to the improved ability to explore different local optima. We thought about testing robustness by modifying the humanoid environment, for instance by freezing some variables in the action space to force the robot into a different forward movement mode. This is, however, non-trivial as it impacts the robot's balance. We agree, as RL algorithms are being deployed in real-world setups, ensuring the robustness of these agents is becoming increasingly important. We added a sentence in the conclusion to encourage the community to invest more in this direction.

Reviewer comment (Experiments):
How does the armotized version of the method perform on all benchmarks? Figure 9 (right) only shows one performance curve with no indication which benchmark it is.

Author response:
**Amortized $S^2AC$**

Results are added to Fig. 8 ($S^2AC(\phi,\theta,\psi)$). There is a small tradeoff in performance for 3 times faster inference compared to the non-amortized version ($S^2AC(\phi,\theta)$). Yet, the improvement over the baselines is still statistically significant (Fig.8 f-j). We added more details on the amortized version in Sec. 3.4. Note that the amortization is just a trick to improve the **test-time**. Specifically, amortized S2AC is just training a feedforward network $f_{\psi}(s,z)$ (Eq.24-25) to mimic the rolled-out (non-amortized policy) as test time (L.15, Alg2). **During training**, we need to use the unrolled policy to be able to compute the entropy estimate of the sampler in the MaxEnt RL objective (L.11-12 in Alg2). Depending on the application, it can be used as the main method, at test time, i.e., if timely reactivity is important (e.g., robot manipulation tasks). In other setups where the accuracy matters more (e.g., medical diagnosis), the unrolled policy would enable checking the Stein identity to assess convergence.

Reviewer comment (Presentation):
Because the armotized version is much faster and performs similarly, I wonder why it is not the main method, but its details are hidden in the appendix.

Author response:
**Amortized $S^2AC$**

Results are added to Fig. 8 ($S^2AC(\phi,\theta,\psi)$). There is a small tradeoff in performance for 3 times faster inference compared to the non-amortized version ($S^2AC(\phi,\theta)$). Yet, the improvement over the baselines is still statistically significant (Fig.8 f-j). We added more details on the amortized version in Sec. 3.4. Note that the amortization is just a trick to improve the **test-time**. Specifically, amortized S2AC is just training a feedforward network $f_{\psi}(s,z)$ (Eq.24-25) to mimic the rolled-out (non-amortized policy) as test time (L.15, Alg2). **During training**, we need to use the unrolled policy to be able to compute the entropy estimate of the sampler in the MaxEnt RL objective (L.11-12 in Alg2). Depending on the application, it can be used as the main method, at test time, i.e., if timely reactivity is important (e.g., robot manipulation tasks). In other setups where the accuracy matters more (e.g., medical diagnosis), the unrolled policy would enable checking the Stein identity to assess convergence.

Reviewer comment (Experiments):
I would be curious how this version performs on the toy problems (Fig. 5-7).

Author response:
**Results of Amortized $S^2AC$ on multi-goal** are reported in Fig.17.

Reviewer comment (Reproducibility):
The main limitation seems to be the slow inference of SVGD. However, the paper argues that this limitation can be addressed by amortization. It would be better if the main manuscript provided more details on how amortization is performed.

Author response:
**[Amortized S$^2$AC]**

We added more details in Sec 4.3 (run-time). The full description is as follows. Unfortunately, we cannot include all the details in the main manuscript because of the strict page limit, and it has been a tough decision to put them into the appendix.

Basically, we train a deepnet $f_{\psi}(s,z)$ to mimic the SVGD dynamics during testing, where $z$ is a random vector that allows mapping the same state to different particles. One way to train $f_{\psi}(s,z)$ is to run SVGD to convergence and train $f_{\psi}(s,z)$ to fit SVGD outputs. This, however, requires collecting a large training set of state action pairs by repeatedly deploying the policy. This might be slow and result in low coverage of the states that are rarely visited by the learned policy and hence lead to poor robustness in case of test time perturbations. We instead propose an incremental approach in which $\psi$ is iteratively adjusted so that the network output $a = f_{\psi}(s,z)$ changes along the Stein variational gradient direction that decreases the KL divergence between the policy and the EBM distribution, i.e.,

$$
\Delta f_{\psi}(z,s) = \frac{1}{m} \sum_{i=1}^{m} k(a_{i}, f_{\psi}(s,z)) \nabla_{a_i} Q(s,a_i) + \alpha \nabla_{a_i}k(a_i,f_{\psi}(s,z))  
$$

Note that $\Delta f_{\psi}$ is the optimal direction in the reproducing kernel Hilbert space, and is thus not strictly the gradient of Eq.(5), but it still serves a good approximation, i.e., $\frac{\partial J}{\partial a_t} \propto \Delta f_{\psi}$, as explained by Wang & Liu (2016). Thus, we can use the chain rule and backpropagate the Stein variational gradient into the policy network according to

$$
\frac{\partial J(s)}{\partial \psi} \propto E_{z} \left[ \Delta f_{\psi}(s,z)  \frac{\partial f_{\psi}(z,s)}{\partial \psi}  \right].
$$

to learn the optimal sampling network parameters $\psi^{*}$. Note that the amortized network takes advantage of a Q-value that estimates the expected future entropy which we compute via unrolling the SVGD steps. The modified \STAC\ algorithm is described in Algorithm 2.

Reviewer comment (Experiments):
It would be great if the paper mentions the scalability of SVGD with respect to dimensionality. How large the method can be scaled in terms of the dimensionality of the problem? Scaling SVGD to higher dimensional spaces seems to be challenging because we need exponentially more particles to represent a distribution.

Author response:
**[Scalability of SVGD with respect to dimensionality]**

Theoretically, SVGD does suffer from the curse of dimensionality for arbitrary distributions. But empirically, S$^2$AC($\phi$, $\theta$) works well on the MuJoCo benchmarks with high-dimension action spaces, e.g., the Humanoid environment with a 17-dimensional action space. This may depend on the specific shape of the distribution, and the fact that we have a **parameterized initialization** of the initial particle distribution has greatly mitigated the issue since it learns to focus on the sub-spaces for action/particle sampling. Notice the difference in performance between the parametrized S$^2$AC($\phi$, $\theta$) and non-parametrized S$^2$AC($\phi$) versions in Fig. 8. In the future, we aim to further improve the scalability by stacking multiple SVGD-induced distributions in a diffusion model style.

Reviewer comment (Novelty):
There are many works looking at multimodality in MaxEnt RL that were not discussed or cited. Moreover, several of these obtain better results than the current work (although they include additional features). For example, one earlier work is "Boosting trust region policy optimization with normalizing flows policy" (https://arxiv.org/pdf/1809.10326.pdf). Their equation (5) seems similar to equations (10) and (11) in the current paper (it's more clear from the proof sketch). Moreover, there are the recent works: "Reparameterized Policy Learning for Multimodal Trajectory Optimization" (https://arxiv.org/pdf/2307.10710.pdf), although it is model-based, and also the work: "Policy Gradient With Serial Markov Chain Reasoning" (https://openreview.net/forum?id=5VHK0q6Oo4M) that is model-free and obtains good performance (it's also based on the MaxEnt principle); this work also included a comparison with a normalizing flow in their appendix, and showed that their method achieves better performance; the experiments in this work were also more substantial than the current work. Finally, there are some other works that can be found in the related works section of this paper: "Leveraging exploration in off-policy algorithms via normalizing flows" (https://arxiv.org/pdf/1905.06893.pdf), "Iterative Amortized Policy Optimization" (https://arxiv.org/pdf/2010.10670.pdf). None of these works were cited, so I think the literature review was insufficient (they cite many works, but miss the most relevant ones).

Author response:
Thank you for sharing the references. We tried to conduct an extensive literature review, but these references were missed because their titles/abstracts do not indicate a strong relation to our work.

We carefully studied them. 

**TL;DR: MaxEnt RL is not just about multimodality**. First, we want to emphasize that our goal is to find a better solution to the MaxEnt RL objective due to its nice properties. Notably, beyond better exploration and multi-modality, MaxEnt RL is provably *robust*. This is achieved by maximizing the future entropy (positive reward setup) and is a highly desired property in real-world applications. Our empirical experiments (Fig. 6) validate that $S^2AC$ is not only multi-modal but also robust. This is, however, not the case for Ref [1-5]. Next, we argue that **Refs [1,2]** solve different objectives that lead to multi-modal policies but do not optimize for robustness. While **Refs [3,4,5]** optimize the MaxEnt RL objective, the learned policies fail to maximize the expected future entropy because of approximations (**Ref [3]**), the instability of training flow models (**Ref [4]**), or modeling the policy as uni-modal Gaussian (**Ref [5]**). We included **a comparison to these methods** in the intro and related work sections. We also **re-implemented Ref [4]** for comparison in both the Multi-goal and MuJoCo experiments (their code is not available). 

Detailed clarifications are as follows.

**Refs [1,2] optimize different objectives from MaxEnt RL.** **Ref [1]** optimizes Trusted Region Policy Optimization with a KL-divergence constraint. **Ref [2]** solves a new ELBO (Eq2) which is inspired but different from the MaxEnt RL objective. Specifically, it introduces additional latent variables and follows a model-based formulation.

**Refs [3,4,5] fail to maximize the future entropy.**
- **SSPG (Ref [3])** models the policy as a Markov chain with transition probabilities being parametrized Gaussians, which enables computing an approximation of the policy gradient (Eq. 13). The experimental results in Figure 4(A) don't demonstrate maximizing the future entropy. Specifically, *the middle plot in Fig. 4(A)* is an identical setup to the Multi-goal environment in Fig. 2 in our paper. As stated by the authors (lines 9-10 of Sec. 4.2), SSPG learns to visit the goals with *similar frequencies*. However, a good solution to the MaxEnt RL objective would be biased towards sampling from the lower modes. Hence, it seems that SSPG has just led to better exploration without robustness benefits. We wanted to further investigate, but the author didn't provide the code (empty link). In contrast, we show that explicitly computing the entropy leads to more optimal MaxEnt RL solutions (Fig. 2), *i.e., to maximizing the expected future entropy*.
- **SAC-NF (Ref [4])** models the policy as a normalizing flow (NF). **We re-implemented the paper** (their code was not provided) and reported results on both the Multi-goal (**Figs. 6-7**) and MuJoCo envs (**Fig. 8**). On Multi-goal, SAC-NF learns to maximize the entropy but, unlike our approach, it misses one mode for smaller $\alpha$ values. We argue that this is due to normalizing flow well-known tendency to quickly collapse to local optima Kobyzev et al. (2020) (in accordance with Fig. 4(A) in Ref[3]). The dispersion term in S2AC encodes an inductive bias to mitigate this issue. On MujoCo, convergence is slower and performance is lower than our approach. **Note that the reported results in Ref [3] are from applying NFs over REDQ (Table 5)**.
- **IAP (Ref [5])** models the policy as a uni-modal Gaussian (Sec. 4.1). Multi-modality is achieved by learning a collection of parameter estimates (mean, variance) through different initializations (Figure 2 and Sec. 3.2). The estimates result in parametrizing different policies. This doesn't result in maximizing the entropy and hence doesn't lead to robustness.

**Eq 5 in Ref [1]** is the Change of Variable Formula (CVF), which we use as a tool given that SVGD is invertible (similarly to normalizing flows). Our theoretical contribution consists of (1) proving that SVGD is invertible which enables the use of CVF, (2) using CVF together with the corollary of Jacobi’s formula to derive a closed-form expression of the SVGD log-likelihood (Eq 10), (3) deriving an equivalent computationally efficient expression of Eq. 10 (evaluates the trace of the Jacobian) that only depends on first order derivatives and vector dot products (Eq. 11) and (4) optimizing for scalability via parameterizing Eq. 11.

Reviewer comment (Experiments):
How does your method compare to other methods that create multi-modal policies in RL? Please include a comparison.

Author response:
Thank you for sharing the references. We tried to conduct an extensive literature review, but these references were missed because their titles/abstracts do not indicate a strong relation to our work.

We carefully studied them. 

**TL;DR: MaxEnt RL is not just about multimodality**. First, we want to emphasize that our goal is to find a better solution to the MaxEnt RL objective due to its nice properties. Notably, beyond better exploration and multi-modality, MaxEnt RL is provably *robust*. This is achieved by maximizing the future entropy (positive reward setup) and is a highly desired property in real-world applications. Our empirical experiments (Fig. 6) validate that $S^2AC$ is not only multi-modal but also robust. This is, however, not the case for Ref [1-5]. Next, we argue that **Refs [1,2]** solve different objectives that lead to multi-modal policies but do not optimize for robustness. While **Refs [3,4,5]** optimize the MaxEnt RL objective, the learned policies fail to maximize the expected future entropy because of approximations (**Ref [3]**), the instability of training flow models (**Ref [4]**), or modeling the policy as uni-modal Gaussian (**Ref [5]**). We included **a comparison to these methods** in the intro and related work sections. We also **re-implemented Ref [4]** for comparison in both the Multi-goal and MuJoCo experiments (their code is not available). 

Detailed clarifications are as follows.

**Refs [1,2] optimize different objectives from MaxEnt RL.** **Ref [1]** optimizes Trusted Region Policy Optimization with a KL-divergence constraint. **Ref [2]** solves a new ELBO (Eq2) which is inspired but different from the MaxEnt RL objective. Specifically, it introduces additional latent variables and follows a model-based formulation.

**Refs [3,4,5] fail to maximize the future entropy.**
- **SSPG (Ref [3])** models the policy as a Markov chain with transition probabilities being parametrized Gaussians, which enables computing an approximation of the policy gradient (Eq. 13). The experimental results in Figure 4(A) don't demonstrate maximizing the future entropy. Specifically, *the middle plot in Fig. 4(A)* is an identical setup to the Multi-goal environment in Fig. 2 in our paper. As stated by the authors (lines 9-10 of Sec. 4.2), SSPG learns to visit the goals with *similar frequencies*. However, a good solution to the MaxEnt RL objective would be biased towards sampling from the lower modes. Hence, it seems that SSPG has just led to better exploration without robustness benefits. We wanted to further investigate, but the author didn't provide the code (empty link). In contrast, we show that explicitly computing the entropy leads to more optimal MaxEnt RL solutions (Fig. 2), *i.e., to maximizing the expected future entropy*.
- **SAC-NF (Ref [4])** models the policy as a normalizing flow (NF). **We re-implemented the paper** (their code was not provided) and reported results on both the Multi-goal (**Figs. 6-7**) and MuJoCo envs (**Fig. 8**). On Multi-goal, SAC-NF learns to maximize the entropy but, unlike our approach, it misses one mode for smaller $\alpha$ values. We argue that this is due to normalizing flow well-known tendency to quickly collapse to local optima Kobyzev et al. (2020) (in accordance with Fig. 4(A) in Ref[3]). The dispersion term in S2AC encodes an inductive bias to mitigate this issue. On MujoCo, convergence is slower and performance is lower than our approach. **Note that the reported results in Ref [3] are from applying NFs over REDQ (Table 5)**.
- **IAP (Ref [5])** models the policy as a uni-modal Gaussian (Sec. 4.1). Multi-modality is achieved by learning a collection of parameter estimates (mean, variance) through different initializations (Figure 2 and Sec. 3.2). The estimates result in parametrizing different policies. This doesn't result in maximizing the entropy and hence doesn't lead to robustness.