Latent Weight Diffusion: Generating policies from trajectories
Diffusion methods;long horizon robotics tasks;Imitation Learning
Yes
The paper is 3 lines over the 10-page limit.
Yes, Other reasons (please specify below)
The paper writing can also be improved at a few places: (1) \tau is not defined when it is first introduced in Sec. 3.1, as well as \varepsilon in Sec. 3.3. (2) In Fig. 8, how is Relative Performance Change defined? Is it defined relative to the policy itself, or on an absolute scale? I find it quite misleading. Also, what is the action parameterization that LWD generates? I assume it is not diffusion?

In Sec. 3.1, it says a_t= \pi(s_t, \theta) + e, where e is normally distributed around 0. Why is the added action noise needed?

Could you comment on the other dimensions of the PCA analysis in Figure 5? Are they not very informative? It might be a good idea to show them in the appendix.

The paper is three lines over the 10-page limit. I suggest moving some of the derivations in Sec. 3 to the appendix, or try to shrink the space at places.
The question of distilling a large multi-task policy into smaller ones for fine-grained control is very relevant in current robotics research. The research proposed by the authors is well-motivated, and the overall setup of applying latent diffusion is intuitive and can potentially lead to strong individual policies while maintaining the diversity.

Based on the proposed approach, some of the experimental results are encouraging. Table 2 shows the benefit of LWD compared to undistilled MLP policies in MetaWorld and demonstrates moderate performance improvement. Figure 8 shows the resilience of LWD to longer action chunk prediction in the PushT task.
The authors propose Latent Weight Diffusion (LWD) that learns a distribution of policy parameters and aims to capture the diversity of trajectory dataset and achieve fine-grained closed-loop control with small model size. The method is mostly based on Hedge et al., 2023, with the experiments using existing trajectory datasets instead. A variational autoencoder (VAE) is used to compress the trajectories into a compact latent space, where a diffusion model is then applied to learn the complex distribution of possible latents. The decoder decodes the network parameters for the policy. LWD is evaluated in simulated environments, and the results demonstrate that LWD can reconstruct the behavior of the original policies and achieve improved performance compared to larger models or baseline with long action chunk prediction in certain settings.
On the technical side, the paper’s novelty is fairly limited since it is a direct application of Hedge et al., 2023 using existing trajectory datasets. While the idea is promising, the experiment results are lacking and suggest important limitations. First, it is unclear how the hypernetwork setup scales up to more complex tasks. While the argument is that LWD can generate smaller policies for individual tasks as Table 2 shows, the authors have not shown convincing results on decomposing a single long-horizon task into subtasks where LWD generates a policy for each part. The experiment on action chunk size (Figure 8) alludes to such possibility but I think dedicated experiments on this setup can strengthen the story significantly and make the method much more relevant for current robotics research. Second, it is unclear how well the hypernetwork setup scales to larger model size when larger size is required, e.g., in more complex tasks. Can the hypernetwork still handle \theta of very high dimensions? e.g., on the order of millions, or can it handle pixel input?

The results with reconstructing the original policy are quite mixed and I have trouble understanding the conclusions made by the authors. It seems that the learned latent cannot differentiate the trajectories well with shorter length (Sec. 4.2). The discussion is fairly lacking. Sec. 4.2 also has a few comments that are not backed with experimental results, e.g., “Surprisingly, we noticed that after training our VAE on the snippets, the decoded policies from randomly snipped trajectories were still faithfully behaving like their original policies.”, and “we noticed that the decoded policies from the trajectory snippets did not perform as well as the original policies”. Please clarify or point to existing results if I misunderstand the comments.

