zengxy0624/diffusion-pusht-obstacles

Diffusion Policy — PushT with Obstacles

Diffusion Policy finetuned from lerobot/diffusion_pusht for the PushT manipulation task with random circular obstacles: the agent pushes a T-shaped block to a goal pose while avoiding 1–3 obstacles per episode.

The base model handles standard PushT but has zero obstacle awareness (0% success, 55% obstacle-hit rate as zero-shot baseline). Finetuning on 101 obstacle-aware demonstrations recovers a working policy.

Results

Checkpoint	Success Rate	Obstacle-Hit Rate
Base (lerobot/diffusion_pusht, zero-shot on obstacle env)	0%	55%
This model (best of 30k finetune steps)	95%	0%

Evaluated on PushTObstacleEnv with 20 episodes per checkpoint, 300 max steps, success threshold 0.95 coverage.

Note: 20 episodes is a noisy estimator (Wilson 95% CI ≈ ±20%). Treat the 95% headline as approximate; a 100-episode re-evaluation is recommended.

Architecture

Inherited from the base model (no architecture changes, only weight finetuning):

Field	Value
Vision backbone	ResNet-18
Image input	3×96×96 (random-cropped to 84×84)
State input	2 (agent_pos)
Action output	2
n_obs_steps	2
horizon	16
n_action_steps	8
Diffusion timesteps	100
Parameters	262,709,026

Training

• Hardware: 1× NVIDIA H100 (NCSA Delta AI), AMP enabled
• Wall time: ~3 hours for 30k steps
• Optimizer: AdamW, β=(0.95, 0.999), wd=1e-6
• LR: 3e-5 (constant after 100-step warmup)
• Batch size: 64
• Dataset: 101 episodes / 15,758 frames @ 10 fps, recorded with mouse teleop in pusht_obstacle_env.py
• Normalization: dataset stats recomputed locally (image mean/std differ from the base PushT distribution due to obstacle pixels)

Usage

Standard inference (recommended)

from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy

policy = DiffusionPolicy.from_pretrained("zengxy0624/diffusion-pusht-obstacles")
policy.eval()

Drop-in compatible with the base model — same input/output schema, just swap the repo id.

Raw checkpoints (every 5k steps)

The full run trajectory is also stored under raw_checkpoints/ in the repo for offline evaluation, ablations, or resuming training. These are PyTorch .pt files in the original training format (NOT safetensors):

File	Contents	Size	Use case
raw_checkpoints/best.pt	model + cfg + success	~1 GB	Inference at peak success
raw_checkpoints/final.pt	model + cfg + step	~1 GB	Last training step
raw_checkpoints/step_{10,15,20,25,30}000.pt	model + optimizer + scheduler + step	~3 GB each	Resume training; per-step ablation

Download e.g. one via:

from huggingface_hub import hf_hub_download
import torch

path = hf_hub_download(
    repo_id="zengxy0624/diffusion-pusht-obstacles",
    filename="raw_checkpoints/step_20000.pt",
)
ckpt = torch.load(path, map_location="cpu", weights_only=False)
# ckpt["model"]: state_dict, ckpt["model_cfg"]: dict, ckpt["step"]: int (or "success" for best.pt)

Note: the schema (model, model_cfg, optimizer, scheduler, step, best_success) is internal to the original training script (finetune.py of this fork). Standard LeRobot tooling does NOT understand .pt files — use the config.json/model.safetensors at the repo root for that.

Limitations

• Only trained on circular obstacles with radius 15 px and 1–3 per episode. Out-of-distribution obstacle counts/shapes are not handled.
• Late-training evaluation showed high variance (occasional collapses to 10% success). The released checkpoint is a single best-of-N draw and may not exactly reproduce 95% on a fresh 100-episode evaluation.
• No EMA was used during training; the base lerobot/diffusion_pusht model was trained with EMA. Adding EMA is a known follow-up.

Citation

If this checkpoint is useful, please cite the original Diffusion Policy work:

@article{chi2023diffusion,
  title={Diffusion Policy: Visuomotor Policy Learning via Action Diffusion},
  author={Chi, Cheng and Feng, Siyuan and Du, Yilun and Xu, Zhenjia and Cousineau, Eric and Burchfiel, Benjamin and Song, Shuran},
  journal={The International Journal of Robotics Research},
  year={2023},
}

And LeRobot:

@misc{cadene2024lerobot,
  title={LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch},
  author={Cadene, Remi and Alibert, Simon and Soare, Alexander and Gallouedec, Quentin and Zouitine, Adil and Wolf, Thomas},
  year={2024},
  url={https://github.com/huggingface/lerobot},
}