safenax 0.4.4__tar.gz → 0.4.6__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: safenax
-Version: 0.4.4
+Version: 0.4.6
 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.4 → safenax-0.4.6}/pyproject.toml

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

{safenax-0.4.4 → safenax-0.4.6}/safenax/__init__.py

@@ -8,7 +8,7 @@ from safenax.portfolio_optimization import (
     PortfolioOptimizationGARCH,
     PortfolioOptimizationCrypto,
 )
-from safenax.frozen_lake import FrozenLake
+from safenax.frozen_lake import FrozenLakeV1, FrozenLakeV2
 from safenax.eco_ant import EcoAntV1, EcoAntV2
 
 
@@ -16,7 +16,8 @@ __all__ = [
     "FragileAnt",
     "PortfolioOptimizationCrypto",
     "PortfolioOptimizationGARCH",
-    "FrozenLake",
+    "FrozenLakeV1",
+    "FrozenLakeV2",
     "EcoAntV1",
     "EcoAntV2",
 ]

safenax-0.4.6/safenax/frozen_lake/__init__.py

@@ -0,0 +1,8 @@
+from safenax.frozen_lake.frozen_lake_v1 import FrozenLakeV1
+from safenax.frozen_lake.frozen_lake_v2 import FrozenLakeV2
+
+
+__all__ = [
+    "FrozenLakeV1",
+    "FrozenLakeV2",
+]

safenax-0.4.4/safenax/frozen_lake.py → safenax-0.4.6/safenax/frozen_lake/frozen_lake_v1.py

@@ -1,4 +1,4 @@
-"""JAX-compatible FrozenLake environment following the gymnax interface."""
+"""JAX-compatible FrozenLakeV1 environment following the gymnax interface."""
 
 from typing import Optional, Tuple, Union, List
 
@@ -69,7 +69,7 @@ MAPS = {
 }
 
 
