DSSP: Diffusion State Space Policy with Full-History Encoding

TL;DR

DSSP introduces a history-conditioned diffusion policy using state space models to improve long-horizon robot manipulation tasks, achieving state-of-the-art results with efficient full-history encoding.

Contribution

The paper proposes DSSP, a novel diffusion policy with full-history encoding via state space models, enhancing long-term task performance in robot manipulation.

Findings

DSSP outperforms existing methods on simulation benchmarks.

DSSP achieves state-of-the-art results with smaller model size.

Hierarchical conditioning effectively captures long-term dependencies.

Abstract

Diffusion-based imitation learning has shown strong promise for robot manipulation. However, most existing policies condition only on the current observation or a short window of recent observations, limiting their ability to resolve history-dependent ambiguities in long-horizon tasks. To address this, we introduce DSSP, a history-conditioned Diffusion State Space Policy that enables efficient, full-history conditioning for robot manipulation. Leveraging the continuous sequence modeling properties of State Space Models (SSMs), our history encoder effectively compresses the entire observation stream into a compact context representation. To ensure this context preserves critical information regarding future state evolution, the encoder is optimized with a dynamics-aware auxiliary training objective. This high-level context representation is then seamlessly fused with recent state observations to form a hierarchical conditioning mechanism for action generation. Furthermore, to maintain architectural consistency and minimize GPU memory overhead, we also instantiate the diffusion backbone itself using an SSM. Extensive experiments across simulation benchmarks and real-world manipulation tasks show that DSSP achieves state-of-the-art performance with a significantly smaller model size, demonstrating superior efficiency of the hierarchical conditioning in capturing crucial information as the history length increases.