Energy-Weighted Flow Matching for Offline Reinforcement Learning

Thank you for your feedback. Here we address your questions.

Q1. What’s the advantage of the proposed method compared with SRPO and SfBC?

A1. The advantage of our proposed method over SRPO and SfBC lies in its formulation as an energy-guided generative model that samples from the distribution $q(x) \propto p(x) \exp(-\beta E(x))$. In contrast:

• SRPO uses the diffusion model primarily as a regularization mechanism for policy optimization, focusing on maximizing the $Q$ function (which acts as the energy function in this context).
• SfBC employs the diffusion model as a data augmentation tool during the planning phase.

Both SRPO and SfBC are limited in that they cannot be extended to more general settings where both maximizing the $Q$ function and aligning with the data distribution are equally important. Our method, however, effectively balances these two objectives, making it suitable for a wider range of generative tasks.

Q2. Can this method be extended to more general tasks?

A2. We appreciate your insights on the potential extensions of our algorithm beyond RL. In response, we conducted experiments on image synthesis and molecular generation, following established benchmarks (DDPM [1], EDM[2]). We have included results in Appendix F.3 and F.4. These results indicate that energy-weighted diffusion can be readily applied to guide complex generative tasks, where standard RL algorithms may not be effective.

Q3. How the div operator is eliminated

A3. Apologies for the typo. We intended to write $\dot{q_t}(x) = -\mathrm{div} \cdot [q_t(x) \hat{u}_t(x)]$. This has been corrected in the revised manuscript.

[1] Ho, Jonathan, Ajay Jain, and Pieter Abbeel. "Denoising diffusion probabilistic models." Advances in neural information processing systems 33 (2020): 6840-6851.

[2] Hoogeboom, Emiel, et al. "Equivariant diffusion for molecule generation in 3d." International conference on machine learning. PMLR, 2022.

Thank you for your feedback and for updating your recommendation!

Regarding your question on whether the energy is being maximized or minimized, we would like to clarify that the difference arises from the specific definitions of the target distribution in different settings.

In the general setting discussed in Section 4, we define the target distribution as $q(x) \propto p(x) \exp(-\beta E(x))$ . Here, when the energy function is defined as the hue value distance, the energy guidance will direct the generative model to find images with smaller hue distances to the target color.

In contrast, in RL, the target policy is defined as $\pi(a | s) \propto \mu(a | s) \exp(\beta Q(s, a))$ . In this case, the reviewer is correct that the guidance seeks to maximize the $Q$-function. However, this difference is merely a matter of the sign of the energy term: $-\beta E(x)$ corresponds to minimizing energy, while $\beta E(x)$ corresponds to maximizing energy. The only thing changed is the sign before the energy function used in CWD training process.

To address this explicitly, we have clarified in the newly uploaded revision that for each application—molecular generation and image synthesis—the energy function is defined consistently with the target distribution. Specifically:

• Image generation: The energy function is defined as the hue distance, which is minimized via $q(x) \propto p(x) \exp(-\beta E(x))$ to guide the generation to a certain color.
• Molecular conformation synthesis: The energy function corresponds to desired molecular properties, which are maximized via $q(x) \propto p(x) \exp(\beta E(x))$.

We hope this explanation addresses your question.

Thank you for your strong support for our paper! We answered your questions as follows:

Q1. How is the variance introduced by the importance sampling?

A1. We acknowledge that using importance sampling to sample from $q(x) \propto p(x) \exp(E(x))$ [1, 2] might introduce additional variance. However, since we only use importance sampling to construct the loss function, we have found empirically that the negative impact of this variance is minimal. Moreover, we want to highlight that the batch size we use for training is relatively small and consistent with the baseline algorithms (64 for both image synthesis and molecular structure generation). Therefore, we do not need to increase the batch size specifically to control variance.

Q2. What are the trade offs that would inform choosing this method over others?

A2. We highlight the trade offs that would inform choosing our methods over the others as follows: our method is the first to directly learn exact energy-guided flow matching or diffusion models without relying on auxiliary models. In contrast, CEP [3] not only requires importance sampling to train its intermediate energy function $E_t(x)$ through contrastive learning (as shown in Eq. (13) of [3]), but it also necessitates additional back-propagation on its auxiliary model $E_t(x)$ . Moreover, methods like CG [4] and CFG [5] lack guarantees for generating the exact distribution, as discussed in Section 4.3.

Q3. Can this method be extended to more general tasks?

A3. We appreciate your insights on the potential extensions of our algorithm beyond RL. In response, we conducted experiments on image synthesis and molecular generation, following established benchmarks (DDPM [6], EDM [7]). Due to time constraints during the rebuttal period, we have included preliminary results in Appendix F.3 and F.4. These results indicate that energy-weighted diffusion can be readily applied to guide complex generative tasks, where standard RL algorithms may not be effective.

Q4. What’s the performance of the algorithm for a different energy landscape?

