Merge pull request #103 from epignatelli/i92

Upgrade to Python>=3.11

Author: epignatelli

.github/workflows/CI.yml

@@ -12,12 +12,13 @@ jobs:
       max-parallel: 5
       matrix:
         os: ["ubuntu"]
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
     continue-on-error: false
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v4
         with:
-          python-version: "3.10"
+          python-version: ${{ matrix.python-version }}
       - name: Setup navix
         run: |
           pip install . -v

examples/purejaxrl/ppo_minigrid.py

@@ -225,14 +225,14 @@ def _loss_fn(params, traj_batch, gae, targets):
                 ), "batch size must be equal to number of steps * number of envs"
                 permutation = jax.random.permutation(_rng, batch_size)
                 batch = (traj_batch, advantages, targets)
-                batch = jax.tree_util.tree_map(
+                batch = jax.tree.map(
                     lambda x: x.reshape((batch_size,) + x.shape[2:]), batch
                 )
-                shuffled_batch = jax.tree_util.tree_map(
+                shuffled_batch = jax.tree.map(
                     lambda x: jnp.take(x, permutation, axis=0), batch
                 )
                 # Mini-batch Updates
-                minibatches = jax.tree_util.tree_map(
+                minibatches = jax.tree.map(
                     lambda x: jnp.reshape(
                         x, [config["NUM_MINIBATCHES"], -1] + list(x.shape[1:])
                     ),

navix/_version.py

@@ -18,5 +18,5 @@
 # under the License.
 
 
-__version__ = "0.7.0"
+__version__ = "0.7.1"
 __version_info__ = tuple(int(i) for i in __version__.split(".") if i.isdigit())

navix/agents/ppo.py

@@ -252,15 +252,15 @@ def update(self, train_state: TrainingState, _) -> Tuple[TrainingState, Dict]:
             ), "batch size must be equal to number of steps * number of envs"
             permutation = jax.random.permutation(rng_1, n_samples)
             samples = (experience, advantages, targets, values)  # (T, N, ...)
-            samples = jax.tree_util.tree_map(
+            samples = jax.tree.map(
                 lambda x: x.reshape((n_samples,) + x.shape[2:]), samples
             )  # (T * N, ...)
-            shuffled_batch = jax.tree_util.tree_map(
+            shuffled_batch = jax.tree.map(
                 lambda x: jnp.take(x, permutation, axis=0), samples
             )  # (T * N, ...)
 
             # One epoch update over all mini-batches
-            minibatches = jax.tree_util.tree_map(
+            minibatches = jax.tree.map(
                 lambda x: jnp.reshape(
                     x, (self.hparams.num_minibatches, -1) + tuple(x.shape[1:])
                 ),

navix/entities.py

@@ -66,7 +66,7 @@ class Entity(Positionable, HasTag, HasSprite):
     To create an entity, use the `create` method."""
 
     def __getitem__(self: T, idx) -> T:
-        return jax.tree_util.tree_map(lambda x: x[idx], self)
+        return jax.tree.map(lambda x: x[idx], self)
 
     @property
     def name(self) -> str:

navix/environments/key_corridor.py

@@ -24,7 +24,6 @@
 import jax
 import jax.numpy as jnp
 from jax import Array
-import jax.tree_util as jtu
 
 from navix import observations, rewards, terminations
 
@@ -116,7 +115,7 @@ def _reset(self, key: Array, cache: Union[RenderingCache, None] = None) -> Times
                     open=jnp.asarray(0),
                 )
             )
-        doors = jtu.tree_map(lambda *x: jnp.stack(x), *doors)
+        doors = jax.tree.map(lambda *x: jnp.stack(x), *doors)
 
         entities = {
             "player": player[None],

navix/experiment.py

@@ -171,10 +171,10 @@ def search(hparam_set_sample):
         # average over seeds
         for i in range(len_search_set):
             print("Logging results for hparam set:", search_set)
-            hparams = jax.tree_map(lambda x: x[i], search_set)
+            hparams = jax.tree.map(lambda x: x[i], search_set)
             config = {**vars(self), **asdict(hparams)}
             wandb.init(project=self.name, config=config, group=self.group)
-            log = jax.tree_map(lambda x: jnp.mean(x[i], axis=0), logs)
+            log = jax.tree.map(lambda x: jnp.mean(x[i], axis=0), logs)
             self.agent.log_on_train_end(log)
             wandb.finish()
         logging_time = time.time() - start_time

navix/grid.py

@@ -26,7 +26,6 @@
 import jax
 import jax.numpy as jnp
 from jax import Array
-import jax.tree_util as jtu
 from flax import struct
 
 

navix/states.py

@@ -22,6 +22,7 @@
 from jax import Array
 import jax.numpy as jnp
 from flax import struct
+from dataclasses import field
 
 
 from .components import (
@@ -63,10 +64,12 @@ class Event(Positionable, HasColour):
         happened (Array): A boolean flag indicating whether the event happened.
         event_type (Array): The type of event that happened."""
 
-    position: Array = jnp.asarray([-1, -1], dtype=jnp.int32)
-    colour: Array = PALETTE.UNSET
-    happened: Array = jnp.asarray(False, dtype=jnp.bool_)
-    event_type: Array = EventType.NONE
+    position: Array = field(
+        default_factory=lambda: jnp.asarray([-1, -1], dtype=jnp.int32)
+    )
+    colour: Array = field(default_factory=lambda: PALETTE.UNSET)
+    happened: Array = field(default_factory=lambda: jnp.asarray(False, dtype=jnp.bool_))
+    event_type: Array = field(default_factory=lambda: EventType.NONE)
 
     def __eq__(self, other: Event) -> Array:
         return jnp.logical_and(
@@ -160,10 +163,10 @@ def record_goal_reached(self, goal: Goal, position: Array) -> EventsManager:
     def record_ball_hit(self, ball: Ball) -> EventsManager:
         """Flags an event when the player is hit by a ball as happened and returns the
         updated events manager.
-        
+
         Args:
             ball (Ball): The ball that hit the player.
-        
+
         Returns:
             EventsManager: The updated events manager."""
         return self.replace(
@@ -178,11 +181,11 @@ def record_ball_hit(self, ball: Ball) -> EventsManager:
     def record_wall_hit(self, wall: Wall, position: Array) -> EventsManager:
         """Flags an event when the player hits a wall as happened and returns the
         updated events manager.
-        
+
         Args:
             wall (Wall): The wall the player hit.
             position (Array): The position of the wall in the grid.
-        
+
         Returns:
             EventsManager: The updated events manager."""
         idx = jnp.where(wall.position == position, size=1)[0][0]
@@ -199,10 +202,10 @@ def record_wall_hit(self, wall: Wall, position: Array) -> EventsManager:
     def record_grid_hit(self, position: Array) -> EventsManager:
         """Flags an event when the player hits a wall as happened and returns the
         updated events manager.
-        
+
         Args:
             position (Array): The position of the wall in the grid.
-        
+
         Returns:
             EventsManager: The updated events manager."""
         return self.replace(
@@ -217,11 +220,11 @@ def record_grid_hit(self, position: Array) -> EventsManager:
     def record_lava_fall(self, lava: Lava, position: Array) -> EventsManager:
         """Flags an event when the lava falls as happened and returns the
         updated events manager.
-        
+
         Args:
             lava (Lava): The lava that fell.
             position (Array): The position of the lava in the grid.
-        
+
         Returns:
             EventsManager: The updated events manager."""
         idx = jnp.where(lava.position == position, size=1)[0][0]
@@ -238,11 +241,11 @@ def record_lava_fall(self, lava: Lava, position: Array) -> EventsManager:
     def record_key_pickup(self, key: Key, position: Array) -> EventsManager:
         """Flags an event when the player picks up a key as happened and returns the
         updated events manager.
-        
+
         Args:
             key (Key): The key the player picked up.
             position (Array): The position of the key in the grid.
-        
+
         Returns:
             EventsManager: The updated events manager."""
         idx = jnp.where(key.position == position, size=1)[0][0]
@@ -259,11 +262,11 @@ def record_key_pickup(self, key: Key, position: Array) -> EventsManager:
     def record_door_opening(self, door: Door, position: Array) -> EventsManager:
         """Flags an event when the player opens a door as happened and returns the
         updated events manager.
-        
+
         Args:
             door (Door): The door the player opened.
             position (Array): The position of the door in the grid.
-        
+
         Returns:
             EventsManager: The updated events manager."""
         idx = jnp.where(door.position == position, size=1)[0][0]
@@ -284,7 +287,7 @@ def record_door_unlock(self, door: Door, position: Array) -> EventsManager:
         Args:
             door (Door): The door the player unlocked.
             position (Array): The position of the door in the grid.
-        
+
         Returns:
             EventsManager: The updated events manager."""
         idx = jnp.where(door.position == position, size=1)[0][0]
@@ -301,11 +304,11 @@ def record_door_unlock(self, door: Door, position: Array) -> EventsManager:
     def record_ball_pickup(self, ball: Ball, position: Array) -> EventsManager:
         """Flags an event when the player picks up a ball as happened and returns the
         updated events manager.
-        
+
         Args:
             ball (Ball): The ball the player picked up.
             position (Array): The position of the ball in the grid.
-        
+
         Returns:
             EventsManager: The updated events manager."""
         idx = jnp.where(ball.position == position, size=1)[0][0]
@@ -339,21 +342,21 @@ class State(struct.PyTreeNode):
 
     def get_entity(self, entity_enum: str) -> Entity:
         """Get an entity from the state by its enum.
-        
+
         Args:
             entity_enum (str): The enum of the entity to get.
-        
+
         Returns:
             Entity: The entity from the state."""
         return self.entities[entity_enum]
 
     def set_entity(self, entity_enum: str, entity: Entity) -> State:
         """Set an entity in the state by its enum.
-        
+
         Args:
             entity_enum (str): The enum of the entity to set.
             entity (Entity): The entity to set.
-        
+
         Returns:
             State: The updated state."""
         self.entities[entity_enum] = entity

pyproject.toml

@@ -9,7 +9,7 @@ requires = ["setuptools >= 50", "setuptools-scm[toml]>=6.2", "wheel"]
 name = "Navix"
 dynamic = ["version", "dependencies"]
 description = "Accelerated gridworld navigation with JAX for deep reinforcement learning"
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 readme = "README.md"
 license = {file = "LICENSE", name = "Apache-2.0"}
 authors = [