---
license: mit
tags:
- diffusion-policy
- imitation-learning
- pusht
- robotics
- obstacle-avoidance
library_name: diffusion_policy
base_model: real-stanford/diffusion_policy
---

# PushT-with-obstacles — Columbia Diffusion Policy Finetune

A finetune of [Columbia's official Diffusion Policy](https://github.com/real-stanford/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

```python
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

```bash
# 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

```bash
# 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.