Merge pull request #79 from andrearosasco/patch-1

Added LogSoftmax to discrete loss

octo/model/components/action_heads.py

@@ -140,8 +140,8 @@ def discrete_loss(
     labels = discrete_tokenizer(ground_truth_value)
     labels_one_hot = jax.nn.one_hot(labels, logits.shape[-1])
 
-    loss = -jnp.sum(logits * labels_one_hot, axis=-1)
-    loss = masked_mean(loss, mask)
+    loss = jnp.sum(jax.nn.log_softmax(logits, axis=-1) * labels_one_hot, axis=-1)
+    loss = -masked_mean(loss, mask)
 
     # compute accuracy between predicted actions and target actions
     pred_label = jnp.argmax(logits, axis=-1)