README.md

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

# Diffusion Policy — Put the right banana in the pot (UR7e, END-EFFECTOR 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 **10-D absolute
end-effector (EE) targets**: `action = [x, y, z (meters), r1..r6 (6-D rotation), grip]`.
The observation state is likewise **10-D** in the same EE parameterization.

- **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
  (EE-pose action/state variant of the task).
- **This checkpoint:** step **100,000** — the best of the 80k–100k plateau: 100k has the
  **best gripper accuracy (0.966)** and the **best overall L1**, and its `poseMAE` is tied
  within noise with the 80k minimum (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 ~4.5× (0.0250 @8k → 0.112 @100k)**, and LeRobot's auto-report flags this
> as "overfit from ~8k". That verdict is **misleading**. The deployment-relevant **open-loop
> rollout MAE kept IMPROVING then plateaued** (poseMAE 0.0665 → ~0.0368; gripAcc → 0.966) —
> **no destructive overfit through 100k.** **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`.** (Exactly the lesson from the [JOINT sibling](https://huggingface.co/Bigenlight/diffusion_banana_in_pot_joint).)

---

## 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:** 10-D absolute end-effector pose,
  `[x, y, z (meters), r1..r6 (6-D rotation, unitless), grip]`. The gripper channel is
  effectively binary (open/close). This is the **EE** variant; the sibling model uses a
  7-D joint action instead.
- **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 holds out the **LAST 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`. The architecture is **identical to the JOINT sibling** except that the first
(state) and last (action) layers are 10-D instead of 7-D.

**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 `(10,)` |

> 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) |

