Fix

ygoai/rl/jax/utils.py

@@ -9,12 +9,12 @@ def masked_mean(x, valid):
     return x.sum() / valid.sum()
 
 
-def masked_normalize(x, valid, epsilon=1e-8):
+def masked_normalize(x, valid, eps=1e-8):
     x = jnp.where(valid, x, jnp.zeros_like(x))
     n = valid.sum()
     mean = x.sum() / n
     variance = jnp.square(x - mean).sum() / n
-    return (x - mean) / jnp.sqrt(variance + epsilon)
+    return (x - mean) / jnp.sqrt(variance + eps)
 
 
 def categorical_sample(logits, key):