DiffuseManip

From-scratch implementation of Diffusion Policy (Chi et al., 2023) for robotic manipulation, built on robosuite + robomimic. Achieves 100% success rate on Lift (50/50 rollouts) with a low-dimensional observation policy trained from scratch on a single CPU.

Results

Task	Success Rate	Avg Episode Length	BC-RNN reference¹
Lift	100% (50/50)	54 steps (2.7s)	94%
Can	coming soon	—	65%
Square	78% (39/50)	235 steps	26%

¹ BC-RNN numbers from Chi et al. 2023 (Diffusion Policy paper), not reproduced here.

Evaluated with 50 rollouts per task, randomized initial state, 20 Hz control.

Square: 78% — above the BC-RNN reference (26%) and in the Diffusion Policy paper's regime. Getting here required fixing an eval-time observation bug, not more training. The v141 dataset stores the nut's position relative to the end-effector in the eef's local frame (R(q_eef)ᵀ · (nut_pos − eef_pos)), but our robosuite-1.5 obs reconstruction computed it in the world frame (eef_pos − nut_pos). Only 3 of 23 obs dims were affected, but they were fed out-of-distribution on every rollout, pushing normalized obs past −2.6 and capping success at ~40%. Replaying the stored MuJoCo states through robosuite 1.5.2 showed 21/23 dims matched v141 exactly; the corrected eef-frame formula matches all 23 to 0.000000 across 428 frames. The same checkpoint that scored 40% with the buggy obs scores 78% with the fix — no retraining, no dataset regeneration.

What is Diffusion Policy?

Standard behavioral cloning learns a deterministic mapping from observation to action. This fails when expert demonstrations are multimodal — e.g., the expert sometimes grasps from the left, sometimes from the right. A deterministic policy averages these modes and predicts an action straight through the object.

Diffusion Policy instead learns a distribution over actions by reversing a noise process (the same mechanism as DALL-E / Stable Diffusion). At inference, it starts from Gaussian noise and iteratively denoises into a coherent action sequence conditioned on the current observation. This lets it represent and sample from multiple valid strategies.

Key design choices in this implementation:

• 1-D temporal U-Net operating over the action horizon axis (not image spatial axes)
• FiLM conditioning — observation embedding modulates every residual block via learned scale + shift
• DDPM training (ε-prediction, 100 steps), DDIM inference (16 steps, deterministic)
• Receding-horizon control: predict Tp=16, execute Ta=8, replan

Architecture

