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

one more incision before knocking out reward decoding

Browse Source
This commit is contained in:
lucidrains 2025-10-08 06:11:02 -07:00
parent c056835aea
commit f7bdaddbbb
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
dreamer4/dreamer4.py
View File

@ -215,8 +215,10 @@ class SymExpTwoHot(Module):
def bins_to_scalar_value( def bins_to_scalar_value(
self, self,
two_hot_encoding # (... l) logits, # (... l)
normalize = False
): ):
two_hot_encoding = logits.softmax(dim = -1) if normalize else logits
return einsum(two_hot_encoding, self.bin_values, '... l, l -> ...') return einsum(two_hot_encoding, self.bin_values, '... l, l -> ...')
def forward( def forward(