S$2$AC: Energy-Based Reinforcement Learning with Stein Soft Actor Critic

• 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.

• 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.

• MuJoCo experiments: We updated our experiments. Please check the common reply.

1. Technical Issues

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.

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)$.

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

2. Claims in empirical results

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".

Sec 4.1: We agree. We replaced "maximizing the future reward and maximizing the future entropy" with "maximizing the sum of the future reward and future entropy" as suggested.

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.

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

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.

Dear Reviewer 4hdR,

Thank you for updating your recommendation to "accept". Could you please adjust your score to reflect your recommendation? We are eager to engage in further discussions and address any additional concerns you may have.

Best regards

We really appreciate the reviewer for checking our notations and mathematical soundness in great detail. We are delighted to improve the clarity of our paper based on your feedback.

1. $\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.

2. 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.

3. We added the proof to Appendix D. On a high level, given the definition of the soft Q-function (Eq. 17), we introduce the Lagrangian of the MaxEnt objective (Eq. 16) by incorporating the constraint $\int_a \pi(a|s)=1$. Computing the derivative w.r.t $\pi(a|s)$ for a given action $a$ and state $s$ results in the exponential policy in Eq. 19. Finally, we choose the Lagrange multiplier such that the obtained distribution sums to one (Eq. 20). We obtain $\pi(a|s)= \exp{( \frac{1}{\alpha} Q(s,a) )}/Z$. $\alpha$ was previously hidden in $Z$. We revised it to reflect the dependency on $\alpha$ in Eq. 3.

4. That is correct. We revised it.

5. Thanks. We added it.

6. 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)$).

7. 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.

8. 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$.

9. Yes, we added the citation.

10. Thanks. Corrected.

We appreciate the reviewer's highly positive comments.

Literature review of entropy estimation

In Appendix B, we reviewed (1) a non-parametric entropy estimator that leverages samples from a distribution without access to the normalized density in (Sec B.1) as well as (2) entropy estimators for EBMs (under knowledge of the unnormalized density). In summary, previous techniques are either based on heuristic approximation (e.g., computing pairwise distance between samples as a proxy for smoothness), lower bounds, or neural estimators of mutual information.

Differently, our approach is: (1) more principled as it computes the entropy of the actor (as opposed to using off-the-shelf generic tricks) by exploiting the invertibility of the SVGD update and the change of variable formula, (2) more parameter-efficient (we only learn a Gaussian initial distribution).

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.

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.

Experiments. We updated our experiments. Please check the common reply.

• [1] Tang, Yunhao et al. ICLR19

• [2] Huang, Zhiao, et al. 2023.

• [3] Cetin, Edoardo, et al. NeurIPS, 2022.

• [4] Mazoure, Bogdan, et al. CoRL, 2020.

• [5] Marino, Joseph, et al. NeurIPS, 2021.

[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.,

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

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.

[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.