From 294948eff57d3730b1b38014092d719fb1aeb124 Mon Sep 17 00:00:00 2001 From: Yachen Kang Date: Thu, 10 Apr 2025 15:29:32 +0800 Subject: [PATCH] fix losses unpadding bug Use original_action_dim instead of self.config.max_action_dim when processing losses to exclude padded regions. Including padding areas in loss computation would impair the model's action reconstruction capability. --- lerobot/common/policies/pi0/modeling_pi0.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lerobot/common/policies/pi0/modeling_pi0.py b/lerobot/common/policies/pi0/modeling_pi0.py index 7599fa63..43b11456 100644 --- a/lerobot/common/policies/pi0/modeling_pi0.py +++ b/lerobot/common/policies/pi0/modeling_pi0.py @@ -325,7 +325,8 @@ class PI0Policy(PreTrainedPolicy): loss_dict["losses_after_in_ep_bound"] = losses.clone() # Remove padding - losses = losses[:, :, : self.config.max_action_dim] + original_action_dim = self.config.action_feature.shape[0] + losses = losses[:, :, :original_action_dim] loss_dict["losses_after_rm_padding"] = losses.clone() # For backward pass