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

fix storing of agent embedding

Browse Source
This commit is contained in:
lucidrains 2025-10-28 09:36:58 -07:00
parent 903c43b770
commit 46f86cd247
2 changed files with 6 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
dreamer4/dreamer4.py
View File

@ -2255,7 +2255,7 @@ class DynamicsWorldModel(Module):
video = cat((video, next_frame), dim = 2)
rewards = safe_cat((rewards, reward), dim = 1)
acc_agent_embed = safe_cat((acc_agent_embed, agent_embed), dim = 1)
acc_agent_embed = safe_cat((acc_agent_embed, one_agent_embed), dim = 1)
# package up one experience for learning
@ -2397,7 +2397,7 @@ class DynamicsWorldModel(Module):
return_intermediates = True
)
agent_embeds = agent_embeds[..., agent_index, :]
agent_embeds = agent_embeds[..., agent_index, :]
# maybe detach agent embed
@ -2672,7 +2672,9 @@ class DynamicsWorldModel(Module):
# maybe store agent embed
acc_agent_embed = safe_cat((acc_agent_embed, agent_embed), dim = 1)
if store_agent_embed:
one_agent_embed = agent_embed[:, -1:, agent_index]
acc_agent_embed = safe_cat((acc_agent_embed, one_agent_embed), dim = 1)
# decode the agent actions if needed

2
pyproject.toml
View File

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