-class FrozenLake(environment.Environment):
+class FrozenLakeV1(environment.Environment):
     """
     JAX-compatible FrozenLake environment.
 
@@ -81,7 +81,7 @@ class FrozenLake(environment.Environment):
     def __init__(
         self,
         map_name: str = "4x4",
-        desc: Optional[chex.Array] = None,
+        desc: Optional[jax.Array] = None,
         is_slippery: bool = True,
         success_rate: float = 1.0 / 3.0,
         reward_schedule: Tuple[float, float, float] = (1.0, 0.0, 0.0),
@@ -250,9 +250,9 @@ def make_frozen_lake(
     is_slippery: bool = True,
     success_rate: float = 1.0 / 3.0,
     reward_schedule: Tuple[float, float, float] = (1.0, 0.0, 0.0),
-) -> FrozenLake:
+) -> FrozenLakeV1:
     """Factory function to easier initialization."""
-    return FrozenLake(
+    return FrozenLakeV1(
         map_name=map_name,
         is_slippery=is_slippery,
         success_rate=success_rate,

safenax-0.4.6/safenax/frozen_lake/frozen_lake_v2.py

@@ -0,0 +1,302 @@
+"""JAX-compatible FrozenLakeV2 environment following the gymnax interface."""
+
+from typing import Optional, Tuple, Union, List
+
+import chex
+import jax
+import jax.numpy as jnp
+from flax import struct
+from gymnax.environments import environment, spaces
+
+# --- Constants ---
+# Actions
+LEFT = 0
+DOWN = 1
+RIGHT = 2
+UP = 3
+
+# Tile Types (using ASCII values for readability during debugging)
+TILE_START = ord("S")
+TILE_FROZEN = ord("F")
+TILE_THIN = ord("T")
+TILE_GOAL = ord("G")
+
+
+@struct.dataclass
+class EnvState:
+    """Environment state for FrozenLake."""
+
+    pos: int
+    time: int
+
+
+@struct.dataclass
+class EnvParams:
+    """Environment parameters for FrozenLake."""
+
+    desc: jax.Array  # Map description as integer array (nrow, ncol)
+    nrow: int
+    ncol: int
+    is_slippery: bool
+    success_rate: float
+    reward_schedule: jax.Array  # (goal_reward, thin_reward, frozen_reward)
+    max_steps_in_episode: int
+    safe_cost_mean: float
+    safe_cost_std: float
+    thin_cost_base: float
+    thin_shock_prob: float
+    thin_shock_val: float
+
+
+# --- Helper Functions ---
+
+
+def string_map_to_array(map_desc: List[str]) -> jax.Array:
+    """Convert a list of strings into a JAX-friendly integer array."""
+    return jnp.array([[ord(c) for c in row] for row in map_desc], dtype=jnp.int32)
+
+
+# Predefined maps
+MAPS = {
+    "4x4": string_map_to_array(["SFFF", "FTFT", "FFFT", "TFFG"]),
+    "8x8": string_map_to_array(
+        [
+            "SFFFFFFF",
+            "FFFFFFFF",
+            "FFFTFFFF",
+            "FFFFFTFF",
+            "FFFTFFFF",
+            "FTTFFFTF",
+            "FTFFTFTF",
+            "FFFTFFFG",
+        ]
+    ),
+}
+
+
+class FrozenLakeV2(environment.Environment):
+    """
+    JAX-compatible FrozenLake environment with "Thin Ice" cost augmentation.
+
+    The agent controls the movement of a character in a grid world to reach a Goal (G).
+    The map contains two types of traversable terrain that introduce a Mean-Variance trade-off:
+
+    1. Safe Ice ('F'):
+       - Represents a stable but arduous path (deep snow is hard to traverse).
+       - Cost: High Mean (e.g., 2.0), Small Variance (e.g., std=0.1).
+       - Risk: Deterministically expensive, but safe.
+
+    2. Thin Ice ('T', formerly Holes):
+       - Represents a dangerous shortcut, prone to cracking or falling but can glide over with little effort.
+       - Cost: Low Mean (e.g., 1.5), High Variance Shocks (e.g. std=4.35) (ice cracks).
+       - Risk: Usually cheap to traverse, but prone to random, catastrophic cost spikes (cracking ice).
+
+    The agent can also slip perpendicular to its intended direction based on the `is_slippery` and `success_rate` parameters.
+    """
+
+    def __init__(
+        self,
+        map_name: str = "4x4",
+        desc: Optional[jax.Array] = None,
+        is_slippery: bool = True,
+        success_rate: float = 1.0 / 3.0,
+        reward_schedule: Tuple[float, float, float] = (1.0, 0.0, 0.0),
+        safe_cost_mean: float = 2.0,
+        safe_cost_std: float = 0.1,
+        thin_cost_base: float = 0.5,
+        thin_shock_prob: float = 0.05,
+        thin_shock_val: float = 20.0,
+    ):
+        super().__init__()
+
+        if desc is None:
+            desc = MAPS[map_name]
+
+        self.desc = desc
+        self.nrow, self.ncol = desc.shape
+        self.is_slippery = is_slippery
+        self.success_rate = success_rate
+        self.reward_schedule = jnp.array(reward_schedule)
+        self.safe_cost_mean = safe_cost_mean
+        self.safe_cost_std = safe_cost_std
+        self.thin_cost_base = thin_cost_base
+        self.thin_shock_prob = thin_shock_prob
+        self.thin_shock_val = thin_shock_val
+
+        # Determine max steps based on map size convention
+        if map_name == "4x4" or (desc.shape[0] == 4 and desc.shape[1] == 4):
+            self.max_steps = 100
+        else:
+            self.max_steps = 200
+
+    @property
+    def default_params(self) -> EnvParams:
+        return EnvParams(
+            desc=self.desc,
+            nrow=self.nrow,
+            ncol=self.ncol,
+            is_slippery=self.is_slippery,
+            success_rate=self.success_rate,
+            reward_schedule=self.reward_schedule,
+            max_steps_in_episode=self.max_steps,
+            safe_cost_mean=self.safe_cost_mean,
+            safe_cost_std=self.safe_cost_std,
+            thin_cost_base=self.thin_cost_base,
+            thin_shock_prob=self.thin_shock_prob,
+            thin_shock_val=self.thin_shock_val,
+        )
+
+    def step_env(
+        self,
+        key: jax.Array,
+        state: EnvState,
+        action: Union[int, float],
+        params: EnvParams,
+    ) -> Tuple[chex.Array, EnvState, float, bool, dict]:
+        """Perform a single environment step with JIT compatibility."""
+        action = jnp.int32(action)
+
+        # 1. Determine the actual direction (handling slippery logic)
+        rng_slip, rng_cost, key = jax.random.split(key, 3)
+
+        def get_slippery_action(k):
+            # 0: intended, 1: perpendicular left, 2: perpendicular right
+            # Probabilities: [success, (1-success)/2, (1-success)/2]
+            fail_prob = (1.0 - params.success_rate) / 2.0
+            probs = jnp.array([params.success_rate, fail_prob, fail_prob])
+
+            # Map samples (0, 1, 2) to actions (action, action-1, action+1)
+            delta_idx = jax.random.choice(k, jnp.arange(3), p=probs)
+
+            # (action - 1) % 4  <-- Perpendicular Left
+            # (action)          <-- Intended
+            # (action + 1) % 4  <-- Perpendicular Right
+
+            # Using a lookup array for cleaner mapping
+            candidates = jnp.array(
+                [
+                    action,  # Index 0: Success
+                    (action - 1) % 4,  # Index 1: Fail Left
+                    (action + 1) % 4,  # Index 2: Fail Right
+                ]
+            )
+            return candidates[delta_idx]
+
+        actual_action = jax.lax.cond(
+            params.is_slippery,
+            get_slippery_action,
+            lambda k: action,  # Deterministic branch
+            rng_slip,
+        )
+
+        # 2. Calculate Movement (Optimized Vectorized Approach)
+        row = state.pos // params.ncol
+        col = state.pos % params.ncol
+
+        next_row, next_col = self._apply_action(row, col, actual_action, params)
+        next_pos = next_row * params.ncol + next_col
+
+        # 3. Check Rewards and Termination
+        tile_type = params.desc[next_row, next_col]
+
+        is_goal = tile_type == TILE_GOAL
+        is_thin = tile_type == TILE_THIN
+
+        # Reward Schedule: [Goal, Thin, Frozen/Start]
+        # We select index 0, 1, or 2 based on tile type
+        tile_idx = jnp.where(is_goal, 0, jnp.where(is_thin, 1, 2))
+        reward = params.reward_schedule[tile_idx]
+
+        safe_noise = jax.random.normal(rng_cost) * params.safe_cost_std
+        cost_safe = params.safe_cost_mean + safe_noise
+
+        is_shock = jax.random.bernoulli(rng_cost, p=params.thin_shock_prob)
+        cost_thin = params.thin_cost_base + (is_shock * params.thin_shock_val)
+
+        cost = jnp.select([is_goal, is_thin], [0.0, cost_thin], default=cost_safe)
+
+        # 4. Update State
+        # Time limit truncation is handled in Gymnax wrappers usually,
+        # but we track it here for the 'done' flag consistency.
+        new_time = state.time + 1
+        truncated = new_time >= params.max_steps_in_episode
+        done = is_goal | truncated
+
+        new_state = EnvState(pos=next_pos, time=new_time)
+
+        return (
+            self.get_obs(new_state, params),
+            new_state,
+            reward,
+            done,
+            {"cost": cost, "tile_type": tile_type},
+        )
+
+    def reset_env(
+        self,
+        key: chex.PRNGKey,
+        params: EnvParams,
+    ) -> Tuple[chex.Array, EnvState]:
+        """Reset environment to the start position."""
+        # By definition, FrozenLake always starts at (0,0)
+        # If dynamic start positions are needed, one would scan params.desc for 'S'
+        state = EnvState(pos=0, time=0)
+        return self.get_obs(state, params), state
+
+    def get_obs(self, state: EnvState, params: EnvParams) -> chex.Array:
+        """Return scalar observation (flat position index)."""
+        return jnp.array(state.pos, dtype=jnp.int32)
+
+    def _apply_action(
+        self, row: int, col: int, action: int, params: EnvParams
+    ) -> Tuple[int, int]:
+        """Optimized movement logic using delta arrays and clipping."""
+        # Gym Direction Mapping:
+        # 0: Left  (0, -1)
+        # 1: Down  (1,  0)
+        # 2: Right (0,  1)
+        # 3: Up    (-1, 0)
+
+        dr = jnp.array([0, 1, 0, -1])
+        dc = jnp.array([-1, 0, 1, 0])
+
+        new_row = row + dr[action]
+        new_col = col + dc[action]
+
+        # Ensure we stay within the grid
+        new_row = jnp.clip(new_row, 0, params.nrow - 1)
+        new_col = jnp.clip(new_col, 0, params.ncol - 1)
+
+        return new_row, new_col
+
+    @property
+    def name(self) -> str:
+        return "FrozenLake-v2"
+
+    @property
+    def num_actions(self) -> int:
+        return 4
+
+    def action_space(self, params: Optional[EnvParams] = None) -> spaces.Discrete:
+        return spaces.Discrete(4)
+
+    def observation_space(self, params: EnvParams) -> spaces.Discrete:
+        return spaces.Discrete(params.nrow * params.ncol)
+
+
+# --- Factory Helper ---
+
+
+def make_frozen_lake(
+    map_name: str = "4x4",
+    is_slippery: bool = True,
+    success_rate: float = 1.0 / 3.0,
+    reward_schedule: Tuple[float, float, float] = (1.0, 0.0, 0.0),
+) -> FrozenLakeV2:
+    """Factory function to easier initialization."""
+    return FrozenLakeV2(
+        map_name=map_name,
+        is_slippery=is_slippery,
+        success_rate=success_rate,
+        reward_schedule=reward_schedule,
+    )

safenax-0.4.4/tests/test_frozen_lake.py → safenax-0.4.6/tests/test_frozen_lake_v1.py

@@ -1,8 +1,8 @@
 import pytest
 import jax
 import jax.numpy as jnp
-from safenax.frozen_lake import FrozenLake, EnvParams
-from safenax.frozen_lake import (
+from safenax.frozen_lake.frozen_lake_v1 import FrozenLakeV1, EnvParams
+from safenax.frozen_lake.frozen_lake_v1 import (
     TILE_START,
     TILE_FROZEN,
     TILE_HOLE,
@@ -16,7 +16,7 @@ from safenax.frozen_lake import (
 @pytest.fixture
 def env():
     """Initializes the environment."""
-    return FrozenLake()
+    return FrozenLakeV1()
 
 
 @pytest.fixture
@@ -58,7 +58,7 @@ def params_slippery(params_deterministic: EnvParams) -> EnvParams:
 
 
 def test_jit_compilation(
-    env: FrozenLake, rng: jax.Array, params_deterministic: EnvParams
+    env: FrozenLakeV1, rng: jax.Array, params_deterministic: EnvParams
 ):
     """Verifies that reset and step can be JIT compiled without errors."""
     reset_jit = jax.jit(env.reset)
@@ -85,7 +85,7 @@ def test_map_parsing():
 
 
 def test_deterministic_reach_goal(
-    env: FrozenLake, rng: jax.Array, params_deterministic: EnvParams
+    env: FrozenLakeV1, rng: jax.Array, params_deterministic: EnvParams
 ):
     """
     Path: Start(0) -> Right(1) -> Down(3/Goal).
