Maximum Entropy Reinforcement Learning via Energy-Based Normalizing Flow

We thank the reviewer’s time and effort spent on the review, and would like to respond to the reviewer’s questions as follows.

---

Comments

C1. Despite the fact that the proposed method is the first to apply EBFlow to reinforcement learning (RL), the novelty of the algorithm may appear weak, as there are several existing RL methods based on flow-based models. ... If such points exist, highlighting them would help to emphasize the contributions of the authors' proposed technique.

Response: As described in Lines 41-53 of Section 1 and Lines 166-174 of Section 3 of our paper, the proposed framework possesses the following two unique features. To the best of our knowledge, none of the existing MaxEnt RL frameworks, including those with flow-based models (e.g., [1-3]), possess the following characteristics:

• Our framework enables the exact calculation of the soft value function, which improves the accuracy of the optimization process.
• Our framework integrates the policy evaluation step and policy improvement step into a single training step, which streamlines the overall training process.

Previous actor-critic frameworks without these features may suffer from inaccurate soft value estimation (i.e., Eqs. (6) and (7)) and estimation error when minimizing the policy improvement loss (i.e., Eq. (5)). We substantiate our claim through an extensive set of experiments conducted on the MuJoCo and Omniverse Isaac Gym benchmarks as depicted in Fig. 3 and 4 of the main manuscript.

[1] Haarnoja et al. Latent Space Policies for Hierarchical Reinforcement Learning, ICML 2018.
[2] Mazoure et al. Leveraging exploration in off-policy algorithms via normalizing flows, CoRL 2019.
[3] Ward et al. Improving Exploration in Soft-Actor-Critic with Normalizing Flows Policies, ICML Workshop 2019.

---

C2. Introducing new frameworks generally necessitates providing sufficient mathematical proof to validate their efficacy, even for basic techniques or simple conjectures. ... However, the paper does not provide adequate mathematical justification for these techniques. Including such proofs would greatly enhance the confidence in the proposed method.

Response: We respectfully remind the reviewer that our method is mathematically supported by multiple theories.

Our main methodology is theoretically supported by Proposition 3.1 in Lines 181-183, which verifies the validity of our selections of $Q_\theta$ and $V_\theta$. According to Theorem 3 in [1], which is described in Section 2.1, the policy of MEow converges to the optimal policy when the soft Bellman error in Eq. (4) is minimized.

The theoretical justification for learnable reward shifting (LRS) is provided in Lines 212 to 219. Lines 214 and 215 show that the soft value function can be calculated exactly (Proposition A.3). In Lines 217 and 218, we mathematically demonstrate that the action distribution remains unchanged after applying LRS. The newly defined $Q_\theta^b$ and $V_\theta^b$ follow Theorem 3 in [1] and converge to the optimal policy when minimizing the soft Bellman error in Eq. (4).

The theoretical justification for SCDQ is given in Lines 226 to 233. SCDQ allows for the implementation of clipped double Q-learning [2] without duplicating the policy in MEow. In Eq. (13) and Lines 228-230, we mathematically show how SCDQ prevents the introduction of additional soft value functions. The mechanisms by which clipped double Q-learning reduces training variances and overestimation are analyzed in TD3’s original paper [2].

We would also like to emphasize that other techniques and claims in this work are theoretically supported with proofs. The deterministic inference technique is supported by Proposition A.4 in Lines 605-609. The theoretical analysis of the soft value estimation methods is offered in Proposition A.1 and Remark A.2 in Lines 565-568. These technical results support the validity of our method and provide motivation for formulating our approach.

[1] Haarnoja et al. Reinforcement Learning with Deep Energy-Based Policies. ICML 2017.
[2] Fujimoto et al. Addressing Function Approximation Error in Actor-Critic Methods. ICML 2018.

---

C3. It would be beneficial for the paper to compare its methods with more recent state-of-the-art algorithms. ... Such comparisons would greatly enhance the paper by placing the proposed method in the context of the latest advancements.

