Local Manifold Approximation and Projection for Manifold-Aware Diffusion Planning

Thank you for the detailed review and constructive feedback on our work. We especially appreciate the reviewer pointing out the clarity issue regarding our claims. Please find our detailed response below.

• “The claim "The current sample is then projected onto this subspace, thereby remaining on the manifold of valid behaviors." at line 69 seems not precise.”

We fully agree with the reviewer that our LoMAP method does not necessarily guarantee the validity of behaviors generated by diffusion models. Accordingly, we will tone down this claim to avoid ambiguity and overstatement. Specifically, we will revise:

"thereby remaining on the manifold of valid behaviors."

to a more accurate and precise description:

"thereby significantly reducing off-manifold deviations and improving the likelihood of generating valid behaviors."

• “The claim "Consequently, combining LoMAP with minority guidance can help uncover feasible yet unexplored solutions that might otherwise remain inaccessible to standard diffusion planners." at line 403 seems not precise.”

We apologize for the confusion caused by this statement. Our original intention was to highlight that standard diffusion planners typically require a large number of samples to generate trajectories located in low-density regions of the data distribution, thus limiting exploration within a restricted sampling budget.

We aimed to clarify that one way to overcome this limitation—though not our contribution—is through minority guidance [1], which explicitly guides sampling towards low-density regions. Our contribution, LoMAP, further enhances this process by ensuring that trajectories sampled under minority guidance remain feasible and closer to the data manifold, allowing efficient generation of trajectories from low-density regions without extensive sampling.

We appreciate the reviewer highlighting this ambiguity and will clarify this point carefully in the final version of the paper.

• “what is $n$ in line 161 defined?”

We apologize for this confusion. This was a typo. To clarify, the original sentence:

"… $\boldsymbol{\tau}^i$ is inherently concentrated around a ($n-d$) dimensional manifold $\mathcal{M}_i$."

should be corrected to the following description:

"… $\boldsymbol{\tau}^i$ is inherently concentrated around a ($d-k$)-dimensional manifold $\mathcal{M}_i$."

Thank you for pointing out this typo.

• “In line 225 the paper claims that Prop. 3.2 indicates that the manifold dev. problem is stronger for long-horizon tasks. How should I infer this from the mathematical result within the proposition?”

Proposition 3.2 states that the guidance gap scales at least proportionally to the square root of the dimensionality $d$ of the trajectory representation:

$\Delta_{\mathrm{guidance}}\bigl(\boldsymbol{\tau}^i\bigr) \ge \frac{\,c\,}{\sqrt{1-\alpha_i}}\sqrt{d}$

We clarify that the dimensionality $d$ here directly relates to the length of the planning horizon. Specifically, in diffusion-based planning, trajectories are represented by concatenating states and actions across the entire planning horizon. Therefore, the dimension $d$ grows linearly with the length of the planning horizon. Consequently, the lower bound provided by Proposition 3.2 indicates that as the planning horizon increases, the guidance gap inevitably grows, which implies a higher likelihood and severity of manifold deviation in longer-horizon tasks.

• “My main concern regarding the paper regards the risk of ending up obtaining a sampler excessively regularized from offline data as explained within my point (2) of Claims And Evidence section.”

Thank you for raising this important point. We acknowledge that our approach inherently regularizes trajectories towards the offline data manifold. However, our primary objective in this work is to ensure the generation of feasible, high-return trajectories for offline RL tasks particularly in settings where safety and reliability are crucial. This highlights a fundamental trade-off between trajectory feasibility and exploration freedom.

To clarify this explicitly, we will revise the manuscript to discuss this limitation in detail, clearly identifying scenarios where LoMAP provides the most substantial benefits (e.g., safety-critical domains or tasks prioritizing feasibility over aggressive exploration).

Furthermore, we believe LoMAP can be effectively combined with trajectory stitching and data augmentation techniques [2, 3], which enrich offline datasets with diverse, synthetic trajectories, which we leave for future work.

References:

[1] Um et al., Don't Play Favorites: Minority Guidance for Diffusion Models. In ICLR, 2024.

[2] Li et al., DiffStitch: Boosting Offline Reinforcement Learning with Diffusion-based Trajectory Stitching. In ICML, 2024.

[3] Chen et al., Extendable Long-Horizon Planning via Hierarchical Multiscale Diffusion. arXiv, 2025.

