This commit is contained in:
lucidrains 2025-10-27 08:34:22 -07:00
parent 349a03acd7
commit f808b1c1d2
2 changed files with 5 additions and 4 deletions

View File

@ -2244,10 +2244,11 @@ class DynamicsWorldModel(Module):
if not exists(experience.is_truncated):
experience.is_truncated = full((batch,), True, device = latents.device)
mask_for_gae = lens_to_mask(experience.lens, time)
if exists(experience.lens):
mask_for_gae = lens_to_mask(experience.lens, time)
rewards = rewards.masked_fill(mask_for_gae, 0.)
old_values = old_values.masked_fill(mask_for_gae, 0.)
rewards = rewards.masked_fill(mask_for_gae, 0.)
old_values = old_values.masked_fill(mask_for_gae, 0.)
# calculate returns

View File

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