Response: A performance comparison between MEow and two latest online RL frameworks, DIPO [1] published in 2023 and S2AC [2] published in 2024, is presented in Table 1 in the global comment. The results show that MEow outperforms DIPO and S2AC with noticeable margins in five MuJoCo environments.

[1] Yang et al. Policy representation via diffusion probability model for reinforcement learning, 2023.
[2] Messaoud et al. S2AC: Energy-Based Reinforcement Learning with Stein Soft Actor-Critic, ICLR 2024.

---

Questions

Q1. According to the appendix, the proposed method is said to be 2.3 times slower than SAC. Given this difference in speed, a time-wise comparison of the algorithm's performance would be very helpful to understand its efficiency relative to SAC.

Response: In response to the reviewer's request, we extended the training time of SAC and provided a timewise comparison between MEow and SAC evaluated on the Hopper-v4 environment in Fig. 3 of the attached PDF file in the global comment. The results show that MEow converges to a policy that achieves a better average return and exhibits greater stability.

---

The response to Q2 is provided in the global comment.

We appreciate the reviewer’s valuable feedback and effort spent on the review, and would like to respond to the reviewer’s questions as follows.

---

Comments

C1. There are some unreasonable aspects in the organization of the content. For example, Section 3.2 appears somewhat abrupt; it is suggested to move the content of A.4 into the main text. The discussion about inference in Section A.3 is also very important and should be included in the main text.

Response: We appreciate the reviewer's suggestion and would be glad to incorporate the contents of Sections A.3 and A.4 into the main manuscript to improve the paper's quality. We will implement this arrangement in the final revision.

---

C2. The experimental section does not sufficiently highlight the advantages of the MEow method. The experiments in Section 4.1 on toy examples are not enough. Can comparisons be made from the perspectives of inference time, training time, representation capability, etc.?

Response: We appreciate the reviewer’s suggestion and have expanded the experiments in Section 4.1 to include a comparison of inference and training times in the Humanoid environment (MuJoCo benchmark), as shown in the following table. The results demonstrate that performing Monte Carlo estimation and MCMC sampling at each step imposes a computational burden during both the training and inference phases.

Table. Training and inference time comparison. $M$ denotes the number of samples used for soft value estimation and the number of steps in the MCMC sampling process. The results are evaluated on NVIDIA V100 GPUs.

Regarding the representation capability of our method, we have provided an analysis in Section A.5.4. The results indicate that our method possesses the ability to model multi-modal action distributions.

---

C3. While the integration of policy evaluation and improvement steps simplifies the training process conceptually, the implementation details of the EBFlow framework may introduce additional complexity that needs to be managed.

Response: The implementation of EBFlow can be simplified by integrating existing normalizing flow model libraries (e.g., [1,2]). Since log determinant calculation is already supported by many normalizing flow packages, the only required modification is the addition of an if-else condition during the forward pass to separate the energy function and the normalizing constant. We have made our code implementation available in an anonymous repository, with the corresponding link provided in Lines 699-700 of Section A.6.

[1] Stimper et al. normflows: A PyTorch Package for Normalizing Flows, Journal of Open Source Software 2023.
[2] Durkan et al. nflows: normalizing flows in PyTorch. 2020.

---

Questions

Q1. GFlowNet points out that the MaxEnt RL paradigm has deficiencies in terms of diversity. Specifically, MaxEnt RL tends to achieve target states that are "further" from the initial state. Does MEow also inherit this defect of the MaxEnt RL paradigm?

Response: We would like to highlight that the primary focus of this paper is on investigating the parameterization and its implications on MaxEnt RL. Comparing different RL frameworks (e.g., GFlowNet versus MaxEnt RL) falls outside the scope of this work. We acknowledge that further investigation into the topic of diversity within this context is an interesting direction for future research, and welcome any additional references the reviewer may provide to enrich this potential line of inquiry.

---

Q2. Energy-based normalizing flow is a generative model, and its original loss is to maximize log-likelihood. ... but the loss still follows the MaxEnt RL paradigm. In this case, could the title of this paper be potentially misleading?

