PushT-with-obstacles — Columbia Diffusion Policy Finetune

A finetune of Columbia's official Diffusion Policy on a custom PushT environment with random obstacles. Trained for 500 epochs (~115k optimizer steps) on a single RTX 5080, starting from Columbia's train_2/epoch=1850/test_mean_score=0.898 PushT image-policy checkpoint.

Source code: https://github.com/Zengxy0624/xarm7_collect/tree/main/push_t/columbia_ft

Headline numbers (n=200 final eval)

Metric	Value	95% CI
`success_rate` (max coverage ≥ 0.95)	0.490	[0.42, 0.56]
`obstacle_hit_rate`	0.175	—
`mean_score` (mean max coverage)	0.676	—

Statistically tied with the parallel lerobot finetune baseline (success 0.465 ± 0.07). Bottleneck identified as dataset size (101 episodes), not training stack — see docs/03_lerobot_vs_columbia.md.

File index

`checkpoints/`

columbia_ft_best_epoch300_score0.695.ckpt — Best ckpt picked by Columbia's TopK manager (monitor=test/mean_score). Use this for inference / as a finetune starting point.
columbia_ft_latest_epoch499.ckpt — Final epoch 499 weights (overfit, val_loss 70× higher than best). Useful only for studying overfitting trajectory.

Each ckpt is ~4 GB. Format: dill-pickled torch state dict containing state_dicts.{model, ema_model, optimizer} + cfg (the full Hydra training config). Load with torch.load(path, pickle_module=dill, weights_only=False, map_location="cpu").

`data/`

pusht_obstacles_dp.zarr.tar.gz — Training data: 101 episodes, 15758 frames @ 10 fps, 96×96 RGB images with obstacles already drawn into the rendered scene. Schema matches diffusion_policy.dataset.pusht_image_dataset.PushTImageDataset:
- data/img (N, 96, 96, 3) uint8
- data/state (N, 5) float32 — [agent_x, agent_y, block_x, block_y, block_θ]
- data/action (N, 2) float32
- meta/episode_ends (E,) int64

Untar before use: tar -xzf pusht_obstacles_dp.zarr.tar.gz.

`logs/`

train_logs.jsonl — Per-step + per-epoch metrics for the full 500-epoch run. ~115k JSON lines. Each line has train_loss, global_step, epoch, lr. Lines at epoch boundaries also have val_loss, train_action_mse_error, test/mean_score, test/success_rate, test/obstacle_hit_rate. Use to plot training curves or replay rollout history.

`eval/`

columbia_ft_eval_n200.json — Final n=200 evaluation of the best ckpt against obstacle_set=legacy. Headline result. Contains per-seed reward + aggregated mean_score / success_rate / obstacle_hit_rate.
columbia_zero_shot_eval_n50.json — Zero-shot baseline: Columbia's stock PushT ckpt directly applied to the obstacle env (no finetune). success=0.06, obstacle_hit=0.72. Reference for "how much did finetune help".
columbia_official_pusht_eval_n50.json — Sanity check: Columbia's stock ckpt on Columbia's stock PushT env (no obstacles). success=0.76, mean_score=0.86. Reproduces paper Table I within 5%; verifies the training stack itself is correctly set up.

`docs/`

01_stack_setup.md — How the dp conda env was built, what versions diverge from upstream and why (torch 2.8 for sm_120 GPUs, shapely 2.x, huggingface_hub<0.26). Includes the exact reproduction commands.
02_finetune_results.md — Full training trajectory analysis. Per-rollout table, three training phases (rapid acquisition → plateau → overfit), train_loss vs val_loss divergence (val rises 70× while train falls 135×). Identifies the small-data overfitting signature.
03_lerobot_vs_columbia.md — Cross-stack comparison with the lerobot finetune baseline (n=200 vec-env eval on both). Statistical analysis of why the previously reported lerobot 0.61 was a small-sample artifact (n=100, CI ±13pp). Discussion of which stack to pick for downstream work.

Minimum loading example

import dill, torch
from diffusion_policy.workspace.train_diffusion_unet_hybrid_workspace import (
    TrainDiffusionUnetHybridWorkspace,
)

# Download from HF
from huggingface_hub import hf_hub_download
ckpt_path = hf_hub_download(
    repo_id="zengxy0624/pusht-obstacles-columbia-ft",
    filename="checkpoints/columbia_ft_best_epoch300_score0.695.ckpt",
)

# Load
payload = torch.load(ckpt_path, map_location="cpu", pickle_module=dill, weights_only=False)
cfg = payload["cfg"]  # full Hydra training config
workspace = TrainDiffusionUnetHybridWorkspace(cfg)
workspace.load_payload(payload)

# Use the EMA model for inference (this is what the rollout scores measured)
policy = workspace.ema_model
policy.eval().to("cuda")

# Inference: see push_t/columbia_ft/eval_zero_shot.py in the source repo

Reproducing the training run

# 1. Clone source
git clone https://github.com/Zengxy0624/xarm7_collect.git
cd xarm7_collect

# 2. Set up dp env (see push_t/columbia_ft/requirements_dp.txt for exact pip install commands)
conda create -n dp python=3.9 pip=22.2.2 -y
conda activate dp
pip install setuptools==65.5.0 wheel==0.38.4
pip install -r push_t/columbia_ft/requirements_dp.txt --extra-index-url https://download.pytorch.org/whl/cu128
pip install --no-deps robomimic==0.2.0

# 3. Download Columbia's PushT image-policy ckpt (used as init)
bash push_t/columbia_ft/download_columbia_ckpt.sh

# 4. Get the dataset (this HF repo)
hf download zengxy0624/pusht-obstacles-columbia-ft data/pusht_obstacles_dp.zarr.tar.gz \
    --local-dir push_t/data --local-dir-use-symlinks=False
cd push_t/data && tar -xzf pusht_obstacles_dp.zarr.tar.gz && cd ../..

# 5. Train
git clone https://github.com/real-stanford/diffusion_policy.git ~/diffusion_policy
export PYTHONPATH=$PWD:$HOME/diffusion_policy:$PYTHONPATH
cd ~/diffusion_policy
python train.py \
    --config-dir=$HOME/.../push_t/columbia_ft/config \
    --config-name=image_pusht_obstacle_ft.yaml \
    hydra.run.dir=outputs/columbia_ft_$(date +%Y%m%d_%H%M%S)

Training takes ~3.5 hours on RTX 5080 (faster on H100). See docs/01_stack_setup.md for full setup notes.

Evaluating the released ckpt

# Sync vec-env eval (matches the published n=200 number)
python -m push_t.columbia_ft.eval_zero_shot \
    --checkpoint /path/to/columbia_ft_best_epoch300_score0.695.ckpt \
    --output-dir push_t/eval_results/repro_n200 \
    --n-test 200 --obstacle-set legacy

Citation / credits

Underlying algorithm: Chi et al. 2023, "Diffusion Policy: Visuomotor Policy Learning via Action Diffusion" — https://diffusion-policy.cs.columbia.edu/
Init checkpoint: Columbia official PushT image-policy run train_2/epoch=1850 from https://diffusion-policy.cs.columbia.edu/data/experiments/image/pusht/diffusion_policy_cnn/
This finetune + obstacle env extension: zengxy0624 / 2026-05

License

MIT (matches Columbia's diffusion_policy repo). The PushT environment is derived from Yunsheng Chi's original PushT.

Downloads last month: -; Downloads are not tracked for this model. How to track

Video Preview

Robotics