Thank you for the positive review and constructive feedback. Please find our detailed response below.

• “With the IVF method for faster nearest neighbors as described in Appendix D, what was the wall clock time to generate a single trajectory for a task like AntMaze? How expensive is the procedure with naive nearest neighbors?”

Thank you for highlighting this important practical aspect. We measured the wall-clock time required by Hierarchical Diffuser (HD) with LoMAP in the AntMaze task, averaging the generation time per trajectory using a single NVIDIA A10G GPU, with the number of neighbors set to $K=10$.

As shown above, the naive KNN search is prohibitively expensive, requiring over 21 seconds to generate a single trajectory, with the vast majority of time dedicated to the KNN search itself. In contrast, our IVF-based approximate KNN significantly reduces runtime, with only a negligible fraction of the total time spent on the nearest neighbor search.

• “Could you include QGPO to the results tables? It is a technique which trains a time dependent energy model for unbiased guidance, so perhaps can serve as a sort of upper bound in performance.”

Thank you for suggesting QGPO as an additional baseline. We agree that including QGPO can indeed serve as a valuable upper-bound comparison for performance. We appreciate this insightful recommendation and will incorporate the QGPO results into the final version of our paper.

• “Cosine distance metric function may not be good heuristics for other problems (especially with high dimensional state spaces).”

As the reviewer correctly points out, cosine distance may not be suitable for certain high-dimensional state spaces, especially pixel-based environments. Exploring effective manifold approximation methods for more complex, high-dimensional state spaces remains an important direction for future work. We believe one promising approach would be to combine LoMAP with latent trajectory representations [1, 2]. We will clarify this limitation explicitly in the final version.

• "An additional references to asymptotically unbiased guidance strategies (perhaps in the appendix)"

Thank you for this valuable suggestion. We will discuss asymptotically unbiased guidance strategies further in the related work section of the final version.

References

[1] Co-Reyes et al., Self-consistent trajectory autoencoder: Hierarchical reinforcement learning with trajectory embeddings. In ICML, 2018.

[2] Jiang et al., Efficient Planning in a Compact Latent Action Space. In ICLR, 2023.

Thank you for the detailed review and constructive feedback on our work. Please find our detailed responses below.

• Not a strong improvement over chosen baselines

We appreciate the valuable feedback provided by the reviewer. However, we respectfully disagree with the subjective judgment that there was “not strong” improvement. We believe the performance improvements demonstrated by our method are meaningful. Specifically, our approach consistently outperforms all baselines in terms of average performance across three established benchmarks, despite the simplicity of our method.

• Missing baselines

We sincerely thank the reviewer for pointing out several interesting works. However, we would like to clarify why the suggested baselines are not directly suitable for comparison in our setting:

• [1] assumes human-provided guidance in the form of sketches or demonstrations, which fundamentally differs from our scenario where no such human intervention is assumed. Additionally, [1] relies on a diffusion-based policy rather than a diffusion-based planner, making direct comparisons inappropriate. The reviewer's comment, "It also enables minority guidance without introducing constraint violations, as demonstrated in your Figure 5," is also not entirely accurate. The trajectories in [1] were generated using a diffusion policy trained with random walks rather than being goal-conditioned as in our experiments.
• [2] assumes the availability of differentiable constraints. For instance, applying [2] directly to the maze environment would require prior knowledge about the exact location of walls, making the comparison unfair.
• [3] focuses specifically on manipulation scenarios guided by language instructions, which significantly deviates from the tasks and settings considered in our work, making it an unsuitable baseline.

---

• Figure suggestion

Thank you for the suggestions! We will incorporate them into the final version of our paper.

• What is the intuition behind the LoMAP

Please see response for reviewer XEyQ for the details.

• Impact of the number of PCA components

We have conducted additional experiments by varying the number of PCA components ($K$). As shown in the results below, LoMAP demonstrates robustness provided that K is above a certain threshold. We select $K=10$ as it offers computational efficiency comparable to $K=20$, yet achieves similar performance.

• “Have you tried applying LoMAP only in early steps or only in later steps?”

Yes, we explored applying LoMAP exclusively in either early or later stages of the diffusion sampling process. We observed that applying LoMAP only in early steps was not particularly effective. In contrast, applying LoMAP during intermediate to later diffusion steps yielded the most significant improvements. This indicates that the mismatch between the model’s reverse transition and the true data distribution becomes most pronounced during these intermediate steps, aligning with similar observations reported in prior study [4].

