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

fix masking for multiple agent tokens

Browse Source
This commit is contained in:
lucidrains 2025-10-08 08:26:44 -07:00
parent ed0918c974
commit 4c2ed100a3
2 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
dreamer4/dreamer4.py
View File

@ -1548,7 +1548,7 @@ class DynamicsWorldModel(Module):
reward_embeds = pad_at_dim(reward_embeds, (1, -pop_last_reward), dim = -2, value = 0.) # shift as each agent token predicts the next reward reward_embeds = pad_at_dim(reward_embeds, (1, -pop_last_reward), dim = -2, value = 0.) # shift as each agent token predicts the next reward
agent_tokens = einx.add('b t ... d, b t', agent_tokens, reward_embeds) agent_tokens = einx.add('b t ... d, b t d', agent_tokens, reward_embeds)
# main function, needs to be defined as such for shortcut training - additional calls for consistency loss # main function, needs to be defined as such for shortcut training - additional calls for consistency loss
@ -1589,8 +1589,8 @@ class DynamicsWorldModel(Module):
attend_kwargs = dict(use_flex = use_flex, softclamp_value = self.attn_softclamp_value, device = device) attend_kwargs = dict(use_flex = use_flex, softclamp_value = self.attn_softclamp_value, device = device)
space_seq_len = ( space_seq_len = (
1 # action / agent token + 1 # signal + step
+ 1 # signal + step + self.num_agents # action / agent tokens
+ self.num_register_tokens + self.num_register_tokens
+ num_spatial_tokens + num_spatial_tokens
) )

2
pyproject.toml
View File

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