[envpool] Upgrade JAX 0.9 and harden state buffer fallback

Author: Trinkle23897

envpool/__init__.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 """EnvPool package for efficient RL environment simulation."""
 
+import numpy as np
+
 import envpool.entry  # noqa: F401
 from envpool.registration import (
   list_all_envs,
@@ -24,7 +26,11 @@
   register,
 )
 
-__version__ = "0.8.4"
+# Gym 0.26 still references np.bool8, which NumPy 2 removed.
+if not hasattr(np, "bool8"):
+  np.bool8 = np.bool_  # type: ignore[attr-defined]
+
+__version__ = "0.9.0"
 __all__ = [
   "register",
   "make",

envpool/core/state_buffer_queue.h

@@ -140,6 +140,10 @@ class StateBufferQueue {
    */
   std::vector<Array> Wait(std::size_t additional_done_count = 0) {
     std::unique_ptr<StateBuffer> newbuf = stock_buffer_.Get();
+    if (newbuf == nullptr) {
+      newbuf = std::make_unique<StateBuffer>(batch_, max_num_players_, specs_,
+                                             is_player_state_);
+    }
     std::size_t pos = done_ptr_.fetch_add(1);
     std::size_t offset = pos % queue_size_;
     auto arr = queue_[offset]->Wait(additional_done_count);

envpool/python/xla_template.py

@@ -22,7 +22,7 @@
 from jax import numpy as jnp
 from jax.core import ShapedArray
 from jax.interpreters import xla
-from jax.lib import xla_client  # type: ignore[attr-defined,unused-ignore]
+from jax._src.lib import xla_client
 
 _core = cast(Any, core)
 _xla = cast(Any, xla)
@@ -54,10 +54,13 @@ def _make_xla_function(
   specs: Tuple[Tuple[Any, ...], Tuple[Any, ...]],
   capsules: Tuple[Any, Any],
 ) -> Callable:
-  if not hasattr(_xla, "backend_specific_translations"):
+  if not (
+    hasattr(_core, "Primitive") and
+    hasattr(_xla, "backend_specific_translations")
+  ):
     raise RuntimeError(
       "XLA is unavailable because this JAX version removed the legacy "
-      "backend translation API used by envpool."
+      "primitive/translation APIs used by envpool."
     )
   in_specs, out_specs = specs
   in_specs = _normalize_specs(in_specs)

envpool/registration.py

@@ -18,10 +18,15 @@
 from typing import Any, Dict, List, Tuple
 
 import gym
+import numpy as np
 from packaging import version
 
 base_path = os.path.abspath(os.path.dirname(__file__))
 
+# Gym 0.26 still references np.bool8, which NumPy 2 removed.
+if not hasattr(np, "bool8"):
+  np.bool8 = np.bool_  # type: ignore[attr-defined]
+
 
 class EnvRegistry:
   """A collection of available envs."""

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = envpool
-version = 0.8.4
+version = 0.9.0
 author = "EnvPool Contributors"
 author_email = "sail@sea.com"
 description = "C++-based high-performance parallel environment execution engine (vectorized env) for general RL environments."

third_party/pip_requirements/requirements-dev-lock.txt

@@ -24,8 +24,8 @@ gymnasium==1.2.3
 h5py==3.16.0
 idna==3.11
 imageio==2.37.3
-jax==0.4.38
-jaxlib==0.4.38
+jax==0.9.2
+jaxlib==0.9.2
 jinja2==3.1.6
 labmaze==1.0.6
 llvmlite==0.46.0
@@ -39,7 +39,7 @@ mujoco==3.6.0
 mujoco-py==2.1.2.14
 networkx==3.6.1
 numba==0.64.0
-numpy==1.26.4
+numpy==2.2.6
 nvidia-cublas-cu12==12.8.4.1
 nvidia-cuda-cupti-cu12==12.8.90
 nvidia-cuda-nvrtc-cu12==12.8.93
@@ -55,7 +55,7 @@ nvidia-nccl-cu12==2.27.5
 nvidia-nvjitlink-cu12==12.8.93
 nvidia-nvshmem-cu12==3.4.5
 nvidia-nvtx-cu12==12.8.90
-opencv-python-headless==4.11.0.86
+opencv-python-headless==4.12.0.88
 opt-einsum==3.4.0
 optree==0.19.0
 packaging==26.0

third_party/pip_requirements/requirements-dev.txt

@@ -1,11 +1,11 @@
 setuptools
 wheel
-numpy<2
+numpy>=2,<3
 dm-env>=1.6
 gym>=0.26.2
 gymnasium>=1.2.3
 optree>=0.19.0
-jax[cpu]>=0.4.38,<0.5
+jax[cpu]>=0.9.2,<0.10
 absl-py
 packaging>=26.0
 tqdm

third_party/pip_requirements/requirements-release-lock.txt

@@ -7,10 +7,10 @@ farama-notifications==0.0.4
 gym==0.26.2
 gym-notices==0.1.0
 gymnasium==1.2.3
-jax==0.4.38
-jaxlib==0.4.38
+jax==0.9.2
+jaxlib==0.9.2
 ml-dtypes==0.5.4
-numpy==1.26.4
+numpy==2.4.3
 opt-einsum==3.4.0
 optree==0.19.0
 packaging==26.0

third_party/pip_requirements/requirements-release.txt

@@ -1,9 +1,9 @@
 setuptools
 wheel
-numpy<2
+numpy>=2,<3
 dm-env>=1.6
 gym>=0.26.2
 gymnasium>=1.2.3
 optree>=0.19.0
-jax[cpu]>=0.4.38,<0.5
+jax[cpu]>=0.9.2,<0.10
 packaging>=26.0