Sample from What You See: Visuomotor Policy Learning via Diffusion Bridge with Observation-Embedded Stochastic Differential Equation

TL;DR

BridgePolicy introduces a diffusion-bridge approach that embeds observations into the stochastic dynamics of visuomotor policies, enabling more precise control by starting sampling from an observation-informed prior.

Contribution

The paper presents a novel diffusion-bridge formulation for imitation learning that integrates heterogeneous observations directly into the stochastic process, improving robotic control performance.

Findings

Outperforms state-of-the-art generative policies across benchmarks

Improves control precision and reliability in real-world tasks

Effectively unifies visual and state inputs for diffusion models

Abstract

Imitation learning with diffusion models has advanced robotic control by capturing the multi-modal action distributions. However, existing methods typically treat observations only as high-level conditions to the denoising network, rather than integrating them into the stochastic dynamics of the diffusion process itself. As a result, the sampling is forced to begin from random noise, weakening the coupling between perception and control and often yielding suboptimal performance. We propose BridgePolicy, a generative visuomotor policy that directly integrates observations into the stochastic dynamics via a diffusion-bridge formulation. By constructing an observation-informed trajectory, BridgePolicy enables sampling to start from a rich and informative prior rather than random noise, substantially improving precision and reliability in control. A key difficulty is that diffusion bridge normally connects distributions of matched dimensionality, while robotic observations are heterogeneous and not naturally aligned with actions. To overcome this, we introduce a multi-modal fusion module and a semantic aligner to unify the visual and state inputs and align the observations with action representations, making diffusion bridge applicable to heterogeneous robot data. Extensive experiments across 52 simulation tasks on three benchmarks and 5 real-world tasks demonstrate that BridgePolicy consistently outperforms state-of-the-art generative policies.