I am fairly curious about the hypernetwork decoder setup, but the paper does not provide details on the architecture and how it handles very large output space (the decoded policy parameters). There is also no appendix for such experimental details.
Yes
No ethics review needed.
1. Please address the above points in Weaknesses.
2. In all experiments, how many diffusion steps are typically needed to decode the policy network?
1. It is indeed interesting to explore the idea of latent diffusion in imitation learning following the success in [1]. The diffused latent is then decoded into policy hyperparameters via hypernetwork [3], which is significant since the diffusion model can capture diverse and multi-modality network parameter space, therefore enabling both state and multi-task conditioning. 
   - According to the experiment results, although not extensively compared with Diffusion policy in a state-conditioned setting, LWD is shown to work better than simple MLP in a multi-task setting as a proof-of-concept.

2. Sec 4.2 is a very insightful experiment. It sheds light on task difficulties and similarities. This setup can be used further to analyze/visualize task characteristics for imitation learning. 

[1] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[2] Chi, Cheng, et al. "Diffusion policy: Visuomotor policy learning via action diffusion." The International Journal of Robotics Research (2023): 02783649241273668.

[3] Ha, David, Andrew Dai, and Quoc V. Le. "Hypernetworks." arXiv preprint arXiv:1609.09106 (2016).
The proposed method, Latent Weight Diffusion (LWD), uses a two-stage process involving a variational autoencoder (VAE) coupled with a latent diffusion model. First, LWD encodes demonstration trajectories into a latent space using a VAE, then learns the distribution over these latent representations using diffusion in this latent space. A hypernetwork decodes the latent representations into policy parameters, generating closed-loop policies that are state or task conditioned. This approach leverages task-specific conditioning to maintain task generalization and efficiency, producing policy networks with reduced model sizes. The authors evaluate LWD on the Metaworld MT10 and D4RL benchmarks. They compare its performance on diverse tasks, such as multi-task learning and long-horizon robotic control, against baselines, including multi-task MLPs.
1. Related works on Sec. 2.2 should additionally cite [2, 3] for the statement "...to learn smooth cost functions for the joint optimization of grasp and motion plans."

2. There are some issues with the experiment outline for benchmarking LWD. For example:
    - Diffusion policy [3] should be compared with LWD in Table 2 to strongly confirm the architecture choice efficiency claimed in the paper. For example, the input of Diffusion policy [3] should be robot state concatenating with the task indicator, and Diffusion policy [3] should be trained with data from all tasks.
    - Ablation on network-size of LWD should be considered vs. success rate with tasks in [1]. For example, another column of LWD with 370.4k parameters should be added to Table 2 to study the effect of a bigger policy network for LWD.
    - Sec 4.1. purpose is to see generalization in behaviors. However, comparing foot contact distributions between models for ablation study is not meaningful. Adding a comparison of state variances/counting different success modes of pick & place or pushing tasks in [1] between models is more insightful.
    - On a side note, Fig. 3 would be much clearer if merged into one figure (probably with a smaller scatter size), making it easier to see the foot contact distribution overlaps.

[1] Yu, Tianhe, et al. "Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning." Conference on robot learning. PMLR, 2020.

[2] Carvalho, Joao, et al. "Motion planning diffusion: Learning and planning of robot motions with diffusion models." 2023 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS). IEEE, 2023.

[3] Luo, Yunhao, et al. "Potential based diffusion motion planning." ICML (2024).

[4] Chi, Cheng, et al. "Diffusion policy: Visuomotor policy learning via action diffusion." The International Journal of Robotics Research (2023): 02783649241273668.
Yes
No ethics review needed.
See the weakness.
1. This paper provides a novel method for generating the network parameters with diffusion models.
2. The experiments showing the distribution modeling of the generated network parameter is compressive.
The paper "Latent Weight Diffusion (LWD): Generating Policies from Trajectories" introduces a novel approach for imitation learning in robotics. Current methods often rely on trajectory-based diffusion models, which can be computationally demanding and less effective in capturing long action sequences. To address these challenges, LWD diffuses policy parameters instead of trajectories. The model first encodes demonstration trajectories into a latent space using a Variational Autoencoder (VAE). Then, a diffusion model learns the distribution within this latent space, allowing the generation of compact, task-specific policies via a hypernetwork decoder.
1. Table 2 is unclear. It is not clear what is an MLP policy. Does it predict a sequence of future actions or just a single-step action? If it is the latter, it is unfair to compare these MLP policies to the proposed LWD method, since LWD predicts a sequence of future actions.

