Fix weak typing

Author: RunnersNum40

oryx/policy/actor_critic/mlp.py

@@ -53,14 +53,14 @@ def __init__(
     ):
         if isinstance(env.action_space, Discrete):
             act_size = int(env.action_space.n)
-            self.log_std = jnp.array([])
+            self.log_std = jnp.array([], dtype=float)
         elif isinstance(env.action_space, Box):
             if env.action_space.shape:
                 act_size = int(jnp.prod(jnp.asarray(env.action_space.shape)))
-                self.log_std = jnp.full((act_size,), log_std_init)
+                self.log_std = jnp.full((act_size,), log_std_init, dtype=float)
             else:
                 act_size = "scalar"
-                self.log_std = jnp.array(log_std_init)
+                self.log_std = jnp.array(log_std_init, dtype=float)
         else:
             raise NotImplementedError(
                 f"Action space {type(env.action_space)} not supported."

oryx/space/base_space.py

@@ -69,8 +69,8 @@ class Discrete(AbstractSpace[Int[Array, ""]]):
     def __init__(self, n: Int[ArrayLike, ""], start: Int[ArrayLike, ""] = 0):
         assert n > 0, "n must be positive"  # pyright: ignore
 
-        self._n = jnp.asarray(n)
-        self.start = jnp.asarray(start)
+        self._n = jnp.array(n, dtype=float)
+        self.start = jnp.array(start, dtype=float)
 
     @property
     def n(self) -> Int[Array, ""]:
@@ -133,8 +133,8 @@ def __init__(
         high: Float[ArrayLike, " ..."],
         shape: tuple[int, ...] | None = None,
     ):
-        low = jnp.asarray(low)
-        high = jnp.asarray(high)
+        low = jnp.asarray(low, dtype=float)
+        high = jnp.asarray(high, dtype=float)
         if shape is None:
             low, high = jnp.broadcast_arrays(low, high)
             shape = low.shape
@@ -361,8 +361,8 @@ def __init__(self, ns: tuple[int, ...], starts: tuple[int, ...] = (0,)):
         assert len(ns) == len(starts), "ns and starts must have the same length"
         assert all(n > 0 for n in ns), "all n must be positive"
 
-        self.ns = jnp.asarray(ns)
-        self.starts = jnp.asarray(starts)
+        self.ns = jnp.array(ns, dtype=float)
+        self.starts = jnp.array(starts, dtype=float)
 
     @property
     def shape(self) -> tuple[int, ...]: