diff --git a/dreamer4/dreamer4.py b/dreamer4/dreamer4.py
index 2f1c20f..c174ad2 100644
--- a/dreamer4/dreamer4.py
+++ b/dreamer4/dreamer4.py
@@ -215,8 +215,10 @@ class SymExpTwoHot(Module):
 
     def bins_to_scalar_value(
         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 -> ...')
 
     def forward(