From d0b58190dab254ea18f814c03129f102c8b68a66 Mon Sep 17 00:00:00 2001 From: Cotton Hu <1821141394@qq.com> Date: Thu, 26 Feb 2026 00:36:31 +0800 Subject: [PATCH] fix(policies): support dp train when n_obs_steps=1 (#2430) Co-authored-by: hukongtao Co-authored-by: Steven Palma --- src/lerobot/policies/diffusion/modeling_diffusion.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/lerobot/policies/diffusion/modeling_diffusion.py b/src/lerobot/policies/diffusion/modeling_diffusion.py index 7525c9252..314ca369c 100644 --- a/src/lerobot/policies/diffusion/modeling_diffusion.py +++ b/src/lerobot/policies/diffusion/modeling_diffusion.py @@ -142,6 +142,9 @@ class DiffusionPolicy(PreTrainedPolicy): """Run the batch through the model and compute the loss for training or validation.""" if self.config.image_features: batch = dict(batch) # shallow copy so that adding a key doesn't modify the original + for key in self.config.image_features: + if self.config.n_obs_steps == 1 and batch[key].ndim == 4: + batch[key] = batch[key].unsqueeze(1) batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4) loss = self.diffusion.compute_loss(batch) # no output_dict so returning None