💻 computer science

DiSPo: Diffusion-SSM based Policy Learning for Coarse-to-Fine Action Discretization

The paper proposes DiSPo, a novel diffusion-state space model (SSM) based policy that leverages Mamba to learn from diverse coarse demonstrations and generate varying control scales of actions, achieving superior success rates and inference efficiency in coarse-to-fine skill learning compared to traditional baselines.

Original authors: Nayoung Oh, Jaehyeong Jang, Moonkyeong Jung, Daehyung Park

Published 2026-02-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Nayoung Oh, Jaehyeong Jang, Moonkyeong Jung, Daehyung Park

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to perform a delicate task, like threading a needle or passing a clamp through a tiny pipe.

The Problem: The "Blurry" vs. "Sharp" Video Dilemma
Traditionally, to teach a robot these skills, you need to show it a high-definition, slow-motion video of a human doing it perfectly. If you only show the robot a low-resolution, fast-forwarded video (coarse data), the robot gets confused. It sees the big picture but misses the tiny, critical details needed to avoid crashing.

Most current AI methods are like a camera that only has one zoom level. If you feed it a blurry video, it tries to guess the details, but it often ends up jittery or crashing. If you feed it a sharp video, it works well, but recording that much high-quality data is expensive, slow, and hard to store.

The Solution: DiSPo (The "Smart Zoom" Robot)
The paper introduces DiSPo (Diffusion-SSM based Policy). Think of DiSPo as a robot with a magical, adaptive zoom lens and a super-memory.

Here is how it works, broken down into simple metaphors:

1. The "Smart Zoom" (Coarse-to-Fine)

Imagine you are looking at a map.

Old Robots: If you give them a map of the whole country, they get lost when they need to turn a corner in a specific street. If you give them a map of just the street, they don't know where the country is.
DiSPo: DiSPo can look at a zoomed-out map (coarse data) to understand the general route, and then instantly "zoom in" to see the tiny cracks in the pavement (fine details) exactly when it needs to make a precise move. It doesn't need a high-definition video for the whole journey; it only needs high-definition data for the tricky parts.

2. The "Two Superpowers" (Diffusion + SSM)

DiSPo combines two different types of AI brains:

The Diffusion Artist (The "Denoiser"): Imagine a sculptor starting with a block of noisy, messy clay. They slowly chip away the noise to reveal a perfect statue. In AI terms, DiSPo starts with a random guess of what to do and slowly "cleans it up" step-by-step until it finds the perfect movement. This makes it very good at learning complex, artistic movements.
The SSM Librarian (The "Memory Keeper"): This is the "Mamba" part. Imagine a librarian who can read a book and remember the plot perfectly without needing to re-read every page. Traditional AI needs to re-read the whole history to remember what happened. The SSM is efficient; it remembers the "flow" of the action very well, even if the action happens at different speeds.

3. The "Magic Dial" (Step-Scale Factor)

This is the paper's biggest innovation.

The Analogy: Imagine you are driving a car. Sometimes you are on a long, straight highway (you can drive fast and take big steps). Sometimes you are in a crowded parking lot (you need to move slowly and take tiny, careful steps).
How DiSPo uses it: DiSPo has a "speed dial."
- When the robot is in open space, it turns the dial to "Coarse" (big steps, fast processing).
- When it approaches a wall or a button, it turns the dial to "Fine" (tiny steps, high precision).
- Crucially: It can learn this from a video that was recorded at a slow, low-quality speed, and then decide on its own to move quickly in safe areas and slowly in dangerous areas.

4. The "Fake Practice" Trick (Pseudo-Demonstrations)

How does the robot learn to move finely if it was only shown slow, coarse videos?

The Analogy: Imagine you are learning to play a piano piece. You only have a slow, muffled recording of a master playing it.
The Trick: DiSPo listens to the slow recording, guesses what the fast notes should sound like, and creates a "practice version" for itself. It then checks its guess against the slow recording to make sure it didn't go off-track. It repeats this until it can play the song at full speed, even though it only had the slow recording to start with. This saves a massive amount of time and storage space.

The Real-World Result

The researchers tested this on a real robot arm (a UR5e).

The Task: Passing a clamp through a tiny pipe without touching the sides, and pressing a button without hitting the wall.
The Outcome:
- Old Robots: They either crashed into the pipe or stopped awkwardly because they couldn't figure out the tiny details from the coarse data.
- DiSPo: It successfully passed the clamp and pressed the button with high precision. It moved fast when it was safe, and slowed down to a microscopic precision when it was near the obstacle.

Summary

DiSPo is like a robot that learns from a "sketch" but draws a "photograph."
It takes rough, low-quality demonstrations, understands the big picture, and uses its special "zoom" ability to fill in the missing high-speed, high-precision details on the fly. This means we don't need to spend hours recording perfect, slow-motion videos to teach robots; we can teach them with rougher, faster data, and they will figure out the fine details themselves.

1. Problem Statement

The paper addresses a critical gap in Imitation Learning (IL) for robotics: the inability of current methods to effectively learn and reproduce skills at multiple levels of granularity (multi-granularity learning) from diverse demonstrations.

