Fix action padding key at SmolVLA (#1717)

Issue https://github.com/huggingface/lerobot/issues/1707

Action padding mask is set at LeRobotDataset as f"{key}_is_pad".

Wrong key doesn't raise any errors, however, padding mask is ignored,
resulting wrong attention at around the edges of an episode
when multi step actions is enabled (aka. action horizon is greater
than 1).

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
This commit is contained in:
H.Yamada
2026-03-12 04:12:21 +09:00
committed by GitHub
parent 19c6adef85
commit f311ca3dce
@@ -374,7 +374,7 @@ class SmolVLAPolicy(PreTrainedPolicy):
lang_tokens = batch[f"{OBS_LANGUAGE_TOKENS}"]
lang_masks = batch[f"{OBS_LANGUAGE_ATTENTION_MASK}"]
actions = self.prepare_action(batch)
actions_is_pad = batch.get("actions_id_pad")
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)
loss_dict["losses_after_forward"] = losses.clone().mean().item()