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

will organize the unembedding parameters under the actor optimizer

Browse Source
This commit is contained in:
lucidrains 2025-10-11 06:55:57 -07:00
parent 563b269f8a
commit 02558d1f08
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
dreamer4/dreamer4.py
View File

@ -344,6 +344,12 @@ class ActionEmbedder(Module):
self.continuous_action_unembed = Parameter(torch.randn(num_continuous_actions, unembed_dim, 2) * 1e-2) self.continuous_action_unembed = Parameter(torch.randn(num_continuous_actions, unembed_dim, 2) * 1e-2)
def embed_parameters(self):
return set([*self.discrete_action_embed.parameters(), *self.continuous_action_embed.parameters()])
def unembed_parameters(self):
return set([*self.discrete_action_unembed.parameters(), *self.continuous_action_unembed.parameters()])
@property @property
def device(self): def device(self):
return self.discrete_action_offsets.device return self.discrete_action_offsets.device