The Challenge: Real-world manipulation tasks (e.g., screwing, welding, insertion) require a mix of coarse movements (large positioning) and fine movements (precise adjustments).
Limitations of Existing Methods:
- Traditional IL (e.g., DMPs): Rely on fixed dynamics models or specific input frequencies, struggling to adapt to varying temporal scales.
- Neural IL (e.g., Transformers, Diffusion): While powerful, they typically learn from a single fixed frequency. To achieve fine-grained control, they require high-frequency demonstrations, leading to massive storage and computational overhead.
- State Space Models (SSMs): While efficient (e.g., Mamba), standard implementations use fixed step sizes, making them ill-suited for capturing complex, multi-modal behaviors across different granularities.

The goal is to create a policy that can learn from coarse (low-frequency) demonstrations and generate fine-grained (high-frequency) actions on demand, adapting to user-intended control scales without retraining.

2. Methodology: DiSPo Framework

The authors propose DiSPo, a novel policy combining Diffusion Models (for representation power) and State Space Models (SSMs), specifically Mamba, (for efficient, flexible discretization).

Core Architecture

Diffusion-SSM Hybrid: DiSPo uses a Denoising Diffusion Probabilistic Model (DDPM) framework where the denoising network is built upon stacked DiSPo blocks. Each block is a variant of the Mamba block.
Input Conditioning: The model takes:
- Noisy action sequences ( $a^{(k)}$ ).
- Observation sequences ( $o$ ).
- Diffusion step embedding ( $k$ ).
- Step-scale factors ( $r_t$ ): A user-defined or predicted vector controlling the temporal granularity of the output.

Key Innovations

Step-Scaled Mamba Block:
- Unlike standard Mamba which has fixed parameters, DiSPo introduces a step-scaling mechanism.
- The SSM parameters ( $\Delta, B, C$ ) are dynamically adjusted based on the input step-scale factor $r_t$ .
- Equation: $\Delta^{(i)}_t = r_t \cdot \text{SoftPlus}(f^{(i)}_\Delta(\ddagger u^{(i)}_t))$ .
- This allows the model to effectively "zoom in" or "zoom out" on the temporal resolution of the action sequence, enabling the generation of actions at arbitrary frequencies.
Multi-Granularity Learning Scheme:
- Pretraining (Sample-Rate Augmentation): The model is trained on datasets augmented with random step-scale factors. It learns to reconstruct actions at various frequencies from a single dataset.
- Fine-Tuning (Pseudo-Demonstration Generation): To generate high-frequency actions that were not present in the original coarse dataset, the authors propose a loop:
  1. Generate a coarse action sequence using the pretrained model.
  2. Replace the coarse actions with the original high-frequency demonstration data (where available) and add noise.
  3. Denoise to create a "pseudo-demonstration" at the target high frequency.
  4. Fine-tune the model on these pseudo-demonstrations to refine high-frequency details.
Granularity Predictor ( $\phi_r$ ):
- An MLP that predicts the optimal step-scale factor $r_t$ based on the current observation sequence, allowing the robot to automatically switch between coarse (efficient) and fine (precise) control modes online.

3. Key Contributions

First Integration of Mamba and Diffusion for Granularity: This is the first work to modulate the discrete model of Mamba for fine-grained robotic manipulation, enabling flexible action discretization.
Coarse-to-Fine Learning from Coarse Data: The method successfully learns to generate high-frequency, precise actions starting from low-frequency (coarse) demonstrations, eliminating the need for expensive high-frequency data collection.
Efficiency: By generating coarse motions in non-critical regions and fine motions only where necessary, DiSPo reduces inference overhead and computational cost.
New Benchmarks: The authors introduce coarse-to-fine imitation learning benchmarks (Clamp Passing, Passage Passing, Button Touch) to evaluate multi-granularity capabilities.

4. Experimental Results

The authors evaluated DiSPo on three simulated benchmarks and two real-world tasks (using a UR5e manipulator).

Simulation Performance:
- Success Rate: DiSPo achieved >81% success rate across all demonstration frequencies (2.5 Hz to 20 Hz), significantly outperforming baselines (DiffusionPolicy-C/T, D-Ma, VQ-BeT).
- Coarse-to-Fine Gap: When trained on 2.5 Hz data and tested at 20 Hz, baselines failed (near 0% success), while DiSPo maintained high performance.
- Mixed-Frequency Learning: In tasks with mixed 2.5 Hz and 5 Hz data, DiSPo achieved 93% success on the "Button Touch" task, whereas baselines struggled with inappropriate action speeds.
Real-World Experiments:
- Tasks: Clamp passing (inserting a ring into a pipe) and Button touching.
- Results: DiSPo successfully performed precise collision-free insertions and button touches.
- Comparison: The baseline (D-Ma) often caused collisions or stopped prematurely, while DiSPo generated the necessary fine-grained corrections.
- Efficiency: Using the data-driven step-scaling predictor, DiSPo reduced the number of required action steps by 32% in the button touch task with only a minor drop in success rate, proving its ability to optimize inference cost.

5. Significance and Impact

Data Efficiency: DiSPo solves the "data bottleneck" in robotics. It allows robots to learn complex, high-precision skills from easily collectable, low-frequency human demonstrations, removing the barrier of needing high-frequency teleoperation data.
Adaptive Control: The ability to modulate action granularity online means robots can be more energy-efficient and computationally efficient by using coarse control for open spaces and switching to fine control only for critical interactions.
Scalability: The integration of SSMs (Mamba) provides linear computational complexity, making the approach scalable to long-horizon tasks compared to Transformer-based diffusion policies.

In conclusion, DiSPo represents a significant step forward in making robotic imitation learning more flexible, data-efficient, and capable of handling the diverse temporal scales required for dexterous manipulation.