From 83276eeb2f2d63b07991ab2eb8b4e9e5b45f6274 Mon Sep 17 00:00:00 2001 From: Jade Choghari Date: Wed, 14 Jan 2026 14:53:18 +0000 Subject: [PATCH] loss naming --- src/lerobot/policies/pi05_full/modeling_pi05.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/lerobot/policies/pi05_full/modeling_pi05.py b/src/lerobot/policies/pi05_full/modeling_pi05.py index 0348e6550..de5ca1955 100644 --- a/src/lerobot/policies/pi05_full/modeling_pi05.py +++ b/src/lerobot/policies/pi05_full/modeling_pi05.py @@ -1009,7 +1009,7 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch` flow_loss = F.mse_loss(u_t, v_t, reduction="none") return { - "flow_loss": flow_loss.mean(), + "flow_mse_loss": flow_loss.mean(), "action_ce_loss": fast_loss.mean(), "subtask_ce_loss": subtask_loss, "loss": flow_loss.mean() + 0.1 * subtask_loss.mean() + 0.05 * fast_loss.mean(), # TODO: jadechoghari: check weights @@ -1501,7 +1501,7 @@ class PI05FullPolicy(PreTrainedPolicy): # Prepare detailed loss dictionary for logging detailed_loss_dict = { "loss": loss.item(), - "flow_loss": loss_dict["flow_loss"].mean().item(), + "flow_mse_loss": loss_dict["flow_loss"].mean().item(), "subtask_ce_loss": loss_dict["subtask_ce_loss"].item(), "action_ce_loss": loss_dict["action_ce_loss"].item(), }