Response: We believe that our paper title accurately reflects the key concept of the proposed methodology. As discussed in EBFlow’s original paper [1], a normalizing flow can be represented as an energy-based model and optimized using energy-based training objectives. Specifically, the authors in [1] employ score matching (which minimizes Fisher divergence) as an alternative to maximum likelihood training (which minimizes KL divergence) to optimize EBFlow. MEow can be viewed as an extension of EBFlow to the MaxEnt RL domain, where a policy is also represented as an energy-based model [2]. By synthesizing these concepts, we parameterize the policy using an EBFlow and demonstrate that the policy can be optimized by minimizing the soft Bellman errors.

[1] Chao et al. Training Energy-Based Normalizing Flow with Score-Matching Objectives. NeurIPS 2023.
[2] Haarnoja et al. Reinforcement Learning with Deep Energy-Based Policies. ICML 2017.

---

Q3. Figure A5 shows that the MEow network uses a hypernetwork module. Why was this module introduced? Is this a common architecture for energy-based normalizing flow?

Response: Hypernetworks are commonly employed in conditional normalizing flow models for modeling the weights in the transformations (see [1] for more details). In our approach, the policy is modeled as a state-conditioned normalizing flow, with hypernetworks utilized to encode state information. As for implementation, we adopt the existing implementation of conditional normalizing flow [2], which incorporates a hypernetwork architecture.

[1] Kobyzev et al. Normalizing Flows: An Introduction and Review of Current Methods. TPAMI 2019.
[2] Stimper et al. normflows: A PyTorch Package for Normalizing Flows, Journal of Open Source Software 2023.

---

The responses to Q4-Q6 are provided in the global comment.

We appreciate the reviewer’s valuable feedback and effort spent on the review and would like to respond to the reviewer’s questions as follows.

---

Comments

C1. (1) There is a potential concern about the empirical investigation: The proposed method is tuned over the target smoothing parameter, while the baselines are not tuned. (2) Why is the target smoothing parameter tuned per environment? What would happen if the proposed method also used the same, fixed smoothing parameter?

Response: (1) Our baselines use the refined hyperparameters of Stable Baseline 3 (SB3) as mentioned in Line 285. Further modifications to the hyperparameters may not result in performance improvement. For a fair evaluation, we provide the results of SAC (the best-performing baseline) tuned with different target smoothing factors used in MEow (i.e., $\tau=0.005$, $0.003$, $0.0005$, and $0.0001$) in Fig. 1 of the attached PDF file in the global comment. The results show that SB3’s original hyperparameter ($\tau=0.005$) is the most effective.

(2) We appreciate the reviewer's question regarding the target smoothing parameter $\tau$. According to Fig. 1 in the global comment, our current approach requires different values of $\tau$ to achieve good performance. However, we acknowledge the importance of determining a more generalized parameter setting. As part of our future work, we plan to expand our experimental framework to explore a broader range of $\tau$ values to identify a generally applicable $\tau$ across diverse environments.

---

C2. This is a minor point. The organization of the paper can be improved. Specifically, the background section is too long, which takes up the space for more interesting discussions presented in the appendix. If space allows, the main text should include the network architecture in Figure A5, as it may be unfamiliar to a lot of readers. In addition, the limitation, especially the computational aspect of the proposed approach, should be discussed in the main paper.

Response: We appreciate the reviewer's suggestion and concur that a more concise Section 2, coupled with the inclusion of discussions on MEow's model architecture and limitations in the main manuscript, will enhance the paper's quality. These modifications will be incorporated in the final revision.

---

C3. Other suggestions. (1) It would be great if there is at least some qualitative comparison between the proposed method and baselines using MaxEnt RL in the multi-goal environment. (2) Typos between Lines 254 and 276: The reference to Figure 1 should be to Figure 2.

Response: (1) We appreciate the reviewer's suggestion. A qualitative comparison among MEow, SAC, and SQL in the multi-goal environment is presented in Fig. 2 of the attached PDF file in the global comment. (2) We appreciate the reviewer’s attention to detail. The typo and figure reference will be corrected in the final revision.

