Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Crossway Diffusion: Improving Diffusion-based Visuomotor Policy via Self-supervised Learning

Xiang Li, Varun Belagali, Jinghuan Shang, Michael S. Ryoo

TL;DR

Crossway Diffusion enhances diffusion-based visuomotor imitation by adding a state decoder and a self-supervised reconstruction objective that regularizes intermediate representations. The method introduces an intersection transformation to jointly leverage the diffusion denoising path and state reconstruction, yielding consistent performance gains over Diffusion Policy across simulated and real-world tasks, including a notable improvement on Transport, mh. Ablation studies validate the importance of the reconstruction task and the specific state-decoder designs, while also discussing inference-speed considerations and qualitative robustness. These results suggest that SSL-informed representation learning can substantially improve diffusion-based policies in visually rich robotic control settings, with practical impact on real-world manipulation tasks.

Abstract

Sequence modeling approaches have shown promising results in robot imitation learning. Recently, diffusion models have been adopted for behavioral cloning in a sequence modeling fashion, benefiting from their exceptional capabilities in modeling complex data distributions. The standard diffusion-based policy iteratively generates action sequences from random noise conditioned on the input states. Nonetheless, the model for diffusion policy can be further improved in terms of visual representations. In this work, we propose Crossway Diffusion, a simple yet effective method to enhance diffusion-based visuomotor policy learning via a carefully designed state decoder and an auxiliary self-supervised learning (SSL) objective. The state decoder reconstructs raw image pixels and other state information from the intermediate representations of the reverse diffusion process. The whole model is jointly optimized by the SSL objective and the original diffusion loss. Our experiments demonstrate the effectiveness of Crossway Diffusion in various simulated and real-world robot tasks, confirming its consistent advantages over the standard diffusion-based policy and substantial improvements over the baselines.

Crossway Diffusion: Improving Diffusion-based Visuomotor Policy via Self-supervised Learning

TL;DR

Crossway Diffusion enhances diffusion-based visuomotor imitation by adding a state decoder and a self-supervised reconstruction objective that regularizes intermediate representations. The method introduces an intersection transformation to jointly leverage the diffusion denoising path and state reconstruction, yielding consistent performance gains over Diffusion Policy across simulated and real-world tasks, including a notable improvement on Transport, mh. Ablation studies validate the importance of the reconstruction task and the specific state-decoder designs, while also discussing inference-speed considerations and qualitative robustness. These results suggest that SSL-informed representation learning can substantially improve diffusion-based policies in visually rich robotic control settings, with practical impact on real-world manipulation tasks.

Abstract

Sequence modeling approaches have shown promising results in robot imitation learning. Recently, diffusion models have been adopted for behavioral cloning in a sequence modeling fashion, benefiting from their exceptional capabilities in modeling complex data distributions. The standard diffusion-based policy iteratively generates action sequences from random noise conditioned on the input states. Nonetheless, the model for diffusion policy can be further improved in terms of visual representations. In this work, we propose Crossway Diffusion, a simple yet effective method to enhance diffusion-based visuomotor policy learning via a carefully designed state decoder and an auxiliary self-supervised learning (SSL) objective. The state decoder reconstructs raw image pixels and other state information from the intermediate representations of the reverse diffusion process. The whole model is jointly optimized by the SSL objective and the original diffusion loss. Our experiments demonstrate the effectiveness of Crossway Diffusion in various simulated and real-world robot tasks, confirming its consistent advantages over the standard diffusion-based policy and substantial improvements over the baselines.
Paper Structure (30 sections, 6 equations, 13 figures, 10 tables)

This paper contains 30 sections, 6 equations, 13 figures, 10 tables.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Behavioral Cloning
  4. Diffusion Models
  5. Diffusion Models for Policy Learning
  6. Method
  7. State Decoder
  8. Intersection Transformation
  9. Crossway Diffusion Loss
  10. Experiment
  11. Task and Dataset
  12. Main Results
  13. Evaluation Metrics
  14. Simulated Experiments
  15. Real-world Experiments
  16. ...and 15 more sections

Figures (13)

  • Figure 1: Trajectory generation formulation of Diffusion Policy. This figure shows a case where $T_s=2$ and $T_a=4$.
  • Figure 2: Left: Architecture of Crossway Diffusion. We introduce a state decoder to the existing Diffusion Policy chi2023diffusion as well as an auxiliary reconstruction objective $\mathcal{L}_{\textit{Recon.}}$. The state decoder takes a transformed intermediate representation 'intersection' to reconstruct the input states. Right: Transformation applied to 'intersection' for the visual state decoder.
  • Figure 3: Architecture of visual state decoder. Numbers in the blocks indicate the number of output channels except that ConvTranspose doubles the spatial resolution while keeping the number of channels unchanged.
  • Figure 4: Visual reference for all tasks
  • Figure 5: Duck Lift task under different obstructions. Two rows show the images captured by two cameras respectively. The red arrow in (e) is used to indicate the position of the duck.
  • ...and 8 more figures