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

@@ -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