Tighten test style

Author: RunnersNum40

tests/algorithms/test_on_policy.py

@@ -15,8 +15,6 @@
 
 
 class OnPolicyAlgorithm[ActType, ObsType](AbstractOnPolicyAlgorithm[ActType, ObsType]):
-    """Minimal concrete on-policy algorithm just for testing."""
-
     env: AbstractEnvLike[ActType, ObsType]
     policy: AbstractActorCriticPolicy[Float, ActType, ObsType]
     state_index: eqx.nn.StateIndex[None] = eqx.nn.StateIndex(None)

tests/algorithms/test_ppo.py

@@ -2,13 +2,16 @@
 
 import equinox as eqx
 import jax
+import pytest
 from jax import numpy as jnp
 from jax import random as jr
 
 from oryx.algorithm import PPO
 from oryx.policy.actor_critic import CustomActorCriticPolicy
 from oryx.space import Box
+from tests.buffers import empty_buffer
 from tests.envs import EchoEnv
+from tests.policies import ConstantPolicy
 
 
 def _make_ppo(num_steps: int = 8) -> tuple[PPO, eqx.nn.State]:
@@ -86,3 +89,95 @@ def test_learn_integration(self):
         pol_state = state.substate(policy)
         pol_state, act = policy.predict(pol_state, obs, key=jr.key(5))
         assert env.action_space.contains(act)
+
+    def test_ppo_loss_clip_vs_unclipped_and_entropy(self):
+        buf = empty_buffer(n=6)
+        pol = ConstantPolicy(new_value=5.0, logp=0.0, entropy_val=2.0)
+
+        loss_c, stats_c = PPO.ppo_loss(
+            pol,
+            buf,
+            normalize_advantages=False,
+            clip_coefficient=0.1,
+            clip_value_loss=True,
+            value_loss_coefficient=0.5,
+            state_magnitude_coefficient=0.0,
+            entropy_loss_coefficient=0.01,
+        )
+        policy_loss = -1.0
+        value_loss_clipped = 0.5 * (0.1**2)
+        entropy = 2.0
+        expected_total = policy_loss + 0.5 * value_loss_clipped - 0.01 * entropy
+
+        assert float(stats_c.policy_loss) == pytest.approx(policy_loss)
+        assert float(stats_c.value_loss) == pytest.approx(value_loss_clipped)
+        assert float(stats_c.entropy_loss) == pytest.approx(entropy)
+        assert float(stats_c.state_magnitude_loss) == pytest.approx(0.0)
+        assert float(stats_c.total_loss) == pytest.approx(expected_total, rel=1e-6)
+        assert float(stats_c.approx_kl) == pytest.approx(0.0)
+
+        loss_u, stats_u = PPO.ppo_loss(
+            pol,
+            buf,
+            normalize_advantages=False,
+            clip_coefficient=0.1,
+            clip_value_loss=False,
+            value_loss_coefficient=0.5,
+            state_magnitude_coefficient=0.0,
+            entropy_loss_coefficient=0.01,
+        )
+        assert float(stats_u.value_loss) == pytest.approx(12.5)
+        expected_total_unclipped = policy_loss + 0.5 * 12.5 - 0.01 * entropy
+        assert float(stats_u.total_loss) == pytest.approx(
+            expected_total_unclipped, rel=1e-6
+        )
+        assert expected_total_unclipped != pytest.approx(expected_total)
+
+
+def _setup_algo(*, anneal: bool):
+    key = jr.key(0)
+    env = EchoEnv(space=Box(-jnp.ones(2), jnp.ones(2)))
+    policy, _ = eqx.nn.make_with_state(CustomActorCriticPolicy)(env=env, key=key)
+    algo, state = eqx.nn.make_with_state(PPO)(
+        env=env,
+        policy=policy,
+        num_steps=8,
+        num_epochs=2,
+        num_mini_batches=2,
+        learning_rate=1e-3,
+        anneal_learning_rate=anneal,
+    )
+    return algo, state
+
+
+def _collect_once(algo, state, *, key):
+    state, carry = algo.initialize_iteration_carry(state, key=key)
+    state, _, buf, _ = algo.collect_rollout(
+        algo.policy, state, carry.step_carry, key=jr.split(key)[0]
+    )
+    return state, buf
+
+
+def test_learning_rate_anneals_down():
+    algo, state = _setup_algo(anneal=True)
+    state, buf = _collect_once(algo, state, key=jr.key(1))
+
+    lr0 = float(algo.learning_rate(state))
+    state, _, _ = algo.train(state, algo.policy, buf, key=jr.key(2))
+    lr1 = float(algo.learning_rate(state))
+    state, _, _ = algo.train(state, algo.policy, buf, key=jr.key(3))
+    lr2 = float(algo.learning_rate(state))
+
+    assert lr0 > lr1 > lr2
+
+
+def test_learning_rate_constant_when_no_anneal():
+    algo, state = _setup_algo(anneal=False)
+    state, buf = _collect_once(algo, state, key=jr.key(4))
+
+    lr0 = float(algo.learning_rate(state))
+    state, _, _ = algo.train(state, algo.policy, buf, key=jr.key(5))
+    lr1 = float(algo.learning_rate(state))
+
+    assert lr0 == pytest.approx(1e-3)
+    assert lr1 == pytest.approx(lr0)

