Crossway Diffusion: Improving Diffusion-based Visuomotor Policy via Self-supervised Learning (2307.01849v3)

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.

• The paper presents Crossway Diffusion, a novel approach that jointly optimizes diffusion loss with a self-supervised state reconstruction objective.
• It integrates a state decoder that reconstructs raw inputs from intermediate representations to improve visuomotor policy learning.
• Experiments on simulated and real robot tasks show up to a 15.7% success rate improvement and robust performance under distractions.

The paper presents "Crossway Diffusion," a method designed to enhance diffusion-based visuomotor policy learning. It builds upon existing diffusion policy approaches by incorporating a state decoder and an auxiliary self-supervised learning (SSL) objective. The key idea is to reconstruct input states (raw image pixels and other state information) from the intermediate representations generated during the reverse diffusion process. This reconstruction is enforced through a specifically designed state decoder. The whole model is optimized jointly using the SSL objective and the standard diffusion loss.

Here's a breakdown of the key components and contributions:

• Problem: Standard diffusion-based policies for robot imitation learning can be improved, particularly in terms of visual representation learning.
• Method: Crossway Diffusion introduces:
  • State Decoder: A neural network that reconstructs raw image pixels and other state information from intermediate representations within the reverse diffusion process. This forces the model to learn better intermediate representations.
  • Self-Supervised Learning (SSL) Objective: A loss function that measures the difference between the reconstructed states and the original input states, encouraging accurate reconstruction. This loss is combined with the original diffusion loss to train the entire model.
  • Intersection Transformation: Transformation applied to the intermediate diffusion representation before feeding it to the state decoder.
• Architecture: The Crossway Diffusion model consists of a state encoder, action encoder, action decoder (same as Diffusion Policy), and the new state decoder. The intermediate representation is dubbed as "intersection," because both flows of information pass through this representation.
• Experiments:
  • Evaluated on simulated robot tasks from Robomimic and a Push-T task from IBC.
  • Demonstrated on real-world robot manipulation tasks.
  • Compared against Diffusion Policy and Implicit Behavior Cloning (IBC).
• Results:
  • Crossway Diffusion consistently outperforms the baseline Diffusion Policy and IBC.
  • Shows significant improvement (e.g., 15.7% in success rate on "Transport, mh") on tasks with varied demonstration proficiency.
  • Qualitative results show good image reconstruction, suggesting effective representation learning.
  • The method exhibits robustness to distractions such as unseen objects and partial occlusions in real-world settings.
  • Ablation studies validate the design choices, including the state decoder architecture and the auxiliary reconstruction objective.
• Ablation Studies:
  • Different designs of the state decoder
  • Predicting the future state versus reconstructing the current state with the state decoder
  • Using contrastive learning instead of image reconstruction as an SSL objective.
• Key Findings from Ablations:
  • Forcing the two flows of information (denoising and state reconstruction) to intersect is important. Design D, where these two flows were disentangled, showed worse performance.
  • Reconstructing the current state is more beneficial than predicting future states as an auxiliary objective.
  • Not all SSL objectives are beneficial. Contrastive learning as an auxiliary loss performed worse than the image reconstruction auxiliary task.
• Contributions:
  • A novel method (Crossway Diffusion) for improving diffusion-based visuomotor policies.
  • Extensive experiments on simulated and real-world tasks.
  • Detailed ablation studies to justify the design choices.

In essence, Crossway Diffusion leverages self-supervised learning through state reconstruction to improve the visual representations learned by diffusion-based policies, leading to better performance in robot imitation learning tasks.