pusht-obstacles-lerobot-finetune

Diffusion Policy checkpoints finetuned from lerobot/diffusion_pusht on a custom 101-episode PushT-with-obstacles dataset (legacy obstacle preset: circle radius=15, 1-3 obstacles per episode).

Source code: github.com/Zengxy0624/xarm7_collect (see push_t/ and push_t/doc/).

Files

File	Source run	Pipeline	Best success (n=50)	Best hit	Step
lerobot_run3.5_best.pt	jid2238141	fixed	0.62	0.12	44000
lerobot_run3.5_final.pt	jid2238141	fixed	0.62 (final eval)	0.12	50000
lerobot_run2_best.pt	jid2235676	broken (pre-fix)	0.60	0.22	38000

The "fixed" pipeline applies three corrections matching LeRobot's official training recipe (ImageNet image stats, padding-loss mask, diffusers.get_scheduler('cosine')) that the original code accidentally bypassed by constructing LeRobotDataset outside make_dataset().

See push_t/doc/ft_jid2237070_2238141.md for the full eval timeline (25 evals across 50k training steps), phase analysis, and A/B comparison vs the broken-pipeline run.

When to use which file

• lerobot_run3.5_best.pt — default; best success and lowest obstacle-hit rate.
• lerobot_run3.5_final.pt — only for resume training. Same model performance as best.pt (62% at both step 44k and step 50k); carries end-of-cosine-schedule state.
• lerobot_run2_best.pt — A/B reference. Trained with the buggy pre-fix pipeline, useful for ablations.

Loading

All ckpts contain numpy stats arrays, so PyTorch 2.6+ requires weights_only=False:

import torch
from huggingface_hub import hf_hub_download

ckpt_path = hf_hub_download(
    repo_id="zengxy0624/pusht-obstacles-lerobot-finetune",
    filename="lerobot_run3.5_best.pt",
)
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)

print(ckpt.keys())          # ['model', 'model_cfg', 'stats', 'success']
print(ckpt["success"])      # 0.62
print(ckpt["model_cfg"])    # full DiffusionConfig dict (ResNet-18 + GroupNorm + 96x96)

For full evaluation, use push_t/test.py from the source repo, which handles the processor pipeline (ImageNet stats override, deterministic cudnn) correctly.

Training details

• Base model: lerobot/diffusion_pusht (ImageNet-pretrained ResNet-18 backbone, GroupNorm, spatial-softmax 32 keypoints, U-Net 1D, DDPM 100 timesteps)
• Architecture: image_shape=[3, 96, 96], crop_shape=[84, 84], n_obs_steps=2, horizon=16, n_action_steps=8, down_dims=[512, 1024, 2048]
• Hyperparameters: bs=64, lr=3e-5, warmup_steps=100, weight_decay=1e-6, grad_clip=10.0, AMP fp16
• Schedule: linear warmup -> cosine decay over 50000 steps
• Eval: 50 episodes per checkpoint, every 2000 training steps
• Hardware: 1x H100 SXM5 on NCSA Delta AI

Limitations

• Eval n=50 gives Wilson 95% CI of approximately +/-14%. Reported success rates should be treated as point estimates with this uncertainty.
• Trained only on the legacy circular-obstacle preset. Will not generalize to the alternative mixed preset (rectangles, triangles, walls) without retraining.
• Both runs show val_loss divergence consistent with small-dataset overfitting; the 101-episode dataset is the bottleneck, not the training stack.

Citation

If you use these checkpoints, please cite:

• The base Diffusion Policy paper: Cheng Chi et al. "Diffusion Policy: Visuomotor Policy Learning via Action Diffusion" (RSS 2023, IJRR 2024).
• The LeRobot framework that produced the base ckpt and was used for these experiments.