cyy/Motrixlab
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: VAE loss /num_mini_batches, KL use raw sum (align with upstream)

Browse Source
This commit is contained in:
8x54zj-m
2026-06-30 15:38:26 +08:00
parent b0f1da4596
commit aa17ea853e
1 changed files with 5 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py
View File

@@ -192,17 +192,16 @@ class DreamWaQPPO(PPO):
mse = nn.functional.mse_loss mse = nn.functional.mse_loss
estimation_loss = mse(code_vel, vel_target) estimation_loss = mse(code_vel, vel_target)
reconstruction_loss = mse(decode, obs_target) reconstruction_loss = mse(decode, obs_target)
# KL 散度：-0.5 * sum(1 + logvar - mean^2 - exp(logvar)) # KL 散度：上游用 raw sum（非 mean），然后除以 num_mini_batches
# clamp logvar 防止 exp 溢出
logvar_latent = torch.clamp(logvar_latent, -20.0, 10.0) logvar_latent = torch.clamp(logvar_latent, -20.0, 10.0)
kl_loss = -0.5 * torch.sum( kl_loss = -0.5 * torch.sum(
1 + logvar_latent - mean_latent.pow(2) - logvar_latent.exp(), dim=-1 1 + logvar_latent - mean_latent.pow(2) - logvar_latent.exp()
).mean() )
# 与上游完全一致：除以 num_mini_batches
autoenc_loss = ( autoenc_loss = (
estimation_loss + reconstruction_loss + self.vae_beta * kl_loss estimation_loss + reconstruction_loss + self.vae_beta * kl_loss
) ) / self.num_mini_batches
# 防止 NaN 传播
if torch.isnan(autoenc_loss) or torch.isinf(autoenc_loss): if torch.isnan(autoenc_loss) or torch.isinf(autoenc_loss):
return torch.tensor(0.0, device=obs_batch.device) return torch.tensor(0.0, device=obs_batch.device)
return autoenc_loss return autoenc_loss