hongshaorou/dreamer4
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

oops

Browse Source
This commit is contained in:
lucidrains 2025-10-18 10:31:48 -07:00
parent b6aa19f31e
commit 79a1b1c46e
2 changed files with 5 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
dreamer4/dreamer4.py
View File

@ -530,10 +530,8 @@ class ActionEmbedder(Module):
if exists(continuous_action_types):
continuous_action_unembed = continuous_action_unembed[continuous_action_types]
if isinstance(pred_head_index, int):
continuous_action_unembed = continuous_action_unembed[:, pred_head_index]
else:
continuous_action_unembed = continuous_action_unembed.index_select(1, pred_head_index)
if exists(pred_head_index):
continuous_action_unembed = continuous_action_unembed.index_select(1, pred_head_index)
continuous_action_mean_log_var = einsum(embeds, continuous_action_unembed, '... d, na mtp d two -> mtp ... na two')
@ -2503,12 +2501,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()

2
pyproject.toml
View File

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