Diffusion Policy (UNet + DDPM) — R1 Lite "insert flower"
Baseline checkpoints for the ODE-policy comparison on the Galaxea R1 Lite right arm. This is the original
Diffusion Policy recipe (Chi et al., 2023): a 1-D conditional UNet predicts the noise of a 16-step
action window (8 past + 8 future) under a DDPM schedule, conditioned on 8 frames of two cameras and
proprioception. Trained with exactly the same data, encoder, batch, learning rate and schedule as the ODE
policy (at237299966/ODE_policy_insert_flowers)
and the flow-matching baseline (at237299966/fm_dit_insert_flowers).
Files
| Path | Content |
|---|---|
checkpoints/{50,100,150,200,250,300}.ckpt |
Full training checkpoints (755 MB each): model, ema_model, optimizer state dicts + the pickled Hydra config. Serve the EMA weights (the loader's default). |
config/config.yaml, config/overrides.yaml |
Resolved training config and the command-line overrides of this run |
logs/logs.json.txt |
Per-step / per-epoch training log (loss, lr, epoch time, val_loss every 5 epochs) |
SHA256SUMS |
Checksums of every uploaded file |
Checkpoints are torch.save(..., pickle_module=dill) payloads written by the ODE_Policy training runner;
loading needs that codebase (roboverse_learn.il.policies.dp.ddpm_unet_image_policy.DiffusionUnetImagePolicy)
and torch.load(..., pickle_module=dill, weights_only=False) — only load checkpoints you trust.
Data
at237299966/r1lite_insert_flower_single:
LeRobot v2.1, 95 episodes / 23,034 frames at 15 Hz, single task "insert flower".
| Input / output | Definition |
|---|---|
| Cameras | head_rgb (720×1280) and right_wrist_rgb (360×640), padded to square then resized to 256×256 (same geometry as the robot client) |
agent_pos |
right-arm joint positions (6, rad) + right gripper opening (1, mm) |
| Action | right-arm absolute joint targets (6, rad) + gripper command (1, mm; 1 = closed, 100 = open), as recorded |
| Window | 8 observation frames; the policy predicts 16 steps and rows 7..14 (the current step and the next 7) are the action chunk |
Training
| Model | ResNet18 (GroupNorm, random crop) per camera + proprioception, 8 frames → global condition (FiLM); ConditionalUnet1D with down_dims=[48, 96, 192], kernel 5 (24.80M parameters outside the 22.35M encoder) |
| Objective | DDPM, 100 train timesteps, squared-cosine schedule, ε-prediction, sample clipping; 100 denoising steps at inference |
| Optimisation | AdamW, lr 1.5e-4, cosine schedule over 300 epochs, warm-up 500 steps, EMA |
| Batch | 2 × A100-40GB (DDP), 48 per GPU = 96 global; 236 steps/epoch, 70,800 steps in total |
| Wall clock | 9 h 12 min (110 s/epoch) |
| epoch | 0 | 10 | 50 | 100 | 150 | 200 | 250 | 300 |
|---|---|---|---|---|---|---|---|---|
| train loss | 0.82 | 0.034 | 0.013 | 0.0071 | 0.0048 | 0.0026 | 0.0015 | 0.0012 |
| val loss | 0.54 | 0.037 | 0.022 | 0.029 | 0.046 | 0.046 | 0.118 | 0.068 (ep 295) |
The validation loss (one held-out episode, raw weights) is lowest at epoch 40 (0.0185) and rises afterwards, so compare early and late checkpoints on the robot. Loss values are not comparable across methods.
Deployment
Same server and wire protocol as the ODE policy; the robot side is the unmodified openpi-r1lite closed-loop
client (15 Hz, executes row 0 of each (8, 14) reply):
python -m roboverse_learn.il.deploy.serve_policy \
--checkpoint checkpoints/100.ckpt --host 0.0.0.0 --port 8000 --device cuda:0 --replan_every 8
Latency caveat: with the 100-step DDPM sampler one prediction takes ~667 ms on an A100, ten times the
66.7 ms tick of the 15 Hz loop, so the robot would stall on every re-planning tick. Real-time use needs a
faster sampler on the same weights (e.g. DDIM with ~10–16 steps, as in the original Diffusion Policy
deployment); the server does not offer one yet. Code: ODE_Policy repository, neural_ode/real/ and
roboverse_learn/il/deploy/.