2. Figure 8 is also unclear. The authors didn't show the success rates of diffusion policy and LWD. They just show the "relative performance change" curve. I hope the author can show the success rate also.

3. Why is LWD close-loop and why does the author say diffusion policy is open-loop? I guess the authors want to say that LWD can generate a new policy every $H_a$ step. However, diffusion policy can also regenerate new action sequences after every $K$ step where $K$ is a hyperparameter. I think here the "close-loop" is an inaccurate and unprofessional word.

4. From my understanding, the author didn't compare the success rates between LWD and diffusion policy on any task. So the claimed contribution in the Abstract and Introduction (LWD offers the benefits of considerably smaller policy networks during inference and requires fewer diffusion model queries) is not supported by evidence.
Yes
No ethics review needed.
1. The authors use the long-horizon experiment to demonstrate closed-loop stability, but this could also be due to the diffusion model's limited capacity to capture the high-dimensional distribution fully. Could the authors compare the performance of the two models under disturbances to show how the learned closed-loop policy improves stability? Can it improve stability with disturbances?
    
2. When Diffusion Policy is running online, it can use the previous solution to warm-start the next step's action sequence. Since in the neighbor state, the policy shouldn’t change drastically. In LWD, do the authors think the same design could be possible?
The combination of diffusion models with hypernetworks is quite novel and enables fast online inference. LWD is able to generate diverse behaviors while still maintaining low inference latency.
This paper improves the online inference time of diffusion-generated policies by using diffusion to generate a hypernetwork instead of only outputting an action sequence. The policy is then reconstructed and inferred at higher frequency.
It is unclear whether policy stability is improved without further interaction with the environment. Specifically, without counterfactual reasoning or data on failed actions, how can the generated policy recover from states that deviate from the nominal path? If the author could empirically verify this or provide more rationale, it would make the approach more sound.
Latent diffusion method that can learn a distribution of policies from a demonstration dataset.
@inproceedings{
anonymous2024latent,
title={Latent Weight Diffusion: Generating policies from trajectories},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=XLCqhdaMpy},
note={under review}
}
With the increasing availability of open-source robotic data, imitation learning has emerged as a viable approach for both robot manipulation and locomotion. Currently, large generalized policies are trained to predict controls or trajectories using diffusion models, which have the desirable property of learning multimodal action distributions. However, generalizability comes with a cost — namely, larger model size and slower inference. Further, there is a known trade-off between performance and action horizon for Diffusion Policy (i.e., diffusing trajectories):
fewer diffusion queries accumulate greater trajectory tracking errors. Thus, it is common practice to run these models at high inference frequency, subject to robot computational constraints.

To address these limitations, we propose Latent Weight Diffusion (LWD), a method that uses diffusion to learn a distribution over policies for robotic tasks, rather than over trajectories. Our approach encodes demonstration trajectories into a latent space and then decodes them into policies using a hypernetwork. We employ a diffusion denoising model within this latent space to learn its distribution. We demonstrate that LWD can reconstruct the behaviors of the original policies that generated the trajectory dataset. LWD offers the benefits of considerably smaller policy networks during inference and requires fewer diffusion model queries. When tested on the Metaworld MT10 benchmark, LWD achieves a higher success rate compared to a vanilla multi-task policy, while using models up to ∼18x smaller during inference. Additionally, since LWD generates closed-loop policies, we show that it outperforms Diffusion Policy in long action horizon settings, with reduced diffusion queries during rollout.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Diffusion methods long horizon robotics tasks Imitation Learning
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/e338d20a3dbd21ade88b4fd69f37617fe2a2e286.pdf
applications to robotics, autonomy, planning
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Latent Weight Diffusion: Generating policies from trajectories
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission