From 0dc6c7265e37d5457be2477e253c42a6950eb62a Mon Sep 17 00:00:00 2001 From: Alexander Soare Date: Wed, 8 May 2024 14:47:28 +0100 Subject: [PATCH] use loggable data types in act return dict --- lerobot/common/policies/act/modeling_act.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lerobot/common/policies/act/modeling_act.py b/lerobot/common/policies/act/modeling_act.py index 5ff25fea..a795d87b 100644 --- a/lerobot/common/policies/act/modeling_act.py +++ b/lerobot/common/policies/act/modeling_act.py @@ -101,7 +101,7 @@ class ACTPolicy(nn.Module, PyTorchModelHubMixin): F.l1_loss(batch["action"], actions_hat, reduction="none") * ~batch["action_is_pad"].unsqueeze(-1) ).mean() - loss_dict = {"l1_loss": l1_loss} + loss_dict = {"l1_loss": l1_loss.item()} if self.config.use_vae: # Calculate Dā‚–ā‚—(latent_pdf || standard_normal). Note: After computing the KL-divergence for # each dimension independently, we sum over the latent dimension to get the total @@ -110,7 +110,7 @@ class ACTPolicy(nn.Module, PyTorchModelHubMixin): mean_kld = ( (-0.5 * (1 + log_sigma_x2_hat - mu_hat.pow(2) - (log_sigma_x2_hat).exp())).sum(-1).mean() ) - loss_dict["kld_loss"] = mean_kld + loss_dict["kld_loss"] = mean_kld.item() loss_dict["loss"] = l1_loss + mean_kld * self.config.kl_weight else: loss_dict["loss"] = l1_loss