tests/buffers/__init__.py

@@ -0,0 +1,29 @@
+from jax import numpy as jnp
+
+from oryx.buffer import RolloutBuffer
+
+
+def empty_buffer(n: int = 8):
+    observations = jnp.zeros((n,))
+    actions = jnp.zeros((n,))
+    rewards = jnp.zeros((n,))
+    terminations = jnp.full((n,), False)
+    truncations = jnp.full((n,), False)
+    log_probs = jnp.zeros((n,))
+    values = jnp.zeros((n,))
+    states = jnp.zeros((n,))
+    returns = jnp.zeros((n,))
+    advantages = jnp.ones((n,))
+
+    return RolloutBuffer(
+        observations=observations,
+        actions=actions,
+        rewards=rewards,
+        terminations=terminations,
+        truncations=truncations,
+        log_probs=log_probs,
+        values=values,
+        states=states,
+        returns=returns,
+        advantages=advantages,
+    )

tests/buffers/test_rollout.py

@@ -6,10 +6,9 @@
 from jax import random as jr
 
 from oryx.buffer import RolloutBuffer
-from oryx.distribution import Normal
-from oryx.policy import AbstractActorCriticPolicy
 from oryx.utils import filter_scan
 from tests.envs import EchoEnv
+from tests.policies import NormalPolicy
 
 
 class TestRolloutBuffer:
@@ -108,43 +107,13 @@ def test_batches_handles_non_divisible_batch_size(self):
 
 
 class TestRolloutBufferWithEnv:
-
-    class _SimplePolicy(AbstractActorCriticPolicy):
-        env: EchoEnv
-        state_index: eqx.nn.StateIndex[None] = eqx.nn.StateIndex(None)
-
-        def __init__(self, env: EchoEnv):
-            self.env = env
-
-        @property
-        def action_space(self):
-            return self.env.action_space
-
-        @property
-        def observation_space(self):
-            return self.env.observation_space
-
-        def extract_features(self, state, observation):
-            return state, jnp.asarray(observation)
-
-        def action_dist_from_features(self, state, features):
-            return state, Normal(features, jnp.asarray(1.0))
-
-        def value_from_features(self, state, features):
-            return state, jnp.asarray(0.0)
-
-        def reset(self, state):
-            return state
-
     def test_env_policy_collection_and_batching(
         self, num_steps: int = 8, batch_size: int = 2
     ):
         reset_key, scan_key, batch_key = jr.split(jr.key(0), 3)
 
         env = EchoEnv()
-        policy, state = eqx.nn.make_with_state(TestRolloutBufferWithEnv._SimplePolicy)(
-            env=env
-        )
+        policy, state = eqx.nn.make_with_state(NormalPolicy)(env=env)
 
         state, obs, _ = env.reset(state, key=reset_key)
 

tests/envs/test_base_env.py

@@ -6,18 +6,16 @@
 from jax import random as jr
 
 from oryx.env import AbstractEnv, AbstractEnvLike
+from oryx.wrapper import Identity
 from tests.envs import DummyEnv
 
 
 class TestAbstractEnvLike:
     def test_cannot_instantiate_abstract(self):
-        """Direct instantiation of abstract bases must fail."""
         with pytest.raises(TypeError):
             AbstractEnvLike()  # pyright: ignore
 
     def test_missing_methods(self):
-        """A subclass that omits a required method should be abstract."""
-
         class NoStepEnv(AbstractEnvLike):
             def reset(self, state, *, key):
                 raise NotImplementedError
@@ -39,13 +37,10 @@ def close(self): ...
 
 class TestAbstractEnv:
     def test_cannot_instantiate_abstract(self):
-        """Direct instantiation of abstract bases must fail."""
         with pytest.raises(TypeError):
             AbstractEnv()  # pyright: ignore
 
     def test_missing_methods(self):
-        """A subclass that omits a required method should be abstract."""
-
         class NoStepEnv(AbstractEnv):
             def reset(self, state, *, key):
                 raise NotImplementedError
