Cleaner tests

Author: RunnersNum40

oryx/distribution/__init__.py

@@ -172,6 +172,14 @@ def __init__(
             tanh = bijectors.Tanh()
             self.distribution = distributions.Transformed(normal, tanh)
 
+    @property
+    def loc(self) -> Float[Array, " dims"]:
+        return self.distribution._distribution.loc
+
+    @property
+    def scale(self) -> Float[Array, " dims"]:
+        return self.distribution._distribution.scale
+
 
 class MultivariateNormalDiag(AbstractDistribution[Float[Array, " dims"]]):
 
@@ -245,6 +253,14 @@ def __init__(
             tanh = bijectors.Tanh()
             self.distribution = distributions.Transformed(mvn, tanh)
 
+    @property
+    def loc(self) -> Float[Array, " dims"]:
+        return self.distribution._distribution.loc
+
+    @property
+    def scale_diag(self) -> Float[Array, " dims"]:
+        return self.distribution._distribution.scale_diag
+
 
 __all__ = [
     "AbstractDistribution",

oryx/policy/actor_critic/custom.py

@@ -37,8 +37,6 @@ class CustomActorCriticPolicy[
           * Box: MLP(feature_size → action_dim) with Normal + squashing
     """
 
-    state_index: eqx.nn.StateIndex[None]
-
     feature_extractor: (
         AbstractModel[[ObsType], FeatureType]
         | AbstractStatefulModel[[ObsType], FeatureType]
@@ -54,6 +52,7 @@ class CustomActorCriticPolicy[
     log_std: Float[Array, " action_size"]
 
     env: AbstractEnvLike[ActType, ObsType]
+    state_index: eqx.nn.StateIndex[None] = eqx.nn.StateIndex(None)
 
     def __init__(
         self,
@@ -134,8 +133,6 @@ def __init__(
         else:
             self.action_model = action_model
 
-        self.state_index = eqx.nn.StateIndex(None)
-
     @property
     def action_space(self) -> AbstractSpace[ActType]:
         return self.env.action_space

oryx/utils.py

@@ -43,6 +43,8 @@ def _f(carry_arr, x):
         carry, y = f(carry, x)
         carry_arr, _static = eqx.partition(carry, eqx.is_array)
 
+        # Assert will be omitted from the compiled code
+        # I tried using `eqx.error(carry_arr, cond)` if but it breaks if the carry includes a key
         assert eqx.tree_equal(
             static, _static
         ), "Non-array carry of filter_scan must not change."

oryx/wrapper/transform_reward.py

@@ -6,7 +6,7 @@
 import equinox as eqx
 from jax import numpy as jnp
 from jax import random as jr
-from jaxtyping import Array, Bool, Float, Key
+from jaxtyping import Array, ArrayLike, Bool, Float, Key
 
 from oryx.env import AbstractEnvLike
 
@@ -93,8 +93,8 @@ class ClipReward[ActType, ObsType](
     def __init__(
         self,
         env: AbstractEnvLike[ActType, ObsType],
-        min: Float[Array, ""] = jnp.asarray(-1.0),
-        max: Float[Array, ""] = jnp.asarray(1.0),
+        min: Float[ArrayLike, ""] = jnp.asarray(-1.0),
+        max: Float[ArrayLike, ""] = jnp.asarray(1.0),
     ):
         self.env = env
         self.min = jnp.asarray(min)

tests/shared/__init__.py

@@ -1,9 +1 @@
-from .envs import DiscreteActionEnv, DummyEnv, EchoEnv, FiniteEpisodeEnv, PassThroughEnv
 
-__all__ = [
-    "DiscreteActionEnv",
-    "DummyEnv",
-    "EchoEnv",
-    "FiniteEpisodeEnv",
-    "PassThroughEnv",
-]

tests/shared/envs.py

@@ -14,9 +14,9 @@ class EchoEnv(AbstractEnv[Float[Array, " n"], Float[Array, " n"]]):
     Echoes the action back as the observation, never terminates.
     """
 
-    state_index: eqx.nn.StateIndex[None]
     action_space: Box
     observation_space: Box
+    state_index: eqx.nn.StateIndex[None] = eqx.nn.StateIndex(None)
 
     def __init__(
         self,
@@ -25,7 +25,6 @@ def __init__(
     ):
         self.action_space = action_space
         self.observation_space = observation_space
-        self.state_index = eqx.nn.StateIndex(None)
 
     def reset(self, state, *, key):
         return state, self.observation_space.sample(key), {}
@@ -92,10 +91,9 @@ class PassThroughEnv(AbstractEnv[Float[Array, ""], Float[Array, ""]]):
     always returns observation = 0.0.
     """
 
-    state_index: eqx.nn.StateIndex[None]
+    state_index: eqx.nn.StateIndex[None] = eqx.nn.StateIndex(None)
 
-    def __init__(self):
-        self.state_index = eqx.nn.StateIndex(None)
+    def __init__(self): ...  # Override dataclass init to avoid unused argument warning
 
     def reset(self, state, *, key):
         return state, jnp.asarray(0.0), {}
@@ -171,12 +169,11 @@ class DiscreteActionEnv(AbstractEnv[Int[Array, ""], Float[Array, " n"]]):
     Never terminates/truncates.
     """
 
-    state_index: eqx.nn.StateIndex[None]
+    state_index: eqx.nn.StateIndex[None] = eqx.nn.StateIndex(None)
     observation_space: Box
     action_space: Discrete
 
     def __init__(self, *, key, n_actions: int = 4, obs_size: int = 3):
-        self.state_index = eqx.nn.StateIndex(None)
         self.observation_space = Box(-jnp.inf, jnp.inf, shape=(obs_size,))
         self.action_space = Discrete(n_actions)
 

tests/shared/models.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import equinox as eqx
+from jax import numpy as jnp
+from jaxtyping import Array
+
+from oryx.model.base_model import AbstractModel, AbstractStatefulModel
+
+
+class Doubler(AbstractModel[[Array], Array]):
+    def __call__(self, x: Array) -> Array:
+        return 2.0 * x
+
+
+class StatefulDoubler(AbstractStatefulModel[[Array], Array]):
+    state_index: eqx.nn.StateIndex[Array]
+
+    def __init__(self):
+        self.state_index = eqx.nn.StateIndex(jnp.array(0.0))
+
+    def __call__(self, state, x):
+        c = state.get(self.state_index)
+        state = state.set(self.state_index, c + 1.0)
+        return state, 2.0 * x
+
+    def reset(self, state):
+        return state.set(self.state_index, jnp.array(0.0))

tests/test_buffers.py

@@ -10,7 +10,7 @@
 from oryx.policy import AbstractActorCriticPolicy
 from oryx.utils import filter_scan
 
-from .shared import EchoEnv
+from .shared.envs import EchoEnv
 
 
 class TestRolloutBufferBasics:
@@ -111,12 +111,11 @@ def test_batches_handles_non_divisible_batch_size(self):
 class TestRolloutBufferWithEnv:
 
     class _SimplePolicy(AbstractActorCriticPolicy):
-        state_index: eqx.nn.StateIndex[None]
         env: EchoEnv
+        state_index: eqx.nn.StateIndex[None] = eqx.nn.StateIndex(None)
 
         def __init__(self, env: EchoEnv):
             self.env = env
-            self.state_index = eqx.nn.StateIndex(None)
 
         @property
         def action_space(self):