@@ -119,7 +119,7 @@ def test_deterministic_reach_goal(
 
 
 def test_deterministic_fall_in_hole(
-    env: FrozenLake, rng: jax.Array, params_deterministic: EnvParams
+    env: FrozenLakeV1, rng: jax.Array, params_deterministic: EnvParams
 ):
     """
     Path: Start(0) -> Down(2/Hole).
@@ -142,7 +142,7 @@ def test_deterministic_fall_in_hole(
 
 
 def test_custom_reward_schedule(
-    env: FrozenLake, rng: jax.Array, params_deterministic: EnvParams
+    env: FrozenLakeV1, rng: jax.Array, params_deterministic: EnvParams
 ):
     """
     Verifies that changing reward_schedule affects the output.
@@ -167,7 +167,9 @@ def test_custom_reward_schedule(
     assert reward == -5.0
 
 
-def test_slippery_dynamics(env: FrozenLake, rng: jax.Array, params_slippery: EnvParams):
+def test_slippery_dynamics(
+    env: FrozenLakeV1, rng: jax.Array, params_slippery: EnvParams
+):
     """
     Statistical test for slippery logic.
     We start at (0,0) and try to move Right.
@@ -201,7 +203,7 @@ def test_slippery_dynamics(env: FrozenLake, rng: jax.Array, params_slippery: Env
     assert 1 in unique_positions
 
 
-def test_truncation(env: FrozenLake, rng: jax.Array, params_deterministic: EnvParams):
+def test_truncation(env: FrozenLakeV1, rng: jax.Array, params_deterministic: EnvParams):
     """Test that episode ends when time limit is reached."""
     # Set max steps to 2
     short_params = params_deterministic.replace(max_steps_in_episode=2)
@@ -224,7 +226,9 @@ def test_truncation(env: FrozenLake, rng: jax.Array, params_deterministic: EnvPa
     assert state.time == 0
 
 
-def test_jit_rollout_scan(env: FrozenLake, rng: jax.Array, params_slippery: EnvParams):
+def test_jit_rollout_scan(
+    env: FrozenLakeV1, rng: jax.Array, params_slippery: EnvParams
+):
     """
     Verifies that the environment can be run in a fully compiled jax.lax.scan loop.
     This ensures no shape mismatches or control flow issues exist in the step logic.
@@ -282,7 +286,9 @@ def test_jit_rollout_scan(env: FrozenLake, rng: jax.Array, params_slippery: EnvP
     assert jnp.all(jnp.isin(rollout_data["reward"], jnp.array([0.0, 1.0])))
 
 
-def test_cost_signal(env: FrozenLake, rng: jax.Array, params_deterministic: EnvParams):
+def test_cost_signal(
+    env: FrozenLakeV1, rng: jax.Array, params_deterministic: EnvParams
+):
     """Ensure cost is 1.0 ONLY when falling into a hole."""
     step_jit = jax.jit(env.step)
     reset_jit = jax.jit(env.reset)

safenax-0.4.6/tests/test_frozen_lake_v2.py

@@ -0,0 +1,167 @@
+"""Unit tests for FrozenLakeV2 environment."""
+
+import jax
+import jax.numpy as jnp
+import pytest
+from safenax.frozen_lake.frozen_lake_v2 import (
+    FrozenLakeV2,
+    EnvState,
+    TILE_FROZEN,
+    TILE_THIN,
+    TILE_GOAL,
+)
+
+# Fix random key for reproducibility
+KEY = jax.random.PRNGKey(42)
+
+
+@pytest.fixture
+def env():
+    """Fixture to create a standard 4x4 FrozenLakeV2 environment."""
+    return FrozenLakeV2(map_name="4x4", is_slippery=True)
+
+
+@pytest.fixture
+def params(env):
+    """Fixture for default environment parameters."""
+    return env.default_params
+
+
+def test_initialization(env, params):
+    """Test that the environment initializes and resets correctly."""
+    key = jax.random.PRNGKey(0)
+    obs, state = env.reset_env(key, params)
+
+    # Check initial state
+    assert state.pos == 0
+    assert state.time == 0
+    assert obs.shape == ()  # Scalar observation
+    assert obs == 0
+
+
+def test_step_mechanics(env, params):
+    """Test basic stepping functionality."""
+    key = jax.random.PRNGKey(0)
+    _, state = env.reset_env(key, params)
+
+    # Take an action (e.g., RIGHT = 2)
+    action = 2
+    obs, next_state, reward, done, info = env.step_env(key, state, action, params)
+
+    # Basic shape/type checks
+    assert isinstance(next_state, EnvState)
+    assert next_state.time == 1
+    assert "cost" in info
+    assert info["cost"] >= 0
+
+
+def test_slippery_dynamics(env, params):
+    """
+    Test that the agent actually slips when is_slippery=True.
+    Starting at (1,1), taking action RIGHT (2) should stochastically result
+    in moving RIGHT, UP, or DOWN.
+    """
+    key = jax.random.PRNGKey(0)
+    keys = jax.random.split(key, 100)  # Run 100 steps
+
+    # Force state to be at (1,1) [pos=5 for 4x4] to allow movement in all directions
+    # Map:
+    # S F F F (0, 1, 2, 3)
+    # F T F T (4, 5, 6, 7)
+    state = EnvState(pos=5, time=0)
+    action = 2  # RIGHT
+
+    # Vmap the step function to run in parallel
+    def step_fn(k):
+        return env.step_env(k, state, action, params)[1].pos
+
+    next_positions = jax.vmap(step_fn)(keys)
+
+    # From 5 (1,1):
+    # Right -> 6 (1,2) [Intended]
+    # Down  -> 9 (2,1) [Slip Right]
+    # Up    -> 1 (0,1) [Slip Left]
+
+    unique_positions = jnp.unique(next_positions)
+
+    # Assert that we ended up in more than just the intended position
+    assert len(unique_positions) > 1
+    assert 6 in unique_positions  # Intended
+    assert (9 in unique_positions) or (1 in unique_positions)  # Slips
+
+
+def test_cost_mean_variance(env, params):
+    """
+    CRITICAL TEST: Verify the Mean-Variance trade-off for VaR-CPO.
+    Safe Ice should have Mean~2.0, Std~0.1.
+    Thin Ice should have Mean~1.5, Std~4.35.
+    """
+    N = 5000  # High sample count for statistical significance
+    key = jax.random.PRNGKey(0)
+    keys = jax.random.split(key, N)
+
+    # We will force the map description to be entirely one type to isolate the cost function logic
+    # regardless of where the agent moves.
+
+    def get_batch_costs(tile_type):
+        """Helper to get costs for N steps on a map filled with `tile_type`."""
+        # Create a dummy map filled with the specific tile type
+        custom_desc = jnp.full((4, 4), tile_type, dtype=jnp.int32)
+        custom_params = params.replace(desc=custom_desc)
+
+        def step_fn(k):
+            # Start at 0, take action 0. Destination will be same tile type.
+            _, _, _, _, info = env.step_env(k, EnvState(0, 0), 0, custom_params)
+            return info["cost"]
+
+        return jax.vmap(step_fn)(keys)
+
+    # 1. Test Safe Ice ('F')
+    safe_costs = get_batch_costs(TILE_FROZEN)
+
+    # Check Safe Statistics: Mean ~ 2.0, Std ~ 0.1
+    assert jnp.abs(jnp.mean(safe_costs) - 2.0) < 0.1, (
+        f"Safe Mean {jnp.mean(safe_costs)} != 2.0"
+    )
+    assert jnp.std(safe_costs) < 0.2, f"Safe Std {jnp.std(safe_costs)} is too high"
+
+    # 2. Test Thin Ice ('T')
+    thin_costs = get_batch_costs(TILE_THIN)
+
+    # Check Thin Statistics: Mean ~ 1.5, Std ~ 4.35
+    assert jnp.abs(jnp.mean(thin_costs) - 1.5) < 0.1, (
+        f"Thin Mean {jnp.mean(thin_costs)} != 1.5"
+    )
+    assert jnp.abs(jnp.std(thin_costs) - 4.35) < 0.1, (
+        f"Thin Std {jnp.std(thin_costs)} is too low (expected ~4.35)"
+    )
+
+    print(
+        f"\nStats Verification:\nSafe: Mean={jnp.mean(safe_costs):.2f}, Std={jnp.std(safe_costs):.2f}"
+    )
+    print(f"Thin: Mean={jnp.mean(thin_costs):.2f}, Std={jnp.std(thin_costs):.2f}")
+
+
+def test_termination(env, params):
+    """Test episode termination logic."""
+    key = jax.random.PRNGKey(0)
+
+    # 1. Test Goal Termination
+    # Force map to be all GOAL
+    goal_desc = jnp.full((4, 4), TILE_GOAL, dtype=jnp.int32)
+    goal_params = params.replace(desc=goal_desc)
+
+    # Agent takes a step on a Goal tile -> should be done
+    _, _, _, done, _ = env.step_env(key, EnvState(0, 0), 0, goal_params)
+    assert done == True
+
+    # 2. Test Time Truncation
+    # Set max_steps to 1
+    short_params = params.replace(max_steps_in_episode=1)
+
+    # Step 1: time becomes 1 -> truncated >= max_steps (1) -> done=True
+    state = EnvState(0, 0)
+    _, next_state, _, done, _ = env.step_env(key, state, 0, short_params)
+
+    assert next_state.time == 1
+    assert done == True

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

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