Use c*2 for version 2

ygoai/rl/jax/agent.py

@@ -273,7 +273,7 @@ class Encoder(nn.Module):
         if self.version == 2:
             x_global = jax.nn.leaky_relu(x_global, negative_slope=0.1)
             x_global = fc_layer(c, dtype=jnp.float32)(x_global)
-            f_global = x_global + GLUMlp(c, dtype=self.dtype, param_dtype=self.param_dtype)(
+            f_global = x_global + GLUMlp(c * 2, dtype=self.dtype, param_dtype=self.param_dtype)(
                 layer_norm(dtype=self.dtype)(x_global))
         else:
             f_global = x_global + MLP((c * 2, c * 2), dtype=self.dtype, param_dtype=self.param_dtype)(x_global)