> **Conditioning dimension (EE-specific):** each observation contributes
> `state 10 + vision 128 = 138-D` (per-camera ResNet18 → 64-D, two cameras → 128-D); with
> `n_obs_steps = 2` the UNet `global_cond` is **276-D** (the JOINT model's is 270-D). This
> 6-D difference is the only structural change from the JOINT config.

**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` | `(10,)` — `[x, y, z (m), r1..r6 (6-D rotation), grip]` |
| `observation.images.cam1` / `cam2` | RGB, network input **360 × 640** |
| `action` | `(10,)` — `[x, y, z (m), r1..r6 (6-D rotation), grip]`, **absolute** EE pose + ~binary gripper |

**Size:** ~**277.9M** parameters, fp32 ~**1.11 GB** — effectively identical to the JOINT
model (they differ only by the 3-D width of the first and last layers).

---

## Training setup

Trained with `lerobot-train` (LeRobot 0.6.1). Values below are from the saved training
config; hyperparameters match the JOINT run.

| Item | Value |
|---|---|
| Policy | `diffusion` (`--policy.type=diffusion`) |
| Dataset | `banana_in_pot_lerobot_v3` (EE action/state), holds out eps 45–50 |
| Batch size | **8** |
| Steps | **100,000**; checkpoints saved every 10,000 |
| Optimizer | **Adam**, `lr = 1e-4`, `betas = [0.95, 0.999]`, `eps = 1e-8`, `weight_decay = 1e-6` |
| LR scheduler | **cosine** (`scheduler_name = cosine`), `num_warmup_steps = 500` |
| Precision | **fp32** (`use_amp = false`) |
| EMA | **none** (no EMA weights in this config) |
| Image transform | on-the-fly `Resize` to `[360, 640]`, deterministic |
| `drop_n_last_frames` | **31** (non-default; see below) |
| `resize_shape` | `[360, 640]` (non-default; see below) |
| Control rate | dataset 30 fps, consecutive frames (step = 1) — no downsampling. obs 2 frames (~0.067 s) / horizon 64 (2.13 s) / exec 32 (1.07 s) @ 30 Hz |
| Eval probe | held-out denoising `eval_loss`; open-loop rollouts run offline per checkpoint |

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

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 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 EE: held-out denoising eval_loss (rising, misleading) vs open-loop rollout MAE (improving then plateauing through 100k)](assets/diffusion_ee_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 first 9 pose dims; `gripAcc` is the binary gripper-open/close accuracy
(threshold 0.5); `overall L1` averages all 10 dims. Numbers persisted at
[`assets/diffusion_ee_openloop_eval.csv`](assets/diffusion_ee_openloop_eval.csv).

| checkpoint | poseMAE (own scale ↓) | gripAcc ↑ | overall L1 ↓ |
|---|---|---|---|
| 10k | 0.06648 | 0.886 | 0.07517 |
| 20k | 0.04499 | 0.923 | 0.05031 |
| 30k | 0.04345 | 0.930 | 0.04839 |
| 40k | 0.03999 | 0.926 | 0.04440 |
| 50k | 0.03786 | 0.942 | 0.04126 |
| 60k | 0.03900 | 0.950 | 0.04131 |
| 70k | 0.03687 | 0.956 | 0.03846 |
| 80k | **0.03674** ⟵ poseMAE min | 0.961 | 0.03773 |
| 90k | 0.03718 | 0.955 | 0.03856 |
| **100k** ⭐ | 0.03717 | **0.966** ⟵ max | **0.03754** ⟵ min |

> **Units caveat — do NOT compare `poseMAE` across the two models.** The EE `poseMAE`
> mixes **meters** (x, y, z) with **6-D rotation** (unitless) channels, so it is on a
> completely different scale from the JOINT model's **radian** `poseMAE` (0.0845). The two
> numbers are not comparable — judge the EE model **only** by its own trend.

**Best checkpoint = the 80k–100k plateau.** `poseMAE` falls 0.0665 → ~0.0368 and then
flattens from ~70k (70k/80k/90k/100k all within ±0.0004 = eval noise); its strict minimum is
**0.03674 @ 80k**. `gripAcc` rises essentially monotonically to **0.966 @ 100k**;
`overall L1` is lowest at **100k (0.03754)**. There is **no open-loop overfitting through
100k**. We upload **100k** because it gives the best gripper and best overall L1 while its
`poseMAE` is tied within noise with the 80k minimum — deploy 100k is safe and strictly
best-gripper. (Pick 80k if you want the strict `poseMAE` minimum; both sit on the plateau.)

### 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 | held-out eval_loss |
|---|---|
| 2k | 0.0331 |
| **8k** | **0.0250** (min) |
| 100k | 0.112 (~4.5× above the min) |

Read naively, the held-out `eval_loss` bottoms at **step 8k** and then rises ~4.5×, and
LeRobot's auto-report calls it **"overfit from ~8k"**. **That verdict is wrong for
deployment:** the same held-out episodes, evaluated by open-loop rollout, get
*monotonically better* out to the 70k–100k plateau.

**Why the two signals disagree:** a diffusion policy is trained to predict the noise added
at a **random** timestep, and `eval_loss` scores exactly that random-timestep
noise-prediction on held-out frames — so it is (a) high-variance/stochastic by construction
and (b) only loosely coupled to closed-loop action quality. But what drives the robot is the
**sampled action** — the *integral* of the full reverse-diffusion trajectory (here DDIM-10).
Those two quantities decorrelate: the network can get "worse" at random-timestep denoising
MSE while the *sampled* action trajectory keeps getting **closer** to ground truth. **The
only faithful held-out metric is to actually sample actions and compare them open-loop**,
which is what `eval_offline.py` does. **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, reconfirmed here on the EE action space.)

---

## Deployment note — research artifact, NOT wired to the robot

**Read this before trying to run the model on hardware.** Unlike the JOINT sibling — whose
7-D joint actions are **directly actuatable** and which is deployed on the real UR7e via the
[`Bigenlight/gello_software`](https://github.com/Bigenlight/gello_software) ROS 2 stack —
this EE model outputs **10-D end-effector poses** (`[x, y, z, r1..r6, grip]`). Those poses
must be converted to joint commands by **inverse kinematics (IK)** before they can drive the
arm.

**There is no IK deploy path in the current stack.** Accordingly this checkpoint is uploaded
as a **research artifact / reference for the EE action space**, **NOT** wired for
real-robot inference. It was **never run on the arm** and there are **no closed-loop
task-success results** for it.

For the actuatable path, use:

- the **JOINT model** [`Bigenlight/diffusion_banana_in_pot_joint`](https://huggingface.co/Bigenlight/diffusion_banana_in_pot_joint)
  (7-D joint actions, directly commandable), and
- the **diffusion deploy node** in
  [`Bigenlight/gello_software`](https://github.com/Bigenlight/gello_software) (ROS 2 Humble,
  UR7e follower, dual RealSense) used to run the JOINT policy.

To deploy *this* EE model you would additionally have to add an IK stage (EE pose → joint
targets, with reachability/limit handling) after `select_action` — that work is out of scope
for this release.

---

## Usage / inference

The snippet below loads the policy and produces a 10-D EE action. **Note:** that action is
an end-effector pose — it needs **IK** before it can drive a robot (see the deployment note
above).

### 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 the EE units.

```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_ee"
device = "cuda"

# (optional) speed up sampling: DDIM with 10 steps instead of the full DDPM schedule.
# Mutate the config BEFORE from_pretrained so the sampler is built with these settings.
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}},
)
```

### Produce actions

Build the observation dict exactly as training did: EE state `(10,)` 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_10,         # (10,) float32: [x,y,z, r1..r6, grip]
#   "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, 10) NORMALIZED
    action = postprocessor(action)         # (1, 10) EE units, on cpu
ee_target = action.squeeze(0).numpy()      # (10,) -> [x, y, z, r1..r6, grip]
# NOTE: ee_target is an END-EFFECTOR POSE. To drive a robot you must first run IK
#       (ee_target[:9] -> joint targets); no IK deploy path ships with this model.
```