---

C4. (1) The limitations of the work are discussed in the appendix, which should be moved to the main text. (2) Another unmentioned potential limitation is the sensitivity to the target network smoothing parameter.

Response: (1) We express our gratitude to the reviewer for the suggestions. A discussion of this work's limitations will be incorporated into the main manuscript in the final revision. (2) An examination of the sensitivity to the target network smoothing parameter is presented in Fig. 1 of the attachment to the global comment. This analysis will be included in the paper's final revision.

---

Questions

Q1. Are deterministic policies used for other algorithms in the evaluation in Figures 3 and 4?

Response: During the evaluation phase, MEow, SAC, TD3, DDPG, and PPO employ deterministic policies. SQL, however, does not support deterministic inference and thus utilizes a stochastic policy for inference.

We appreciate the reviewer’s valuable feedback and effort spent on the review and would like to respond to the reviewer’s questions as follows.

---

Questions

Q1. Typo at line 149, "Jocobian".

Response: Thank you for pointing this out. We will correct the typo in the final revision.

---

Q2. I am not sure if the role of the Learnable Reward Shifting (LRS) is clear to me. The authors say that they were inspired by the reward-shifting literature. However, I don't see the LRS term working the same way as classical reward shaping, as a state-dependent non-linear function ($b_\theta$), without any particular constraints (I don't see any stated) wouldn't preserve the properties of the original functions. I think the role of the learnable reward shifting is to have some form of "baseline" in the Q and V functions that simplifies the estimation process (since it's using a more flexible learning architecture), and so it avoids the explosion of the Jacobian. Could the authors provide clarification on this?

Response: Based on our formulation, the original $Q_\theta$ and $V_\theta$ functions are the residuals between their respective targets and $b_\theta$. We agree with the reviewer that this property of $b_\theta$ helps $Q_\theta$ and $V_\theta$ in predicting their targets during the optimization process.

