EE (10-D) diffusion policy 100k + model card

ae647b1 verified 2 days ago

22.1 kB

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

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

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:

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

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.

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 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 (7-D joint actions, directly commandable), and
the diffusion deploy node in 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.

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.

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

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

Citation

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

Bigenlight
/

diffusion_banana_in_pot_ee

Diffusion Policy — Put the right banana in the pot (UR7e, END-EFFECTOR action space)

Task & data

Train / held-out split

Model architecture

Training setup

Why the two non-default flags are required

Results & the headline finding

The misleading `eval_loss` (the lesson)

Deployment note — research artifact, NOT wired to the robot

Usage / inference

Load the policy (LeRobot 0.6.1)

Produce actions

Reproduce the offline evaluation

Limitations & intended use

Links

Citation

Diffusion Policy — Put the right banana in the pot (UR7e, END-EFFECTOR action space)

Task & data

Train / held-out split

Model architecture

Training setup

Why the two non-default flags are required

Results & the headline finding

The misleading eval_loss (the lesson)

Deployment note — research artifact, NOT wired to the robot

Usage / inference

Load the policy (LeRobot 0.6.1)

Produce actions

Reproduce the offline evaluation

Limitations & intended use

Links

Citation

The misleading `eval_loss` (the lesson)