safenax 0.4.3__tar.gz → 0.4.4__tar.gz

{safenax-0.4.3 → safenax-0.4.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: safenax
-Version: 0.4.3
+Version: 0.4.4
 Summary: Constrained environments with a gymnax interface
 Project-URL: Homepage, https://github.com/0xprofessooor/safenax
 Project-URL: Repository, https://github.com/0xprofessooor/safenax

{safenax-0.4.3 → safenax-0.4.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "safenax"
-version = "0.4.3"
+version = "0.4.4"
 description = "Constrained environments with a gymnax interface"
 readme = "README.md"
 requires-python = ">=3.11"

{safenax-0.4.3 → safenax-0.4.4}/safenax/eco_ant/eco_ant_v1.py

@@ -27,7 +27,8 @@ class EcoAntV1(Ant):
 
     def step(self, state: State, action: jax.Array) -> State:
         # 1. RETRIEVE BATTERY FROM CURRENT OBSERVATION
-        current_battery = state.info["battery"]
+        current_battery_pct = state.obs[-1]
+        current_battery = current_battery_pct * self.battery_limit
 
         # 2. HANDLE STOCHASTICITY
         _, noise_key = jax.random.split(state.info["rng"])
@@ -47,8 +48,7 @@ class EcoAntV1(Ant):
         # 4. PHYSICS STEP
         next_state = super().step(state, noisy_action)
 
-        # 5. AUGMENT STATE
-        # Termination: OR with existing done condition
+        # 5. Termination: OR with existing done condition
         new_done = jnp.max(jnp.array([next_state.done, is_empty]))
 
         # Cost Signal: 1.0 if battery died this step
@@ -61,7 +61,7 @@ class EcoAntV1(Ant):
             **next_state.info,
             "rng": noise_key,
             "cost": cost,
-            "battery": jnp.where(new_done, jnp.array(self.battery_limit), new_battery),
+            "battery": new_battery,
         }
 
         return next_state.replace(obs=new_obs, done=new_done, info=new_info)

{safenax-0.4.3 → safenax-0.4.4}/safenax/eco_ant/eco_ant_v2.py

@@ -27,7 +27,8 @@ class EcoAntV2(Ant):
 
     def step(self, state: State, action: jax.Array) -> State:
         # 1. RETRIEVE BATTERY FROM CURRENT OBSERVATION
-        current_battery = state.info["battery"]
+        current_battery_pct = state.obs[-1]
+        current_battery = current_battery_pct * self.battery_limit
 
         # 2. HANDLE STOCHASTICITY
         _, noise_key = jax.random.split(state.info["rng"])
@@ -42,6 +43,7 @@ class EcoAntV2(Ant):
         # Check constraints
         is_empty = new_battery <= 0.0
         new_battery = jnp.maximum(new_battery, 0.0)
+        new_battery_pct = new_battery / self.battery_limit
 
         # 4. PHYSICS STEP
         next_state = super().step(state, noisy_action)
@@ -49,18 +51,24 @@ class EcoAntV2(Ant):
         # 5. Termination: OR with existing done condition
         new_done = jnp.max(jnp.array([next_state.done, is_empty]))
 
+        # Observation: Append the new battery level to the observation vector
+        new_obs = jnp.concatenate([next_state.obs, jnp.array([new_battery_pct])])
+
         new_info = {
             **next_state.info,
             "rng": noise_key,
             "cost": energy_used,
-            "battery": jnp.where(new_done, jnp.array(self.battery_limit), new_battery),
+            "battery": new_battery,
         }
 
-        return next_state.replace(done=new_done, info=new_info)
+        return next_state.replace(obs=new_obs, done=new_done, info=new_info)
 
     def reset(self, rng: jax.Array) -> State:
         state = super().reset(rng)
 
+        # Append initial battery to observation
+        new_obs = jnp.concatenate([state.obs, jnp.array([1.0])])
+
         # Initialize info
         new_info = {
             **state.info,
@@ -69,7 +77,7 @@ class EcoAntV2(Ant):
             "battery": jnp.array(self.battery_limit),
         }
 
-        return state.replace(info=new_info)
+        return state.replace(obs=new_obs, info=new_info)
 
 
 envs.register_environment(EcoAntV2.name, EcoAntV2)

{safenax-0.4.3 → safenax-0.4.4}/tests/test_eco_ant_v1.py

@@ -87,6 +87,7 @@ def test_termination_and_cost(env: EcoAntV1, key: jax.Array):
     # 1. Manually set battery to near-death in INFO (not Obs)
     # The environment logic reads current_battery = state.info["battery"]
     # So to test termination, we must hack the info dict.
+    env.battery_limit = 0.1  # Temporarily set limit low for test
     new_info = state.info.copy()
     new_info["battery"] = jnp.array(0.1)
 
@@ -102,7 +103,7 @@ def test_termination_and_cost(env: EcoAntV1, key: jax.Array):
     assert next_state.obs[-1] <= 0.0
 
     # Check Raw Battery Floor
-    assert next_state.info["battery"] == 10.0
+    assert next_state.info["battery"] == 0.0
 
     # Check Cost Signal
     assert next_state.info["cost"] == 1.0
@@ -143,7 +144,7 @@ def test_wrapper_step_api(wrapped_env: BraxToGymnaxWrapper, key: jax.Array):
     assert info["battery"] < 10.0
 
 
-def test_wrapper_autoreset_logic(wrapped_env: BraxToGymnaxWrapper, key: jax.Array):
+def test_wrapper_autoreset_logic(key: jax.Array):
     """
     CRITICAL TEST: Verifies AutoResetWrapper behavior on battery death.
 
@@ -154,16 +155,10 @@ def test_wrapper_autoreset_logic(wrapped_env: BraxToGymnaxWrapper, key: jax.Arra
        - Returns `done=True` (signaling the end of the dying episode)
        - Returns `obs` from the NEW reset state (Battery=1.0, not 0.0!)
     """
+    brax_env = EcoAntV1(battery_limit=0.1, noise_scale=0.1)
+    wrapped_env = BraxToGymnaxWrapper(env=brax_env, episode_length=10)
     obs, state = wrapped_env.reset(key)
 
-    # 1. Force near-death
-    new_info = state.info.copy()
-    new_info["battery"] = jnp.array(0.1)
-    # We must also update 'obs' because the Wrapper might use 'obs' for something,
-    # though usually physics uses 'state'. Let's keep obs as is (1.0) to prove
-    # the physics engine checks 'info', not 'obs'.
-    state = state.replace(info=new_info)
-
     # 2. Kill the agent
     action = jnp.ones(wrapped_env.action_size)
     next_obs, next_state, reward, done, info = wrapped_env.step(key, state, action)
@@ -176,7 +171,7 @@ def test_wrapper_autoreset_logic(wrapped_env: BraxToGymnaxWrapper, key: jax.Arra
     # B. The 'info' contains the terminal cost of the DEAD agent
     # But includes the battery reset
     assert info["cost"] == 1.0
-    assert info["battery"] == 10.0
+    assert info["battery"] == 0.0
 
     # C. The 'next_obs' is from the NEW alive agent (AutoReset happened)
     # The observation returned is for the *next* step.
@@ -184,4 +179,4 @@ def test_wrapper_autoreset_logic(wrapped_env: BraxToGymnaxWrapper, key: jax.Arra
 
     # D. The 'next_state' is the NEW alive state
     assert next_state.info["cost"] == 1.0
-    assert next_state.info["battery"] == 10.0
+    assert next_state.info["battery"] == 0.0

{safenax-0.4.3 → safenax-0.4.4}/tests/test_eco_ant_v2.py

@@ -78,18 +78,12 @@ def test_battery_depletion_in_info(env: EcoAntV2, key: jax.Array):
     assert jnp.allclose(start_battery - energy_used, current_battery, atol=1e-5)
 
 
-def test_termination_logic(env: EcoAntV2, key: jax.Array):
+def test_termination_logic(key: jax.Array):
     """
     Verifies that the episode terminates when info['battery'] hits zero.
     """
-    state = env.reset(key)
-
-    # 1. Manually set battery to be very low in the info dict
-    # This simulates a "near death" state
-    low_battery_info = state.info.copy()
-    low_battery_info["battery"] = jnp.array(0.1)
-
-    near_death_state = state.replace(info=low_battery_info)
+    env = EcoAntV2(battery_limit=0.1, noise_scale=0.1)
+    near_death_state = env.reset(key)
 
     # 2. Take a large step to consume > 0.1 energy
     # Action of 1.0s usually consumes ~4.0 energy
@@ -100,7 +94,7 @@ def test_termination_logic(env: EcoAntV2, key: jax.Array):
     assert next_state.done == 1.0
 
     # 4. Verify Battery Floor
-    assert next_state.info["battery"] == 10.0
+    assert next_state.info["battery"] == 0.0
 
 
 def test_stochasticity_impact(env: EcoAntV2, key: jax.Array):

{safenax-0.4.3 → safenax-0.4.4}/uv.lock

@@ -1501,7 +1501,7 @@ wheels = [
 
 [[package]]
 name = "safenax"
-version = "0.4.2"
+version = "0.4.3"
 source = { editable = "." }
 dependencies = [
     { name = "brax" },