`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` is ~binary — threshold at `> 0.5 → close`.

### 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_ee_val_diag/checkpoints/100000/pretrained_model \
  --episodes 45,46,47,48,49,50 \
  --device cuda \
  --scheduler DDIM --num-inference-steps 10 \
  --out eval_out_diffusion_ee_100k
```

`--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`).

---

## Limitations & intended use

- **Research artifact only — not deployable as-is.** The 10-D EE action requires IK to
  actuate and there is **no IK deploy path** in the stack, so this model was **never run on
  hardware** and has **no closed-loop success rate**. For a directly actuatable, deployed
  policy use the [JOINT sibling](https://huggingface.co/Bigenlight/diffusion_banana_in_pot_joint).
- **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 (100k) reaches **held-out poseMAE ≈ 0.0372**
  (EE's own mixed-unit scale) and gripper accuracy ≈ 0.966 in open-loop rollout. These are
  *not* closed-loop task success rates, and the EE `poseMAE` is **not comparable** to the
  JOINT model's radian `poseMAE`.
- **Absolute EE action space:** the policy was only ever conditioned on states near the
  data-collection start pose; any real deployment would additionally need reachability /
  joint-limit handling in the IK stage plus the safety guards used for the JOINT model.
- **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)
- **JOINT sibling model (deployed, actuatable):** [`Bigenlight/diffusion_banana_in_pot_joint`](https://huggingface.co/Bigenlight/diffusion_banana_in_pot_joint)
- **Experiments repo:** [github.com/Bigenlight/banana-in-pot-experiments](https://github.com/Bigenlight/banana-in-pot-experiments)
- **Deployment stack (ROS 2 Humble, JOINT path):** [github.com/Bigenlight/gello_software](https://github.com/Bigenlight/gello_software)
- **Framework:** [LeRobot](https://github.com/huggingface/lerobot) v0.6.1

## Citation

```bibtex
@misc{theo2026bananainpotdiffusionee,
  title        = {Diffusion Policy for "put the right banana in the pot"
                  (UR7e, end-effector action space)},
  author       = {Theo and {Bigenlight}},
  year         = {2026},
  howpublished = {\url{https://huggingface.co/Bigenlight/diffusion_banana_in_pot_ee}},
  note         = {LeRobot 0.6.1 DiffusionPolicy, 10-D EE action, trained on
                  banana_in_pot_lerobot_v3; research artifact (needs IK to actuate)}
}
```

License: **Apache-2.0**.