• “Computational overhead of finding nearest neighbors”

Please see response for reviewer wazP for the details.

• “Does your algorithm require access not only to the diffusion model but also to the dataset the model was trained on?”

LoMAP, as currently presented, assumes the existence of an offline dataset, consistent with other offline RL baselines. For instance, Decision Diffuser (DD) retrains a conditional diffusion planner from scratch using offline datasets, and RGG trains an out-of-distribution (OOD) score predictor utilizing offline datasets. Unlike these methods, our approach leverages the dataset directly without additional training, making it simple and straightforward to apply.

References:

[1] Shi et al., Inference-Time Policy Steering through Human Interactions. In ICRA, 2025.

[2] Lee et al., Learning Diverse Robot Striking Motions with Diffusion Models and Kinematically Constrained Gradient Guidance. arXiv, 2024.

[3] Hao et al., DISCO: Language-Guided Manipulation with Diffusion Policies and Constrained Inpainting. arXiv, 2024.

[4] Na et al., Diffusion Rejection Sampling. ICML, 2024.

Thank you for the detailed review and constructive feedback on this work. We especially appreciate the insightful questions regarding the accurate manifold approximation. Please find our detailed answers below.

• "I do not fully understand why LoMAP can help with the issue of diffusion sampling."

We thank the reviewer for this important question. To clarify, the core intuition behind LoMAP is grounded in the observation that feasible trajectories from offline datasets lie on an intrinsically low-dimensional manifold embedded within a high-dimensional trajectory space. Diffusion-based sampling methods can deviate from this manifold due to inaccuracies in guidance (as analyzed in Section 3.1), leading to infeasible trajectories.

LoMAP addresses this problem by iteratively projecting guided diffusion samples back onto an approximated local manifold derived directly from offline trajectories. Specifically, at each diffusion step, we first compute a denoised estimate of the current trajectory sample using Tweedie's formula. We then retrieve the $k$-nearest neighbors from the offline dataset, which naturally represent trajectories that closely adhere to the true data manifold. By forward-diffusing these neighbor trajectories to the current diffusion timestep, we approximate the local structure of the intermediate manifold ($\mathcal{M}_{i-1}$).

PCA then provides a convenient and effective way to approximate the local manifold around this neighborhood. Because PCA identifies principal directions of variance, it naturally captures the major local geometric structures of the feasible manifold represented by nearby trajectories. By projecting the diffusion sample onto this PCA-derived subspace, we remove the off-manifold components, effectively correcting the artifact trajectories.

• "I'm wondering if the author can provide some evidence this method actually helps with more accurate manifold approximation."

Thanks for the important question. To provide further evidence, we computed the Realism Score [1], which quantifies how closely generated trajectories lie to the true manifold defined by the offline dataset. Specifically, we approximated the true manifold using k-NN hyperspheres constructed from 20,000 offline trajectories and generated 100,000 trajectory samples to measure the average realism score. As shown below, applying LoMAP consistently yields higher realism scores compared to diffusion sampling without LoMAP, clearly demonstrating that our method effectively produce samples closer to the true data manifold.

• “I suggest adding more DT-related baseline methods.”

Thank you for suggesting additional DT-related baselines. Following your advice, we now include comparisons with QDT [2], LPT [3] (guidance-based methods), and WT [4] (waypoint-based method). For MuJoCo locomotion tasks, we used the QDT implementation from the d3rlpy library and the official implementation provided by the authors for LPT, while WT results are taken directly from the original paper [4]. As shown in the table below, despite its simplicity, LoMAP consistently achieves the best average performance across these benchmarks. We will incorporate these additional comparisons into the final version of the paper.

References

[1] Kynkäänniemi et al., Improved precision and recall metric for assessing generative models. In NeurIPS, 2019.

[2] Yamagata et al., Q-learning Decision Transformer: Leveraging Dynamic Programming for Conditional Sequence Modelling in Offline RL. In ICML, 2023.

[3] Kong et al., Latent Plan Transformer for Trajectory Abstraction: Planning as Latent Space Inference. In NeurIPS, 2024.

[4] Badrinath et al., Waypoint Transformer: Reinforcement Learning via Supervised Learning with Intermediate Targets. In NeurIPS, 2023.