square_real100 diffusion checkpoints
Diffusion-policy behaviour cloning on LeeHakHo/square_real100
-- 100 real Franka demonstrations of the Square (nut-on-peg) task, 84px, seq_length 16, batch 64,
1000 epochs, lr 1e-4, four arms trained concurrently on one RTX 4080. mask/train is 90
demonstrations and mask/valid the 10 held out.
There is no simulator behind this dataset, so no success rate is reported here; the only training signal is the diffusion action MSE on those 10 held-out demonstrations.
Arms
One backbone and one config; they differ only in the auxiliary objective attached to the
visual encoder. No rotation is supervised in any arm (rotation_blocks: []): the heads emit the
rotation slots but their error is masked out, so only position is learned.
| folder | aux objective |
|---|---|
baseline/ |
none |
aux_world_frame/ |
object position in the world (robot base) frame (--aux_pose) |
aux_eef_frame/ |
object position in the gripper frame (--aux_eef_frame) |
aux_obj_eef_frame/ |
one head, both of the above (--aux_obj_eef_frame) |
obs/object is FoundationPose tracking, not ground truth: the 1.8 percent of frames with
obs/aux_valid = 0 are dropped from the aux loss and from its target standardization.
Validation
Loss is the action loss for every arm (logged before the aux term is added), so the
four are directly comparable.
| arm | best val loss | at epoch | val loss @ 1000 | train loss @ 1000 |
|---|---|---|---|---|
baseline/ |
0.0698 | 67 | 0.3605 | 0.0066 |
aux_world_frame/ |
0.0719 | 99 | 0.3819 | 0.0067 |
aux_eef_frame/ |
0.0725 | 85 | 0.3916 | 0.0067 |
aux_obj_eef_frame/ |
0.0716 | 34 | 0.3800 | 0.0065 |
Validation bottoms out around epoch 34-99 and rises for the rest of training -- 90 real demonstrations over 1000 epochs is what that looks like. Note that for diffusion policies this metric is a weak proxy: the reference implementation selects checkpoints on TRAINING loss, and training past the validation minimum is the norm rather than a defect.
Epochs
model_epoch_{10,20,...,90,100,200,...,1000}.pth -- 19 per arm, dense where the validation
minimum falls. Checkpoints carry optimizer state and the EMA weights, so each is ~1.53 GiB.