Another important aspect of LRS is that it reduces the magnitude of the soft value function $V_\theta$, which makes our model less susceptible to numerical calculation errors. An analysis of this issue is provided in Section A4. This problem arises due to limited numerical precision. Consider a case where FP32 precision is in use, `MEow (Vanilla)’ could fail to learn a target $V_\{\theta^{\*}\} (s_t)>38, \forall s_t$, since $\exp(-V_\{\theta^{\*}\}(s_t))=\Pi_\{i \in S_l\} \| \det (J_\{g_\{\theta^{\*}\}^i\}(s_t)) |<2^{-126}$ cannot be represented using FP32 precision. Therefore, without shifting the reward function, the loss sometimes becomes undefined values, and can lead to ineffective training (e.g., the green lines in Fig. 6). The reward shifting term can be designed as a (state-conditioned) function or a (non-state-conditioned) value. It can also be learnable or non-learnable. All of these designs (i.e., $b_\{\theta\} (s_t)$, $b$s_t$$, $b_\{\theta\}$, and $b$) can be directly applied to MEow since none of them influences the action distribution. Based on our preliminary experiments, we identified that a learnable state-conditioned reward shifting worked the best and proposed this technique.

---

Q3. I would be interested in seeing a comparison between the learned reward shifting term and the value function (V) in a plot.

Response: We appreciate the reviewer's suggestion. The plot is offered in Fig. 4 of the attached PDF file in the global comment.

We appreciate the reviewer’s valuable feedback and effort spent on the review and would like to respond to the reviewer’s questions as follows.

---

Comments

C1. (1) The calculation of the determinant is usually time-consuming. (2) So the reviewer wonders how the training time of MEow is compared with that of the classical off-policy RL methods like SAC.

Response: (1) Several existing normalizing flow architectures (e.g., [1] used in this work) support efficient calculation of Jacobian determinants, as discussed in Section 2.2 (Lines 138-140). In general, the Jacobian determinant of these architectures can be calculated with time complexity $O(D^2L)$, where $D$ represents the dimension of actions and $L$ is the number of layers in $g_\theta$ (defined in Line 132). This complexity is the same as forward passing a normalizing flow model.

(2) A discussion regarding the training time of MEow and SAC is offered in Section A.7. Although the training time of our method is longer than that of SAC, SAC can not model multi-modal action distribution. For a fair comparison, we compared the training time of SAC-Flow (i.e., [2]) and MEow. We found that the training time of SAC-Flow is 1.13x slower than MEow due to the additional policy improvement updates (i.e., Eq. (5)).

[1] Dinh et al. NICE: Non-linear Independent Components Estimation, ICLR Workshop 2015.
[2] Haarnoja et al. Latent Space Policies for Hierarchical Reinforcement Learning, ICML 2018.

---

C2. In the experiments, MEow is just compared with several classical RL methods for continuous control and lacks the comparison with advanced RL methods [1, 2] with other generative models like diffusion model and consistency model, which also supports the modeling of multi-modal action distribution.

Response: We appreciate the reviewer’s suggestion and have included a performance comparison between MEow and DIPO [1] in Table 1 in the global comment. The results show that MEow outperforms DIPO, even when DIPO uses 100 sequential forward passes during action sampling.

To the best of our knowledge, consistency models have not been applied to online RL setups. In addition, combining consistency models with online MaxEnt RL may be challenging due to the intractability of the entropy calculation.

On the other hand, reference [2] addresses offline RL tasks, which differ from the online RL setting discussed in this paper. The primary distinction is that in offline RL, agents cannot interact with the environment during training. Therefore, reference [2] represents a separate research direction and is distinct from the main focus of this paper.

Finally, we would like to highlight that several online MaxEnt RL methods developed for modeling multi-modal policies are discussed in Section 2.2 and compared in Section 4.2. We use a number of representative MaxEnt RL methods, such as SQL, as our baselines in Fig. 3. We also compare our method with more advanced variants like SAC-Flow (i.e., [3]), denoted as ECFA in Fig. A3. Furthermore, Table 1 in the global comment demonstrates that our method outperforms the latest MaxEnt online RL framework (i.e., S2AC [4]), where S2AC also supports modeling multi-modal action distributions.

[1] Yang et al. Policy representation via diffusion probability model for reinforcement learning, 2023.
[2] Yue et al. Boosting offline reinforcement learning via data rebalancing, 2022.
[3] Haarnoja et al. Latent Space Policies for Hierarchical Reinforcement Learning, ICML 2018.
[4] Messaoud et al. S2AC: Energy-Based Reinforcement Learning with Stein Soft Actor-Critic, ICLR 2024.

---

Questions

Q1. How is the training time of MEow compared with other methods?

Response: The training time comparison is discussed in the response of C1.

---

Q2. How is MEow compared with other advanced RL methods with the utilization of generative-model-based policy?

Response: Please refer to the responses to C2 and Table 1 in the global comment.

---

Q3. Could the authors explain why MEow can learn a good reward shifting rather than a trivial shifting, like for any $s_t$, $b(s_t)=0$?

Response: We appreciate the reviewer's question regarding MEow's ability to learn a good reward shifting rather than a trivial one. The inclusion of a learnable reward shifting term serves a crucial purpose: it prevents numerical errors in the Jacobian determinant calculations, as discussed in Section 3.2 and analyzed in Appendix A4. The efficacy of this approach is demonstrated in Fig. 1 of the main manuscript, where we observe the successful resolution of these numerical issues. Furthermore, Fig. 6 in the manuscript illustrates that this method contributes to improved overall performance. These results collectively indicate that MEow learns a meaningful and beneficial reward shifting, rather than a trivial one.

In addition, according to our observation, $b_\theta$ does not learn to be a constant. Its value changes according to different visited states. To verify this, we sample a trajectory of states and plot the corresponding values of $b_\theta$ for the Hopper-v4 environment during the inference phase in Fig. 5 of the attached PDF file in the global comment.