hangX/lerobot
1
0
Fork 0
forked from tangger/lerobot
Code Pull Requests Activity

Enable logging all the information returned by the forward methods of policies (#151)

Browse Source
This commit is contained in:
Alexander Soare
2024-05-10 07:45:32 +01:00
committed by GitHub
parent b187942db4
commit 1249aee3ac
5 changed files with 12 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
lerobot/common/policies/act/modeling_act.py
View File

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