From 3b86050ab0fcb00cf46c091ee21ab395aad8392f Mon Sep 17 00:00:00 2001 From: Radek Osmulski Date: Mon, 27 May 2024 18:06:26 +1000 Subject: [PATCH] throw an error if config.do_maks_loss and action_is_pad not provided in batch (#213) Co-authored-by: Alexander Soare --- lerobot/common/policies/diffusion/modeling_diffusion.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/lerobot/common/policies/diffusion/modeling_diffusion.py b/lerobot/common/policies/diffusion/modeling_diffusion.py index 2ae03f221..2bf45bb6b 100644 --- a/lerobot/common/policies/diffusion/modeling_diffusion.py +++ b/lerobot/common/policies/diffusion/modeling_diffusion.py @@ -304,7 +304,11 @@ class DiffusionModel(nn.Module): loss = F.mse_loss(pred, target, reduction="none") # Mask loss wherever the action is padded with copies (edges of the dataset trajectory). - if self.config.do_mask_loss_for_padding and "action_is_pad" in batch: + if self.config.do_mask_loss_for_padding: + if "action_is_pad" not in batch: + raise ValueError( + f"You need to provide 'action_is_pad' in the batch when {self.config.do_mask_loss_for_padding=}." + ) in_episode_bound = ~batch["action_is_pad"] loss = loss * in_episode_bound.unsqueeze(-1)