Multi-Task Diffusion Policy for Robotic Manipulation
A single conditional diffusion model trained to solve Push-T, Lift, and Can tasks
I explored and extended a state-based diffusion policy to a multi-task setting. The goal was to train a single shared model to perform three distinct robotic manipulation tasks—Push-T, Lift, and PickPlaceCan—by combining per-task processing with a unified noise-prediction network.
What I built
- Extended a diffusion policy to solve the PyMunk Push-T task and the Robosuite Lift and Can tasks using a single shared 1D U-Net.
- Designed per-task linear projection heads to map native observation features (ranging from 5 to 23 dimensions) into a shared 64-dimensional space.
- Implemented task conditioning by concatenating a one-hot task label to the projected observations.
- Built per-task action encoders and decoders to manage varying action spaces across different robots and environments.
- Created a custom collate function and uniform/weighted sampling strategies to combine and batch training data across the multiple tasks.
- Utilized action chunking during live environment rollouts to execute a horizon of multiple predicted actions without replanning.
Why it matters
Training a single architecture to handle multiple disparate tasks is a major step toward generalist robot agents. By using task-conditioned diffusion models with shared representations, the network learns robust, unified features while seamlessly bridging task-specific observation and action spaces.