wangzixiang
/
lerobot
mirror of https://github.com/huggingface/lerobot.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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.
This commit is contained in:
Yachen Kang 2025-04-10 15:29:32 +08:00 committed by GitHub
parent 5322417c03
commit 294948eff5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
lerobot/common/policies/pi0/modeling_pi0.py
View File

@ -325,7 +325,8 @@ class PI0Policy(PreTrainedPolicy):
loss_dict["losses_after_in_ep_bound"] = losses.clone() loss_dict["losses_after_in_ep_bound"] = losses.clone()
# Remove padding # 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() loss_dict["losses_after_rm_padding"] = losses.clone()
# For backward pass # For backward pass