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 — 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).
• Framework: 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 (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):

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

Diffusion process (noise scheduler):

Normalization (normalization_mapping):

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:

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.

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

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.

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:

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.

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.

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

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 — 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), 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
• Experiments repo: github.com/Bigenlight/banana-in-pot-experiments
• ACT sibling model: Bigenlight/act_banana_in_pot
• Deployment stack (ROS 2 Humble): github.com/Bigenlight/gello_software
• Framework: LeRobot v0.6.1

Citation

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