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

TL;DR

Crossway Diffusion introduces a self-supervised learning approach with a state decoder to enhance diffusion-based visuomotor policies, leading to improved robot imitation performance in both simulated and real-world tasks.

Contribution

It presents a novel method combining a state decoder and SSL objectives to improve diffusion-based robot policies, which was not explored before.

Findings

Significant performance improvements over standard diffusion policies.

Effective in both simulated and real-world robot tasks.

Outperforms baseline methods consistently.

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.