A4. We acknowledge that when the energy function is highly peaky, it becomes challenging to learn and sample effectively. This challenge is well-known in fields like non-convex optimization and Langevin dynamics. In response to your question, we modified the energy function in the 8-Gaussian example and presented the results in Appendix F.1. The results indicate that even with a changed energy landscape, the energy-weighted method is still able to generate samples effectively. Additionally, we note that the energy function used in the existing "chessboard" examples (Figures 2 and 3) is also non-convex with four distinct peaks.

In summary, while we acknowledge that learning an “extremely” peaky energy function is inherently difficult, our energy-weighted methods have demonstrated robust performance across different energy landscapes, including challenging ones.

[1] Cremer, Julian, et al. "Multi-Objective Guidance via Importance Sampling for Target-Aware Diffusion-based De Novo Ligand Generation." ICML'24 Workshop ML for Life and Material Science: From Theory to Industry Applications.

[2] Chen, Huayu, et al. "Offline Reinforcement Learning via High-Fidelity Generative Behavior Modeling." The Eleventh International Conference on Learning Representations.

[3] Lu, Cheng, et al. "Contrastive energy prediction for exact energy-guided diffusion sampling in offline reinforcement learning." International Conference on Machine Learning. PMLR, 2023.

[4] Dhariwal, Prafulla, and Alexander Nichol. "Diffusion models beat gans on image synthesis." Advances in neural information processing systems 34 (2021): 8780-8794.

[5] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance." arXiv preprint arXiv:2207.12598 (2022).

[6] Ho, Jonathan, Ajay Jain, and Pieter Abbeel. "Denoising diffusion probabilistic models." Advances in neural information processing systems 33 (2020): 6840-6851.

[7] Hoogeboom, Emiel, et al. "Equivariant diffusion for molecule generation in 3d." International conference on machine learning. PMLR, 2022.

Q5. What if the data distribution is not the same as the prior distribution?

A5. Similar to diffusion models, our method assumes that the prior distribution matches the data distribution. We acknowledge that the algorithm’s performance may be impacted by irregular data present in the dataset or from a pre-trained model. There is a body of work [5, 6] aimed at filtering such irregular data to improve training outcomes. However, we emphasize that addressing irregular data is not the primary focus of this work. We are confident that these existing methods can be integrated into our framework to handle this challenge effectively.

Q6. What is the effect of estimating an expectation with a minibatch rather than a universal constant across the dataset?

A6. We acknowledge that using a global $\mathbb{E}_{p(x)}[\exp(\beta E(x))]$ as a constant could potentially reduce the variance compared to estimating the expectation over a minibatch. However, our choice to estimate the expectation using a minibatch was made to ensure that the proposed algorithm remains flexible and easily adaptable to different tasks, with minimal hyperparameter tuning or dataset preprocessing. Empirically, we have observed that the behavior of these two methods is quite similar, suggesting that the impact of estimating the expectation using a minibatch is not significant in practice.

[8] Lin, Guang, et al. "Robust Diffusion Models for Adversarial Purification." arXiv preprint arXiv:2403.16067 (2024).

[9] Xiao, Chaowei, et al. "Densepure: Understanding diffusion models for adversarial robustness." The Eleventh International Conference on Learning Representations. 2023.

Thank you for your feedback and for acknowledging our response! We appreciate you pointing out a recent related work. We have included it as part of the discussion on related sampling methods in the revision. Additionally, we would like to provide a more detailed comparison here.

Specifically, [1] addresses sampling from the Boltzmann distribution $p_{\text{target}}$ using a tractable initial distribution $\mu$ in Section 3, and it explored credit assignment via the energy function gradient $\nabla_x \mathcal{E}(x)$ in Section 4.1. In comparison, [1] requires computing the gradient $\nabla_x \mathcal{E}(x)$ through back-propagation for the inductive bias. Therefore it will be slower when evaluating $u(x)$ as defined in Equation (14), similar with the CEP [2]. In our QIPO algorithm, the $u(x)$ is directly learnt from the Q-weighted loss function so the evaluation time should be significantly reduced (see Table 3). Therefore we claimed that our method is the first that directly (without any additional gradient requirements) learn the energy-guided flow.

We appreciate your comment again and hope our response have addressed all your questions!

[1] Improved off-policy training of diffusion samplers, Marcin Sendera, Minsu Kim, Sarthak Mittal, Pablo Lemos, Luca Scimeca, Jarrid Rector-Brooks, Alexandre Adam, Yoshua Bengio, Nikolay Malkin, NeurIPS 2024

[2] Lu, Cheng, et al. "Contrastive energy prediction for exact energy-guided diffusion sampling in offline reinforcement learning." International Conference on Machine Learning. PMLR, 2023.

Thank you for your feedback. We address your questions below.

Q1. What is the contribution of this paper compared with reward-weighted regression EDP and QVPO?

A1. Thank you for pointing out these references. We have included a detailed discussion of these papers in Appendix A. of the revised manuscript. Although the energy-weighted flow matching (and diffusion) loss appears similar to the reward-weighted regression in EDP or QVPO, they are fundamentally different for several reasons:

• Motivation: The energy-weighted diffusion model is derived from controlling the generation process of $q(x) \propto p(x) \exp(-\beta E(x))$, where both the original distribution $p(x)$ and the energy function $\beta E(x)$ are equally important. On the other hand, reward-weighted (policy) regression is formulated as minimizing the loss $\mathbb{E}[f(Q(s, a)) \log \pi_{\theta}(a | s)]$, as discussed in EDP [1]. The primary focus in RL is to maximize the $Q$ function, which differs significantly from generative modeling tasks where the goal is to generate high-quality samples while maintaining low energy. This distinction separates RL objectives from generative modeling objectives.
• Theoretical Insights: Building on our motivation, we provide theoretical guarantees that energy-weighted flow matching can accurately learn the probability distribution $q(x) \propto p(x) \exp(-\beta E(x))$ (in RL is $\pi(a | s) \propto \mu(a | s) \exp(\beta Q(s, a))$. In contrast, neither EDP [1] nor QVPO [2] can precisely learn this energy-guided distribution or provide such theoretical guarantees, as detailed below.
• Methodology: Although weighted regression is commonly used in literature, the Energy-Weighted diffusion differs from EDP and QVPO and provide unique perceptive:
  • Energy-Weighted Diffusion (ours) uses weights of the form $\exp(\beta Q(s, a)) / \mathbb{E}_{a \in \mu(\cdot | s)} \exp(\beta Q(s, a))$, which are guaranteed to generate $q(x) \propto p(x) \exp(-\beta E(x))$.
  • EDP uses a more artifically defined weight, $f(Q(s, a))$. Specifically,
    • $f_{\text{CRR}}(Q(s, a)) = \exp\left(\beta Q(s, a) - \mathbb{E}_{a \in \hat{\pi}} Q(s, a)\right)$, where $\hat{\pi} = \mathcal{N}(\hat{a}_0, I)$ and $\hat{a}_0$ is evaluated using the score network. This formulation oversimplifies the behavioral policy as a Gaussian distribution, making it different from the exact energy-guided distribution. Our approach, in contrast, directly samples from the behavioral policy $\mu(a | s)$.
    • $f_{\text{IQL}}(Q(s, a)) = \exp(Q(s, a) - V(s))$ suffers from a similar issue, as it replaces $\mathbb{E}_{a \in \hat{\pi}} Q(s, a)$ with $V(s)$. This difference prevents it from achieving the exact energy-guided distribution, and IQL also requires a pre-trained value function.
  • QVPO employs a linear weight $Q(s, a)$. Its goal, as discussed by the authors, is to encourage exploration in online RL, where exponential weighting is more suitable for offline RL.
  • As the theoretical analysis presented in Theorem 3.3 and Theorem 4.3, neither EDP nor QVPO can strictly recover the distribution $q(x) \propto p(x) \exp(-\beta E(x))$. While these algorithms improve $Q$ values, this is not the sole objective in generative modeling.
• Objective: Our methods balance alignment with the original policy $p(x)$ and energy guidance $E(x)$. This balance makes our approach applicable to diverse tasks, such as image synthesis and molecular structure generation, as demonstrated in Appendix F.3 and F.4 of the revision. In these tasks, simply maximizing the reward (e.g., generating a pure blue image) is inadequate; the generative model must align with the underlying dataset distribution.
• Empirical Performance: In comparison with EDP in the offline RL setting, our algorithm achieves superior performance, as presented in the table below. This highlights the importance of the theoretical guarantee on the distribution and the contribution of our paper.

[1] Kang B, Ma X, Du C, et al. Efficient diffusion policies for offline reinforcement learning[J]. Advances in Neural Information Processing Systems, 2024, 36.

[2] Ding S, Hu K, Zhang Z, et al. Diffusion-based Reinforcement Learning via Q-weighted Variational Policy Optimization[J]. arXiv preprint arXiv:2405.16173, 2024.

[3] Peters J, Schaal S. Reinforcement learning by reward-weighted regression for operational space control[C]//Proceedings of the 24th international conference on Machine learning. 2007: 745-750.

Q2. What’s the motivation for introducing flow matching to RL?

A2. Flow matching serves as a generalized framework for diffusion models, offering several advantages, including simulation-free implementation and faster convergence compared to diffusion models. These benefits have been demonstrated in our experiments (see Table 3 in our paper). We believe that flow matching provides a more general framework for studying the generative models in RL and also provides a faster convergence rate when generating new actions.

Additionally, we want to emphasize that flow matching has been successfully developed in various fields, such as image synthesis [4] and molecular structure generation [5]. Notably, the recent work on $\mathbf{\pi_0}$ general robotics control [6] utilizes flow matching to manage real-world robotic tasks. We believe this example illustrates the practical significance of flow matching in RL, as it demonstrates the potential to improve control and performance in complex environments.

[4] Lipman, Yaron, et al. "Flow Matching for Generative Modeling." The Eleventh International Conference on Learning Representations. [5] Song, Yuxuan, et al. "Equivariant flow matching with hybrid probability transport for 3d molecule generation." Advances in Neural Information Processing Systems 36 (2024). [6] Black, Kevin, et al. "$\pi_0$: A Vision-Language-Action Flow Model for General Robot Control." arXiv preprint arXiv:2410.24164 (2024).