Observation (To=2 frames, 19-dim)
        │
        ▼
  LowDimEncoder
  Linear → SiLU → Linear
        │
        ▼ obs_cond (256-dim)
        │
        ├────────────────────────────────┐
        │                                │
        ▼                                │
  Noisy actions (Tp=16, 7-dim)    Timestep embedding
        │                                │
        └──────────┬─────────────────────┘
                   │
                   ▼
         ConditionalUNet1D
         ┌─────────────────────┐
         │ Encoder             │
         │  ResBlock(7→256)    │
         │  Downsample ↓2      │
         │  ResBlock(256→512)  │
         │  Downsample ↓2      │
         │  Bottleneck(512→1024│
         │ Decoder             │
         │  Upsample ↑2 + skip │
         │  ResBlock(1024→512) │
         │  Upsample ↑2 + skip │
         │  ResBlock(512→256)  │
         │  Final(256→7)       │
         └─────────────────────┘
                   │
                   ▼
         Predicted noise ε̂

Each ResBlock applies FiLM conditioning from the concatenated (timestep, obs) vector.

Setup

Requirements: Windows/Linux, Python 3.10, CUDA optional (trains on CPU in ~55 min/30 epochs).

conda create -n diffusemanip python=3.10
conda activate diffusemanip

pip install torch torchvision
pip install robosuite==1.5.2
pip install mujoco==2.3.7          # pin: 3.x breaks robosuite 1.5 OSC controller
pip install robomimic --no-deps    # --no-deps avoids Linux-only egl_probe on Windows
pip install h5py numpy

Download the Lift dataset:

python -m robomimic.scripts.download_datasets --tasks lift --dataset_types ph

Usage

Train:

python train.py --hdf5 data/lift/ph/low_dim_v141.hdf5 --task Lift
# Resume from checkpoint:
python train.py --hdf5 data/lift/ph/low_dim_v141.hdf5 --task Lift --resume runs/<name>/last.ckpt

Evaluate:

python eval.py --checkpoint runs/<name>/best.ckpt --task Lift --n-rollouts 50
# Save rollout GIFs:
python eval.py --checkpoint runs/<name>/best.ckpt --task Lift --n-rollouts 3 --save-videos

Run tests:

python test_windowing.py          # data pipeline unit tests (no deps beyond numpy)
python test_diffusion_policy.py   # model architecture unit tests (PyTorch only)

Files

File	Description
datasets.py	HDF5 data loader, sliding-window sampler, per-dimension normalizer
obs_encoders.py	LowDimEncoder (MLP)
diffusion_policy.py	SinusoidalPosEmb, ResidualBlock1D (FiLM), ConditionalUNet1D, GaussianDiffusion, DDIM
train.py	Training loop with EMA, AdamW, checkpointing; --obs-mode {lowdim,image}
eval.py	Rollout harness: receding-horizon control, success detection, GIF export
obs_pipeline.py	M2: one obs path for train+eval (low-dim reconstruction, image pipeline, shared run_rollout)
vision.py	M2: ResNet-18 + GroupNorm + spatial-softmax MultiCameraEncoder
image_dataset.py	M2: lazy per-item HDF5 image windows, worker-safe crop augmentation
scripts/regenerate_image_obs.py	M2: replay states → image+proprio HDF5 under robosuite 1.5.2
test_windowing.py	8-check unit tests for the data pipeline
test_diffusion_policy.py	8-check unit tests for the model (shapes, gradients, DDIM)
test_vision.py / test_image_dataset.py	M2: encoder (5) + image-dataset (7) unit tests
demo_diffusion.py	Animated GIF of DDIM denoising mechanism (model architecture demo)

Key Implementation Notes

Observation space (19-dim, concatenated in order):

• object: cube position (3) + cube quaternion (4) + eef-to-cube vector (3) = 10-dim
• robot0_eef_pos: end-effector position (3)
• robot0_eef_quat: end-effector orientation (4)
• robot0_gripper_qpos: gripper joint positions (2)

Action space: 7-dim OSC_POSE — (ΔX, ΔY, ΔZ, Δroll, Δpitch, Δyaw, gripper)

Hyperparameters (paper defaults):

• Pred horizon Tp=16, obs horizon To=2, action horizon Ta=8
• DDPM T=100, DDIM steps=16 at inference
• AdamW lr=1e-4, batch=256, grad clip=1.0
• EMA decay=0.999

Bugs Fixed Along the Way

Bug	Symptom	Fix
EMA decay too high (0.9999)	75% of shadow was random init → 2% success	Lower to 0.999; auto-detect and skip bad EMA in eval
robosuite 1.4 vs 1.5 obs sign flip (Lift)	object-state dims 7-9 have opposite sign → OOD inputs	Negate dims 7-9 in extract_obs()
mujoco 3.x API break	mj_fullM signature changed, OSC controller crashes	Pin mujoco==2.3.7
PyTorch 2.6 weights_only default	Checkpoint load fails with numpy arrays	Add weights_only=False
robosuite 1.5 load_controller_config moved	ImportError on eval startup	Remove import; suite.make() uses default config automatically
Windows Smart App Control	mujoco.dll blocked (WinError 4551)	Disable SAC in Windows Security settings
robosuite 1.4 vs 1.5 object-state layout (Square)	v141 layout is [nut_pos(3), nut_quat(4), rel_pos_eef_frame(3), rel_quat(4)]; robosuite 1.5 returns a reordered vector AND zeros at reset (obs cache empty) → policy misreads "gripper at nut" on step 0	Reconstruct v141-compatible 14-dim obs from individual named keys (SquareNut_pos, SquareNut_quat, robot0_eef_pos, robot0_eef_quat)
Square relative-position frame bug (the big one)	dims 7-9 are the nut position relative to the eef in the eef's local frame R(q_eef)ᵀ·(nut−eef), but reconstruction used the world frame eef−nut → 3/23 obs dims OOD every rollout, normalized obs hit −2.6, success capped ~40%	Rotate the relative position into the eef frame in _reconstruct_nut_object_obs (eval.py) and _reconstruct_nut_obs (train.py); verified to match v141 to 0.000000. 40% → 78%
EMA mode collapse on short training (Square)	EMA at epoch 39 (4040 steps) averages early "go to peg" behavior with late "grasp then insert" behavior → robot skips pickup, 0% success	Disable EMA for checkpoints trained <10k steps via --no-ema flag
Headless OSMesa render segfault (M2, no GPU)	MUJOCO_GL=osmesa crashed in GL-context init: system mesa's libLLVM needs GLIBCXX_3.4.32 newer than conda's libstdc++; preloading the system lib mixed ABIs → llvmpipe segfault	conda install -c conda-forge mesalib + libOSMesa.so symlink; run image tools with MUJOCO_GL=osmesa PYOPENGL_PLATFORM=osmesa LD_LIBRARY_PATH=<env>/lib

Roadmap

• M0: Environment setup (robosuite + robomimic + MuJoCo)
• M1a: Diffusion Policy on Lift, low-dim obs — 100% success
• M1b: Diffusion Policy on Square — 78% success (eef-frame obs bug fixed; see note above)
• M1b (continued): Can task
• [~] M2: Image observations — pipeline built & unit-tested; training pending GPU (see below)

M2: from state to pixels

The image milestone swaps the low-dim state vector for two RGB cameras (agentview + robot0_eye_in_hand) + 9-dim proprioception. The diffusion core (1-D temporal U-Net, DDPM/DDIM, EMA) is unchanged — only the conditioning path is new. Status: all code is implemented and unit-tested; the actual training runs are GPU-gated (this dev box is CPU-only — a single image epoch is ~8 min, so hundreds of epochs are impractical here) and are deferred.

Design choices (following Chi et al. 2023):

• Regenerate under the eval sim. scripts/regenerate_image_obs.py replays each demo's stored MuJoCo states through the installed robosuite 1.5.2 and reads images straight from the env obs dict — the same path eval uses. Train/eval pixels then agree by construction. This is the M1b eef-frame lesson in pixel form: proprio reproduces v141 to <1e-5, and a verify_pixel_alignment check guards eval.
• One obs pipeline. All obs handling (low-dim reconstruction and image crop/scale/history) lives in obs_pipeline.py, consumed by both training rollouts and eval — no more duplicated logic to patch twice.
• ResNet-18 from scratch, GroupNorm, spatial softmax. vision.py; BatchNorm is banned (assert_no_batchnorm in tests and at train start) because BN + EMA + temporally correlated batches silently degrades DP. One encoder per camera.
• Random crop 84→76 (train) / center crop (eval) as the sole image augmentation.
• Self-describing checkpoints. obs_mode, cameras, crop, and the normalizer live in the checkpoint, so eval reconstructs the exact pipeline with no CLI drift.

Run (image mode needs the OSMesa render prefix — see the bugs table):

python scripts/regenerate_image_obs.py --source data/lift/ph/low_dim_v141.hdf5 \
    --out data/lift/ph/image_v15.hdf5 --env Lift \
    --cameras agentview robot0_eye_in_hand --size 84
python train.py --hdf5 data/lift/ph/image_v15.hdf5 --task Lift --obs-mode image \
    --batch-size 64 --num-workers 4
python test_vision.py && python test_image_dataset.py   # encoder + dataset unit tests

Reference

Chi, C., Feng, S., Du, Y., Xu, Z., Cousineau, E., Burchfiel, B., & Song, S. (2023). Diffusion Policy: Visuomotor Policy Learning via Action Diffusion. RSS 2023. arXiv:2303.04137