From efee611403e6d2a3b500900faf651d7937199282 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Thu, 12 Mar 2026 00:51:31 +0100 Subject: [PATCH] fix(policies): crop losses based on the action dof (#3133) Co-authored-by: Chenning Yu --- src/lerobot/policies/smolvla/modeling_smolvla.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/lerobot/policies/smolvla/modeling_smolvla.py b/src/lerobot/policies/smolvla/modeling_smolvla.py index 048d974af..32165eba8 100644 --- a/src/lerobot/policies/smolvla/modeling_smolvla.py +++ b/src/lerobot/policies/smolvla/modeling_smolvla.py @@ -377,6 +377,8 @@ class SmolVLAPolicy(PreTrainedPolicy): actions_is_pad = batch.get("action_is_pad") loss_dict = {} losses = self.model.forward(images, img_masks, lang_tokens, lang_masks, state, actions, noise, time) + original_action_dim = self.config.action_feature.shape[0] + losses = losses[:, :, :original_action_dim] loss_dict["losses_after_forward"] = losses.clone().mean().item() if actions_is_pad is not None: