feat: use protocol to force all environment states to have a key

jumanji/env.py

@@ -15,14 +15,22 @@
 """Abstract environment class"""
 
 import abc
-from typing import Any, Generic, Tuple, TypeVar
+from typing import Any, Generic, Protocol, Tuple, TypeVar
 
+import jax.random
 from chex import PRNGKey
 
 from jumanji import specs
 from jumanji.types import Action, TimeStep
 
-State = TypeVar("State")
+
+class StateProtocol(Protocol):
+    """Enforce that the State for every Environment must implement a key."""
+
+    key: jax.random.PRNGKey
+
+
+State = TypeVar("State", bound="StateProtocol")
 
 
 class Environment(abc.ABC, Generic[State]):

jumanji/environments/combinatorial/cvrp/types.py

@@ -14,6 +14,8 @@
 
 from typing import TYPE_CHECKING, NamedTuple
 
+import jax.random
+
 if TYPE_CHECKING:  # https://github.com/python/mypy/issues/6239
     from dataclasses import dataclass
 else:
@@ -42,6 +44,7 @@ class State:
     visited_mask: Array  # (problem_size + 1,)
     order: Array  # (2 * problem_size,) - this size is worst-case (visit depot after each node)
     num_total_visits: jnp.int32
+    key: jax.random.PRNGKey = jax.random.PRNGKey(0)
 
 
 class Observation(NamedTuple):

jumanji/environments/combinatorial/knapsack/types.py

@@ -14,6 +14,8 @@
 
 from typing import TYPE_CHECKING, NamedTuple
 
+import jax.random
+
 if TYPE_CHECKING:  # https://github.com/python/mypy/issues/6239
     from dataclasses import dataclass
 else:
@@ -40,6 +42,7 @@ class State:
     used_mask: Array  # (problem_size,)
     num_steps: jnp.int32
     remaining_budget: jnp.float32
+    key: jax.random.PRNGKey = jax.random.PRNGKey(0)
 
 
 class Observation(NamedTuple):

jumanji/environments/combinatorial/tsp/types.py

@@ -14,6 +14,8 @@
 
 from typing import TYPE_CHECKING, NamedTuple
 
+import jax.random
+
 if TYPE_CHECKING:  # https://github.com/python/mypy/issues/6239
     from dataclasses import dataclass
 else:
@@ -38,6 +40,7 @@ class State:
     visited_mask: Array  # (problem_size,)
     order: Array  # (problem_size,)
     num_visited: jnp.int32
+    key: jax.random.PRNGKey = jax.random.PRNGKey(0)
 
 
 class Observation(NamedTuple):

jumanji/environments/games/connect4/types.py

@@ -15,6 +15,7 @@
 from typing import TYPE_CHECKING
 
 import jax.numpy as jnp
+import jax.random
 from chex import Array
 from typing_extensions import TypeAlias
 
@@ -30,6 +31,7 @@
 class State:
     current_player: jnp.int8
     board: Board
+    key: jax.random.PRNGKey = jax.random.PRNGKey(0)
 
 
 @dataclass

jumanji/wrappers.py

@@ -35,10 +35,9 @@
 from jax import jit, random
 
 from jumanji import specs, tree_utils
-from jumanji.env import Environment
+from jumanji.env import Environment, State
 from jumanji.types import Action, TimeStep, restart, termination, transition
 
-State = TypeVar("State")
 Observation = TypeVar("Observation")
 
 # Type alias that corresponds to ObsType in the Gym API