Habitizing Diffusion Planning for Efficient and Effective Decision Making

Thank for reviewing our paper and acknowledging "idea is very straightforward and novel / easy to read / extensive experiments / reasonable ablation studies." We believe the following response can address your concerns.

Q1: Comparison with other generative decision-making baselines

Thank you for your suggestions. Since references 2 and 4 involve visuomotor policies and lack the same benchmarks, we compared our method with:

• Flow to Better: Offline preference-based reinforcement learning via preferred trajectory generation
• Boosting continuous control with consistency policy

The policy of FTB is also deterministic, which is the same as BC/SRPO.

We conducted comparisons on common MuJoCo locomotion tasks (since these work did not present the results on the other benchmarks in our paper, and FTB lacks the results on Medium dataset). Habi achieves the best performance in most tasks and demonstrates a higher overall score across environments.

Importantly, we do not want to convey the message that Habi (variational Bayesian method) is better than flow-based method. Indeed, they are orthogonal and can be combined (e.g., Kingma et al., 2016. "Improved Variational Inference with Inverse Autoregressive Flow"). We appreciate your question and will have a deeper dive into potential collaboration between flow-based method in our future work.

Q2: The proposed method requires a pre-trained diffusion model and also needs to train a critic - which requires more time in training than other "train-from-scratch" methods.

Thank you for your question. This work positions Habi as a general acceleration framework designed to speed up given diffusion planners. Our starting point is to maintain diffusion planning performance while significantly improving decision frequency.

Indeed, Habi's training process is light-weighted,e.g., Habi only takes around 2 hours for Habitual Training using one A100 GPU.

Q3: Though the critic can improve per-step quality, the overall episode return does not seem to be improved much and even degrades later.

There is probably a misunderstanding. We assume you are talking about Figure 7 (please let us know if not the case). Figure 7 illustrates the impact of increasing candidates on performance given the same critic. When there are too many candidates, performance may decline. Therefore, around 5 candidates is recommended. This phenomenon occurs because too many candidates amplify the critic's judgment errors for individual actions.

Q4: Why not report the computation time for BC and SRPO baselines?

We acknowledge that BC and SRPO are super fast, as they used deterministic policy function with single-pass inference (i.e., an MLP). We did not report the computational time since the bottleneck of inference time no longer come from action computation, but rather from other parts such as env.step(action). Thanks for the advice, and we will revise the paper to clarify this.

However, probabilistic generative models like diffusion policies and planners typically achieve superior performance by modeling complex, multi-modal action distributions -- but at significant computational cost. Our work focuses specifically on accelerating probabilistic generative decision-making methods while maintaining their performance advantages over deterministic approaches. Habi addresses this trade-off by accelerating probabilistic generative models while preserving their stochastic nature and performance benefits.

Thank you again for your constructive feedback! We hope the responses above have addressed all your comments. Please kindly let us know if you have additional suggestions, and we would be more than happy to discuss them.

We greatly appreciate your detailed and comprehensive comments, the following responses are to address your concerns.

Q1: Position Habi as an "acceleration technique" rather than a standalone algorithm.

Thank you for the clarification. Actually, we indeed position Habi as a general acceleration framework for diffusion planners in this paper. Our goal was creating an elegant framework that significantly improves decision frequency while maintaining effectiveness on SOTA diffusion planners through brief habitual training (1-2 hours).

Our focus is on decision-frequency, with performance comparisons to demonstrate Habi's effectiveness.

Q2: Comparisons with established distill methods (e.g., soft targets)

Diffusion models' probability distributions aren't directly accessible, making traditional distillation methods using cross-entropy loss (e.g., soft targets) inapplicable. This motivated Habi's development.

We included comprehensive comparisons with two categories of acceleration approaches: numerical acceleration methods (DiffuserLite) and distillation-like frameworks (DTQL). The Direct Distill baseline uses the same architecture as Habi but performs imitation learning directly on state-action pairs.

Q3: Performance improvements seem marginal for simple tasks (Maze2D, MuJoCo)

There is probably a misunderstanding. Our rigorous evaluation methodology (5 training seeds × 500 evaluation seeds) demonstrates these improvements are statistically significant and reproducible, which is supported by:

$\mu_{avg} = \frac{1}{G} \sum_{i=1}^G \mu_i$

$stderr_{avg} =\frac{\sqrt{\frac{1}{G}\sum_{i=1}^G(\mu_i^2 + N_{seed} (stderr_i)^2) - \mu_{total}^2}}{\sqrt{G*N{seed}}},$ where $G$ is the task number, $N_{seed}$ is the number of seed.

