README.md

---
license: apache-2.0
library_name: lerobot
tags:
  - robotics
  - lerobot
  - diffusion-policy
  - imitation-learning
  - ur7e
  - manipulation
pipeline_tag: robotics
datasets:
  - Bigenlight/banana_in_pot_lerobot_v3
---

# Diffusion Policy — Put the right banana in the pot (UR7e, JOINT action space)

A **Diffusion Policy** (visuomotor DDPM, 1D-conv UNet denoiser) trained by imitation
learning to perform the manipulation task *"put the right banana in the pot"* on a
**Universal Robots UR7e** arm with two RGB cameras. Actions are **7-D absolute joint
targets** (6 UR joints in radians + gripper).

- **Policy:** LeRobot `DiffusionPolicy` — per-camera **ResNet18** visual encoder
  (ImageNet-pretrained) + **SpatialSoftmax** keypoints, conditioning a **1D convolutional
  UNet** denoiser. Receding-horizon action generation: `horizon = 64`, `n_obs_steps = 2`,
  `n_action_steps = 32`.
- **Noise model:** **DDPM**, `num_train_timesteps = 100`, `beta_schedule =
  squaredcos_cap_v2`, `prediction_type = epsilon` (ε-prediction), `clip_sample = true`.
- **Trained on:** [`Bigenlight/banana_in_pot_lerobot_v3`](https://huggingface.co/datasets/Bigenlight/banana_in_pot_lerobot_v3)
  — 51 teleoperated episodes / 21,524 frames, UR7e follower + GELLO leader, 2 RGB cameras.
- **This checkpoint:** step **80,000** (best by open-loop held-out MAE; see
  [Results](#results--the-headline-finding)).
- **Framework:** [LeRobot](https://github.com/huggingface/lerobot) v0.6.1.

> **Headline finding (read this first):** on a held-out split the diffusion **denoising
> `eval_loss` ROSE ~5× (0.029 → 0.149)** over training, which naively screams "severe
> overfitting". But the deployment-relevant **open-loop rollout MAE kept IMPROVING** to
> step 80k (0.119 → 0.085 rad). **For a Diffusion Policy the held-out denoising loss is a
> misleading overfit/early-stop signal — select checkpoints by open-loop MAE, not by
> `eval_loss`.** (Contrast the ACT sibling, where the two signals agreed.)

---

## Task & data

**"put the right banana in the pot."** The tabletop holds several distractor objects —
**two bananas, an apple, carrots/peppers, and a slice of watermelon** — plus a **silver
pot**. The operator must grasp the **RIGHT banana** (the target) and place it inside the
pot. Success = the right banana ends up inside the pot. Every demonstration is a success.

- **Dataset:** [`Bigenlight/banana_in_pot_lerobot_v3`](https://huggingface.co/datasets/Bigenlight/banana_in_pot_lerobot_v3)
  (LeRobot v3.0 format).
- **Scale:** **51 episodes / 21,524 frames / 30 fps / ~12 min.**
- **Action / state space:** 7-D absolute joint (6 UR joints in radians + gripper), i.e.
  `[cmd1..cmd6, grip_cmd]`. The gripper channel is effectively binary (open/close).
- **Cameras:** two RGB viewpoints (Intel RealSense D435 + D435if), captured at 1280×720
  (720p) @ 30 fps, **RGB only** (no depth / IR). `cam1 ↔ cam2` order is fixed and must be
  preserved at deploy time.

### Train / held-out split

Training used `--dataset.eval_split=0.117`, which holds out the **LAST
`ceil(51 × 0.117) = 6` episodes (indices 45–50)** as a true validation split and trains on
the other **45** episodes (0–44). The held-out episodes 45–50 are used both for the
in-training denoising `eval_loss` probe and for all offline open-loop evaluation below.

---

## Model architecture

LeRobot `DiffusionPolicy`. All values below are quoted directly from the checkpoint's
`config.json` / `train_config.json`.

**Observation encoder (vision):**

| Item | Value |
|---|---|
| Vision backbone | `resnet18` |
| Pretrained weights | `ResNet18_Weights.IMAGENET1K_V1` (ImageNet) |
| Per-camera encoder | `use_separate_rgb_encoder_per_camera = true` (separate ResNet18 per view) |
| Pooling | **SpatialSoftmax**, `spatial_softmax_num_keypoints = 32` |
| Group norm in encoder | `use_group_norm = false` (keeps BatchNorm from the pretrained backbone) |
| Cameras | 2 × RGB (`observation.images.cam1`, `observation.images.cam2`) |
| Network input resolution | **360 × 640** (`resize_shape = [360, 640]`; see [why](#why-the-two-non-default-flags-are-required)) |
| Crop | **OFF** — `crop_shape = null`, `crop_ratio = 1.0` (`crop_is_random` is moot with no crop) |
| State input | `observation.state`, shape `(7,)` |

> Note: `config.json` records the raw dataset image feature shape as `[3, 720, 1280]`, but
> the on-the-fly `Resize` to `[360, 640]` (`resize_shape`) means the network actually sees
> **360 × 640** frames. See the training section for why this must match at inference.

**Denoiser (conditional 1D-conv UNet):**

| Item | Value |
|---|---|
| Denoiser | 1D convolutional UNet (Diffusion Policy / Janner-style) |
| `down_dims` | `[512, 1024, 2048]` |
| `kernel_size` | `5` |
| `n_groups` (GroupNorm) | `8` |
| `diffusion_step_embed_dim` | `128` |
| FiLM conditioning | `use_film_scale_modulation = true` |
| `horizon` | `64` (prediction horizon, in frames) |
| `n_obs_steps` | `2` (observation context length) |
| `n_action_steps` | `32` (actions executed before replanning) |

**Diffusion process (noise scheduler):**

| Item | Value |
|---|---|
| `noise_scheduler_type` | `DDPM` |
| `num_train_timesteps` | `100` |
| `beta_schedule` | `squaredcos_cap_v2` (cosine) |
| `beta_start` / `beta_end` | `0.0001` / `0.02` |
| `prediction_type` | `epsilon` (predict noise) |
| `clip_sample` | `true`, `clip_sample_range = 1.0` |
| `num_inference_steps` | `null` in config → defaults to the full DDPM schedule at inference unless overridden (evals here used **DDIM, 10 steps** for speed — see [Usage](#usage--inference)) |

**Normalization (`normalization_mapping`):**

| Feature group | Mode |
|---|---|
| `VISUAL` (images) | `MEAN_STD` (ImageNet stats, `use_imagenet_stats = true`) |
| `STATE` (observation.state) | `MIN_MAX` |
| `ACTION` (action) | `MIN_MAX` |

Normalizer statistics are baked into the pre/post-processor pipelines saved alongside the
checkpoint (`policy_preprocessor.json` / `policy_postprocessor.json`), not into
`forward()`.

**I/O summary:**

| I/O | Spec |
|---|---|
| `observation.state` | `(7,)` — UR joints `q1..q6` (radians) + gripper position |
| `observation.images.cam1` / `cam2` | RGB, network input **360 × 640** |
| `action` | `(7,)` — `[cmd1..cmd6, grip_cmd]`, **absolute** joint targets (radians) + ~binary gripper |

---

## Training setup

Trained with `lerobot-train` (LeRobot 0.6.1). Exact invocation:
`train_diffusion_joint_valdiag.sh`. Values below are from that script and the saved
`train_config.json`.

| Item | Value |
|---|---|
| Policy | `diffusion` (`--policy.type=diffusion`) |
| Dataset | `banana_in_pot_lerobot_v3`, `--dataset.eval_split=0.117` (holds out eps 45–50) |
| Batch size | **8** |
| Steps | script requested 100,000; this release / analysis is the run **to 80,000** (`steps: 80000` in `train_config.json`); checkpoints saved every 10,000 |
| Optimizer | **Adam**, `lr = 1e-4`, `betas = [0.95, 0.999]`, `eps = 1e-8`, `weight_decay = 1e-6` |
| Grad clip | `grad_clip_norm = 10.0` |
| LR scheduler | **cosine** (`scheduler.type = diffuser`, `name = cosine`), `num_warmup_steps = 500` |
| Seed | **1000** |
| Precision | **fp32** (`use_amp = false`) |
| EMA | **none** (no EMA weights in this config) |
| Image transform | on-the-fly `Resize` to `[360, 640]`, `max_num_transforms = 1`, deterministic |
| `drop_n_last_frames` | **31** (non-default; see below) |
| `resize_shape` | `[360, 640]` (non-default; see below) |
| Workers | `num_workers = 4`, `prefetch_factor = 4`, `persistent_workers = true` |
| Eval probe | held-out denoising `eval_loss` every 2,000 steps; train loss logged every 200 |
| GPU | single **RTX 3060 12 GB**, ~**9.7 GB** used, ~**2.2 step/s** |
| W&B | disabled |

### Why the two non-default flags are required

Both `resize_shape=[360,640]` and `drop_n_last_frames=31` are **not** the LeRobot defaults;
they are mandatory for this dataset/config and encode real operational knowledge:

1. **`resize_shape=[360,640]` — SpatialSoftmax is shape-rigid.** The Diffusion Policy RGB
   encoder ends in a `SpatialSoftmax` layer whose keypoint geometry is fixed to the spatial
   dimensions of the feature map at build time. The network must therefore be *constructed*
   for the exact input resolution it will ever see. Setting `resize_shape=[360,640]` builds
   the encoder for 360×640 and — combined with crop being **off** (`crop_shape=null`) —
   guarantees the training image path, the offline-eval image path, and any deploy image
   path all feed the encoder identically. A mismatched resolution (or leaving crop on)
   changes the SpatialSoftmax grid and breaks the model. The same 360×640 `Resize` is
   reproduced in `eval_offline.py` (`build_image_transforms`).

2. **`drop_n_last_frames=31` — horizon / n_action off-by-one at episode ends.** The
   trajectory sampler must not draw a window that runs past the end of an episode. With
   `horizon=64`, `n_obs_steps=2`, and `n_action_steps=32`, the last valid start frame in an
   episode has to leave room for the horizon, so the correct number of trailing frames to
   drop is `horizon - n_action_steps - (n_obs_steps - 1) = 64 - 32 - 1 = 31`. Using the
   default (7, tuned for the reference `horizon=16` config) would let the sampler pull
   frames off the end of an episode and corrupt the action targets. **If you change
   `horizon`/`n_obs_steps`/`n_action_steps`, recompute `drop_n_last_frames`.**

---

## Results & the headline finding

![Diffusion joint: held-out denoising eval_loss (blue, rising) vs open-loop rollout MAE (improving to 80k)](assets/diffusion_joint_overfit_diag.png)

Offline **open-loop** evaluation on the held-out episodes **45–50** with `eval_offline.py`
(each logged observation is fed to `select_action`; the predicted action is compared to the
dataset ground truth). Sampling used **DDIM with 10 inference steps**
(`--scheduler DDIM --num-inference-steps 10`) for ~10× faster rollouts; DDIM is a valid
sampler for a DDPM-trained ε model (same beta schedule). `poseMAE` is the mean absolute
error over the 6 joint dims (radians); `gripAcc` is the binary gripper-open/close accuracy
(threshold 0.5); `overall L1` averages all 7 dims.

| checkpoint | poseMAE (rad) | gripAcc | overall L1 |
|---|---|---|---|
| 10k | 0.1193 | 0.729 | 0.1454 |
| 20k | 0.1037 | 0.888 | 0.1078 |
| 30k | 0.0921 | 0.919 | 0.0928 |
| 40k | 0.0907 | 0.949 | 0.0862 |
| 50k | 0.0865 | 0.942 | 0.0832 |
| 60k | 0.0849 | 0.944 | 0.0812 |
| 70k | 0.0855 | 0.951 | 0.0809 |
| **80k** ⭐ | **0.0845** | **0.953** | **0.0796** |

**Best checkpoint = 80k** (tied-best poseMAE, best gripAcc). Open-loop poseMAE improves
monotonically then **plateaus at ~0.085 rad from 60k onward** (60k/70k/80k =
0.0849/0.0855/0.0845, within eval noise); gripper accuracy climbs all the way to **0.953 @
80k**. There is **no open-loop overfitting through 80k**.

### The misleading `eval_loss` (the lesson)

During training the held-out **denoising `eval_loss`** (LeRobot's in-training validation
probe, computed under `policy.eval()` on eps 45–50) did the opposite of the rollout metric:

| step | train loss | held-out eval_loss |
|---|---|---|
| 2k | 0.0310 | 0.0361 |
| 4k | 0.0240 | 0.0289 (min region) |
| 6k | 0.0210 | 0.0303 |
| 20k | 0.0120 | 0.0399 |
| 40k | 0.0080 | 0.0660 |
| 60k | 0.0050 | 0.1272 |
| 80k | 0.0050 | ~0.1487 |

Read naively, the held-out `eval_loss` bottoms near step 4k–6k and then rises ~5×, so an
early-stop rule would pick **~step 6k** and declare "severe overfit". **That
recommendation is wrong for deployment:** the same held-out episodes, evaluated by
open-loop rollout, get *monotonically better* out to 80k.

**Why:** the denoising loss is a per-sample ε-regression on a *randomly re-sampled noise
vector and diffusion timestep at every forward pass* — it is (a) high-variance/stochastic
by construction and (b) only loosely coupled to closed-loop action quality. As the model
sharpens its learned action distribution, the average ε-MSE on unseen frames can rise even
while the *sampled* action trajectories become more accurate. **Takeaway: for a Diffusion
Policy, select checkpoints and early-stop by open-loop rollout MAE, not by held-out
denoising `eval_loss`.** (The ACT sibling did not show this divergence — there the two
signals agreed — so this is a diffusion-specific pitfall.)

---

## Usage / inference

### Load the policy (LeRobot 0.6.1)

Normalization is **not** baked into `forward()` in LeRobot 0.6.1 — it lives in the
pre/post-processor pipelines saved with the checkpoint. `select_action` returns a
**normalized** action; the post-processor converts it back to radians.

```python
import torch
from lerobot.configs import PreTrainedConfig
from lerobot.policies.factory import get_policy_class, make_pre_post_processors

CKPT = "Bigenlight/diffusion_banana_in_pot_joint"
device = "cuda"

# (optional) speed up sampling: DDIM with 10 steps instead of the full DDPM schedule.
cfg = PreTrainedConfig.from_pretrained(CKPT)
cfg.pretrained_path = CKPT
cfg.device = device
cfg.noise_scheduler_type = "DDIM"     # valid sampler for a DDPM-trained epsilon model
cfg.num_inference_steps = 10          # ~10x faster rollouts

policy = get_policy_class(cfg.type).from_pretrained(CKPT, config=cfg)  # -> DiffusionPolicy
policy.to(device)
policy.eval()

preprocessor, postprocessor = make_pre_post_processors(
    policy_cfg=cfg,
    pretrained_path=CKPT,
    preprocessor_overrides={"device_processor": {"device": device}},
)
```

### Run the control loop

Build the observation dict exactly as training did: joint state `(7,)` plus **both**
cameras as RGB CHW tensors in `[0, 1]`, **resized to 360×640** (aspect-preserving
half-resolution). `cam1`/`cam2` must map to the same physical viewpoints as at collection.

```python
policy.reset()          # once at the start of each episode/rollout
preprocessor.reset()
postprocessor.reset()

# obs = {
#   "observation.state":        state_7,          # (7,) float32, radians + gripper
#   "observation.images.cam1":  img1_chw,         # (3, 360, 640) float32 in [0,1]
#   "observation.images.cam2":  img2_chw,         # (3, 360, 640) float32 in [0,1]
#   "task": "put the right banana in the pot",
# }

with torch.inference_mode():
    proc   = preprocessor(obs)             # rename -> add batch dim -> device -> normalize
    action = policy.select_action(proc)    # (1, 7) NORMALIZED
    action = postprocessor(action)         # (1, 7) radians, on cpu
q_target = action.squeeze(0).numpy()       # (7,) -> [cmd1..cmd6, grip_cmd]
```

`select_action` returns **one** action per call from an internal queue. Because
`n_action_steps = 32`, the policy denoises a fresh action sequence, executes 32 actions
from it, then replans (with `n_obs_steps = 2` frames of observation context). Call
`policy.reset()` at the start of every episode to clear that queue. The gripper channel
`grip_cmd` is ~binary — threshold at `> 0.5 → close` and map to your gripper driver.

### Reproduce the offline evaluation

The repo's `eval_offline.py` runs the exact open-loop protocol used for the results table
(same 360×640 `Resize`, same normalization via the saved processors):

```bash
python eval_offline.py \
  --checkpoint outputs/train/diffusion_joint_val_diag/checkpoints/080000/pretrained_model \
  --episodes 45,46,47,48,49,50 \
  --device cuda \
  --scheduler DDIM --num-inference-steps 10 \
  --out eval_out_diffusion_80k
```

`--scheduler DDIM --num-inference-steps 10` gives the ~10× rollout speedup; omit them to
sample with the full trained DDPM schedule (`num_train_timesteps = 100`).

---

## Deployment on a real UR7e

Closed-loop deployment targets a real **UR7e** through the ROS 2 Humble stack in
[**Bigenlight/gello_software**](https://github.com/Bigenlight/gello_software) — the same
stack used to collect this dataset (UR7e follower + GELLO leader, dual RealSense cameras).
LeRobot ships **no UR robot class**, so deployment requires a small **policy deploy node**
that, each control tick (target **30 Hz**):

1. reads the UR7e measured joints + gripper → `observation.state` `(7,)`;
2. grabs both camera frames, BGR→RGB, **resizes to 360×640**, CHW `[0,1]` →
   `observation.images.cam1` / `cam2`;
3. `preprocessor → policy.select_action → postprocessor` → `q_target` (7,);
4. streams `q_target[:6]` to the arm (e.g. `servoJ` via `ur_rtde` / the ROS 2 driver) and
   drives the gripper from `grip_cmd`.

This diffusion policy would need a **deploy node analogous to the ACT one**
([`Bigenlight/act_banana_in_pot`](https://huggingface.co/Bigenlight/act_banana_in_pot)),
with two differences: (a) the action queue length is `n_action_steps = 32` (not ACT's 100),
so it replans ~every 32 ticks; and (b) inference runs a diffusion sampler — use **DDIM /
10 steps** to keep per-replan latency low enough for 30 Hz.

**Safety — actions are ABSOLUTE joint positions:**

1. **Start near the dataset initial pose** before enabling the policy, or the first absolute
   command is a large jump.
2. **First-command jump guard:** if `max(|q_target − getActualQ()|)` exceeds a small
   threshold (~0.15 rad), **abort**.
3. **Clamp per-tick joint change** and clamp to UR software joint limits; run at reduced
   speed for first trials with a hand on the **E-stop**.
4. **`cam1`/`cam2` mapping is fixed** — swap the two views and the policy fails silently.
   Verify wiring every session.

---

## Limitations & intended use

- **Small, single-task lab dataset:** 51 demonstrations, one scene layout, one operator.
  Expect limited generalization to novel object arrangements, lighting, or camera placement.
- **Success-only demonstrations:** no failure/recovery data; not suited as-is for methods
  that need negative examples.
- **Offline metrics only:** the best checkpoint (80k) reaches **held-out poseMAE ≈ 0.085
  rad** and gripper accuracy ≈ 0.953 in open-loop rollout. These are *not* closed-loop task
  success rates — real closed-loop success on hardware has not been measured here and must
  be validated on the arm.
- **Absolute-joint action space** demands the safety guards above; the policy was only ever
  conditioned on states near the data-collection start pose.
- **Not for production.** Intended for research in imitation learning / diffusion policies
  for robot manipulation. Workspace-, robot-, and camera-specific.
- The ResNet18 encoders are **ImageNet-pretrained** (not robotics-pretrained); the UNet
  denoiser is trained from scratch on this task.

---

## Links

- **Dataset:** [`Bigenlight/banana_in_pot_lerobot_v3`](https://huggingface.co/datasets/Bigenlight/banana_in_pot_lerobot_v3)
- **Experiments repo:** [github.com/Bigenlight/banana-in-pot-experiments](https://github.com/Bigenlight/banana-in-pot-experiments)
- **ACT sibling model:** [`Bigenlight/act_banana_in_pot`](https://huggingface.co/Bigenlight/act_banana_in_pot)
- **Deployment stack (ROS 2 Humble):** [github.com/Bigenlight/gello_software](https://github.com/Bigenlight/gello_software)
- **Framework:** [LeRobot](https://github.com/huggingface/lerobot) v0.6.1

## Citation

```bibtex
@misc{theo2026bananainpotdiffusion,
  title        = {Diffusion Policy for "put the right banana in the pot"
                  (UR7e, joint action space)},
  author       = {Theo and {Bigenlight}},
  year         = {2026},
  howpublished = {\url{https://huggingface.co/Bigenlight/diffusion_banana_in_pot_joint}},
  note         = {LeRobot 0.6.1 DiffusionPolicy, trained on banana_in_pot_lerobot_v3}
}
```

License: **Apache-2.0**.