@@ -65,7 +60,6 @@ def close(self): ...
             NoStepEnv()  # pyright: ignore
 
     def test_dummy_env_reset_and_step_shapes(self):
-        """`reset` and `step` should respect the declared spaces & signatures."""
         key = jr.key(0)
         env, state = eqx.nn.make_with_state(DummyEnv)(key=key)
 
@@ -86,30 +80,7 @@ def test_dummy_env_reset_and_step_shapes(self):
         assert info == {}
 
     def test_unwrapped_returns_base_env(self):
-        """Even after extra wrapping, `.unwrapped` should reach the original env."""
-        from oryx.wrapper.base_wrapper import AbstractNoRenderOrCloseWrapper
-
-        class IdentityWrapper(AbstractNoRenderOrCloseWrapper):
-            env: DummyEnv
-
-            def __init__(self, env):
-                self.env = env
-
-            def reset(self, state, *, key):
-                return self.env.reset(state, key=key)
-
-            def step(self, state, action, *, key):
-                return self.env.step(state, action, key=key)
-
-            @property
-            def action_space(self):
-                return self.env.action_space
-
-            @property
-            def observation_space(self):
-                return self.env.observation_space
-
         base_env, _ = eqx.nn.make_with_state(DummyEnv)(key=jr.key(0))
-        wrapped_env = IdentityWrapper(base_env)
+        wrapped_env = Identity(base_env)
 
-        assert wrapped_env.unwrapped is base_env, "unwrapped must expose base env"
+        assert wrapped_env.unwrapped is base_env

tests/policies/__init__.py

@@ -1 +1,70 @@
+import equinox as eqx
+from jax import numpy as jnp
 
+from oryx.distribution import Normal
+from oryx.policy.actor_critic import AbstractActorCriticPolicy
+from oryx.space import Box
+from tests.envs import EchoEnv
+
+
+class ConstantPolicy(AbstractActorCriticPolicy):
+    new_value: float
+    logp: float
+    entropy_val: float
+
+    state_index: eqx.nn.StateIndex[None] = eqx.nn.StateIndex(None)
+    action_space = Box(-jnp.inf, jnp.inf, shape=())
+    observation_space = Box(-jnp.inf, jnp.inf, shape=())
+
+    def __init__(self, new_value: float, logp: float = 0.0, entropy_val: float = 0.0):
+        self.new_value = float(new_value)
+        self.logp = float(logp)
+        self.entropy_val = float(entropy_val)
+
+    def extract_features(self, state, observation):
+        return state, jnp.asarray(0.0)
+
+    def action_dist_from_features(self, state, features):
+        raise NotImplementedError
+
+    def value_from_features(self, state, features):
+        return state, jnp.asarray(self.new_value)
+
+    def reset(self, state):
+        return state
+
+    def evaluate_action(self, state, observation, action):
+        return (
+            state,
+            jnp.asarray(self.new_value),
+            jnp.asarray(self.logp),
+            jnp.asarray(self.entropy_val),
+        )
+
+
+class NormalPolicy(AbstractActorCriticPolicy):
+    env: EchoEnv
+    state_index: eqx.nn.StateIndex[None] = eqx.nn.StateIndex(None)
+
+    def __init__(self, env: EchoEnv):
+        self.env = env
+
+    @property
+    def action_space(self):
+        return self.env.action_space
+
+    @property
+    def observation_space(self):
+        return self.env.observation_space
+
+    def extract_features(self, state, observation):
+        return state, jnp.asarray(observation)
+
+    def action_dist_from_features(self, state, features):
+        return state, Normal(features, jnp.asarray(1.0))
+
+    def value_from_features(self, state, features):
+        return state, jnp.asarray(0.0)
+
+    def reset(self, state):
+        return state

tests/test_utils.py

@@ -82,9 +82,6 @@ def _compute(x):
 
 @pytest.mark.parametrize("thread_flag", [False, True])
 def test_debug_wrapper_threading(thread_flag):
-    """
-    Ensure the callback fires and – when requested – in a different thread.
-    """
     wrapped = debug_wrapper(_compute, thread=thread_flag)
 
     @jax.jit
@@ -97,12 +94,12 @@ def f(x):
     time.sleep(0.05)
 
     assert float(result) == pytest.approx(4.14)
-    assert hasattr(_compute, "thread_id"), "callback never executed"
+    assert hasattr(_compute, "thread_id")
 
     if thread_flag:
-        assert _compute.thread_id != main_thread, "expected a different thread"
+        assert _compute.thread_id != main_thread
     else:
-        assert _compute.thread_id == main_thread, "should run on main thread"
+        assert _compute.thread_id == main_thread
 
 
 def _array_collector(x):