Shashank Hegde,Gautam Salhotra,Gaurav S. Sukhatme
Abstract: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.
What problem does this paper attempt to address?
The main problems that this paper attempts to solve are the limitations of existing methods based on Trajectory Diffusion in robot control tasks, specifically including:
1. **Model size and inference speed**: Current general - purpose policy models often require a large model scale in order to improve generalization ability, which leads to a large model volume and slow inference speed.
2. **Trade - off between performance and action time range**: For Diffusion Policy, fewer diffusion queries will accumulate larger trajectory tracking errors. Therefore, high - frequency inference is usually required to ensure performance, but this is limited by the computing resources of the robot.
3. **Performance degradation in long - action time ranges**: When the action time range is long, the performance of methods based on Trajectory Diffusion will significantly degrade due to the influence of modeling errors and trajectory tracking errors.
To solve these problems, the paper proposes a new method - Latent Weight Diffusion (LWD). LWD improves existing methods in the following ways:
- **Learning policy distributions instead of trajectory distributions**: LWD encodes demonstration trajectories into the latent space and uses a hypernetwork to decode these latent representations into policy parameters. This can generate smaller policy networks and reduce the number of diffusion model queries.
- **Closed - loop control**: The policies generated by LWD are closed - loop controllers, which can better respond to environmental changes, thus supporting longer action time ranges.
- **Task - specific policy generation**: LWD can generate task - specific policies according to task identifiers, which not only maintains multi - task generalization ability but also improves task - specific performance.
Through these improvements, LWD performs well in multiple benchmark tests. For example, in the Metaworld MT10 benchmark test, LWD achieves a higher success rate while using a model volume approximately 18 times smaller than that of traditional methods. In addition, LWD also shows better performance in tasks with long - action time ranges and reduces the call frequency of expensive diffusion models.
### Formula summary
- **Policy parameter distribution**:
\[
p(\theta)=\int\int p(\tau,\theta,z)\,dz\,d\theta
\]
\[
p(\tau)=\log\int\int p(\tau|z,\theta)p(\theta|z)p(z)\,dz\,d\theta
\]
- **Modified ELBO**:
\[
mELBO = \sum_{t = 1}^T\left[\mathbb{E}_{q(z|\tau)}\left[\mathbb{E}_{p(\theta|z)}[\log p(a_t|s_t,\theta)]\right]-KL(q(z|\tau)\|p(z))\right]
\]
- **VAE loss function**:
\[
L(\{s_t^k,a_t^k\}_{t = 1}^T|\phi_{enc},\phi_{dec})=C-\frac{1}{2\sigma^2}\sum_{t = 1}^T\mathbb{E}_{q_{\phi_{enc}}(z|\{s_t^k,a_t^k\}_{t = 1}^T)}\left[(a_t^k-\pi(s_t^k,f_{\phi_{dec}}(z)))^2\right]-\beta_{kl}\sum_{i = 1}^{dim(z)}\left(\sigma_i^2+\mu_i^2 - 1-\log\sigma_i^2\right)
\]
These formulas show how LWD learns policy distributions through diffusion models in the latent space and finally generates efficient robot control policies.