Performance Comparison:

Moreover, the improvements becomes more pronounced in complex environments: in Kitchen and Antmaze domains, Habi exhibits 19.1–29.4% stronger performance.

Q4: Necessity of Critic.

Thank you for your question. The critic is essential. Without it, performance in complex environments like Antmaze suffers an 18.5% degradation. Tested with over 500 seeds, the critic consistently provides performance improvements. Meanwhile, the critic is a shallow MLP that doesn't add significant computational burden.

Q5: Questions on network structures, diffusion planners, and latent dimensionality.

Thank you for your suggestion. For high decision speed, MLPs are already sufficiently simple. DV and DQL adequately represent the two types of diffusion planning. Regarding latent dimensionality, our experiments show:

Latent dimension has modest impact on performance. For Antmaze, there's slight improvement as dimension increases. For Maze2D, performance remains stable. As our method does not involve an information bottleneck, we selected dim(z)=256 as a balanced choice providing good performance without unnecessary computational overhead.

W1: Innovation & Q5: Difference between Habi and VAE

We understand your concerns. As acknowledged by other reviewers (qHBM, JDXx), Habi maintains simple and elegance for easy using, while Habi and VAE are different fundamentally:

Function: VAE's latent dimension is smaller than input dimension, forming an information bottleneck; Habi uses a large latent dimension (256) as it doesn't target compact representation.

Latent Bottleneck: VAE's latent dimension is smaller than input dimension, forming an information bottleneck; Habi uses a large latent dimension (256) as it doesn't target compact representation.

Learnable Prior: VAE employs a fixed unit-Gaussian prior; Habi's prior distribution conditions on the current state and is learned.

W2. biologically-inspired rather than technical innovation

Thanks for the suggestion. While the high-level idea of our paper is brain-inspired, we agree that we should more clearly highlight the technical innovations by revising the paper.

W3. Lack of real-world, vision-based deployment.

Thank you for the suggestion. We acknowledge it in Section 6 of our paper. This work focuses on algorithmic contributions using standard offline RL for robust benchmarking in the simulation, which is common practice in the field.

We appreciate your thoughtful comments and acknowledging "Simple yet effective approach / Well-motivated idea / Well written paper / elegant idea and strong performance". We will address all of your concerns as follows.

Q1: Does the expert/teacher planner have to be a diffusion-based model? What's the relationship between the proposed framework and diffusion-planners?

Good point! Theoretically, Habi also supports acceleration for other probabilistic generative decision-making models. However, Habi is especially well-suited for diffusion planners for several important reasons:

1. Many other generative models (like VAEs or flow-based models) are already computationally efficient, but often face performance limitations rather than speed constraints [1, 2]. Their bottleneck, however, is typically quality of decisions, not decision frequency.
2. Diffusion models have emerged as the most powerful and well-performing probabilistic generative models in decision-making, as evidenced by the impact of Diffuser and related work. However, diffusion planners are inherently slow during inference due to their multi-step denoising process, which limits their practical application in real-world scenarios.

Habi addresses this crucial gap between 1) and 2) by accelerating diffusion-based decision-making models while maintaining nearly lossless performance. To our best knowledge, Habi is the first framework to successfully preserve diffusion models' superior performance while achieving decision speeds comparable to traditional generative models.

We appreciate your suggestion and will clarify this important distinction in the revised paper.

Q2: The claim in L107-108 is a bit too strong: "Habi can be used straightforwardly for any diffusion planning and diffusion policy models." While it seems that in the experiments the authors have only implemented HI for one base planner (correct me if I misunderstood).

Actually, we used Habi on two base diffusion decision making models: [1] and [3] (line 738, Appendix Table 3). We selected these two models since they are best performing on corresponding benchmarks.

However, we greatly appreciate this suggestion since we did not realize this clarity issue. We will make more clear statements in the main texts of the revised paper.

Q3: Why is the framework using ztq instead of ztp for training the Critic?

The Prior distribution and Posterior distribution are constrained through KL divergence. Due to the nature of the KL constraint KL(q||p):

$KL(q(z) \parallel p(z)) = \int q(z) \log \frac{q(z)}{p(z)} dz$

The prior $p(z)$ is required to cover the posterior distribution $q(z)$, but points with low probability in the posterior may still be sampled by the prior. Using these low-probability samples to train the critic would introduce noise and potentially mislead the training process.

Therefore, using latents generated from the posterior can reduce the impact of these misleading samples on critic training, ensuring more reliable and accurate critic learning. Intuitively, this shares a roughly similar idea with Teacher Forcing in autoregressive models: During training, instead of feeding the model's own predicted output (akin to prior z in Habi), teacher forcing uses the ground-truth (akin to posterior z in Habit) to prevents compounding errors.

We also make an additional experiment on maze2d to empirically validate this design choice:

It can be seen that using posterior z for critic training works slightly better.

Thank you again for your acknowledgement of Habi and your contribution to our paper. Hope that we have address your concerns.

Reference

[1] What Makes a Good Diffusion Planner for Decision Making? Lu et al. ICLR 2025

[2] CleanDiffuser: An Easy-to-use Modularized Library for Diffusion Models in Decision Making, Dong et al. NeurIPS2024

[3] Diffusion policies as an expressive policy class for offline reinforcement learning, Wang et al. ICLR 2023

Thank you again for your constructive feedback! We hope the responses above have addressed all your comments. Please kindly let us know if you have additional suggestions, and we would be more than happy to discuss them.

We are grateful to your thoughtful comments and acknowledging our work "written very clearly / research question is very important / method makes sense / supported by extensive experimental validation / simple and elegant" as well as bringing the questions. We address your questions as follows.

Q1: GPU for decision-making speed tests.

Thanks for the suggestion. In fact, we already provided the results on both CPU and GPU in Appendix Table 4. We reported the results on CPU as CPUs are more portable (as descent CPUs can be put into mobile phones), thus be more practical when considering local deployment of decision models on individual robots and edge devices.

Q2: Decision frequency of DiffuserLite

We also noticed this. DiffuserLite serves as a key baseline for decision-making acceleration. We are unsure if absolute value differences stem from hardware variations, as we follow its official code for testing. In our experiments, we used consistent hardware (Apple M2 Max, Nvidia A100, AMD EPYC 7V13 64-Core) and software (PyTorch=2.2.2, numpy=1.22.4) to ensure fair and consistent speed test results. Therefore, the relative decision frequency presented in our paper could fairly reflects the differences in decision speed among various generative decision-making methods.

Q3: Relationship with world model and MBRL

Thank you for the suggestion. We treat Habi purely as an acceleration framework for generative decision-making. Its training process does not involve Q-Learning, nor does it interact with virtual reward dynamics. Habi's goal is to maintain the performance of the original generative model while accelerating decision-making within 1-2 hours of Habitual Training. Indeed, Habi is not limited to accelerating model-based planners (e.g., Diffuser, DV) but can also accelerate model-free planners (e.g., IDQL, DQL). We appreciate your suggestion and will discuss Habi's connection to offline MBRL methods (e.g., MOREC) in the related works section.

Q4: Can the ideas of habi be applied to more realistic robot control or robot control environments?

Yes, and thanks for suggesting extensive benchmarks. Beyond the Franka robotic arm manipulation tasks, we evaluate Habi on the Adroit environment (including opening the door, driving the nail, repositioning the pen orientation, and relocating the ball) for dexterous manipulation assessment.

The results demonstrate that Habi maintains comparable performance to diffusion planners across various robot control tasks. In some environments like Door, Hammer, and Kitchen-P, Habi even shows slight improvements. We believe this highlights the robustness of our simple-yet-effective method and its potential for real-world robotic applications.

Q5: Why is the diffusion planner algorithm rarely applied in real robotic control scenarios? As far as I know, most mainstream algorithms are based on diffusion policy.

We first would like to make a disambiguity statement about the "diffusion policy" in your question.

--If you meant the paper Diffusion policy: Visuomotor policy learning via action diffusion from Chi et al.

The premise of this question contains a misconception. Diffusion Policy is actually built upon Diffuser (a classic diffusion planner), with slight modifications to action modeling and visual inputs. Consequently, diffusion planner algorithms have indeed been deployed in robotics, with DP being among the first to bridge the ML-robotics gap and demonstrate results on physical robots.

Furthermore, the subsequent development directions of these two approaches differ: the robotics community typically emphasizes Imitation Learning with expert demonstrations, while the ML community focuses on Offline RL that can derive optimal policies from varied-quality data. The core algorithms, however, share the same diffusion-based foundation.

--If you meant the model-free diffusion policies such as Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning from Wang et al.

We believe one of the most crucial problems of diffusion planners are its heavy computation cost even on a descent GPU: As we have shown in Appendix Table 4, Diffuser take 0.3~0.6 second for one decsion, which is not acceptable to real-world robot. However, research shows diffusion planners can outperform diffusion policies on many tasks [1]. This is why our work focused on accelarating diffusion planning while maintaining its effectiveness.

[1] What Makes a Good Diffusion Planner for Decision Making? Lu et al. ICLR 2025