This commit is contained in:
lucidrains 2025-10-18 10:28:56 -07:00
parent b6aa19f31e
commit b07870333c
2 changed files with 3 additions and 3 deletions

View File

@ -2503,12 +2503,12 @@ class DynamicsWorldModel(Module):
)
if exists(discrete_log_probs):
discrete_log_probs = discrete_log_probs.masked_fill(discrete_mask[..., None], 0.)
discrete_log_probs = discrete_log_probs.masked_fill(~discrete_mask[..., None], 0.)
behavior_clone_loss = behavior_clone_loss - reduce(discrete_log_probs, 'mtp b t na -> b t', 'sum').mean()
if exists(continuous_log_probs):
continuous_log_probs = continuous_log_probs.masked_fill(continuous_mask[..., None], 0.)
continuous_log_probs = continuous_log_probs.masked_fill(~continuous_mask[..., None], 0.)
behavior_clone_loss = behavior_clone_loss - reduce(continuous_log_probs, 'mtp b t na -> b t', 'sum').mean()

View File

@ -1,6 +1,6 @@
[project]
name = "dreamer4"
version = "0.0.38"
version = "0.0.39"
description = "Dreamer 4"
authors = [
{ name = "Phil Wang", email = "lucidrains@gmail.com" }