[envpool] upgrade JAX to 0.9 and harden state buffer fallback

[Trinkle23897]

Summary

• Problem: EnvPool still pins JAX below 0.5 and NumPy below 2, while the dummy env stress test can segfault in the multi-producer state buffer preallocation path.
• Scope: Upgrade the Python dependency window to JAX 0.9 and NumPy 2, add the minimal compatibility shims envpool still needs, and harden StateBufferQueue::Wait() when the background stock buffer temporarily returns nullptr.
• Outcome: The package version moves to 0.9.0, JAX 0.9 / NumPy 2 installs are supported by default, and the recurring dummy_envpool_test exit-139 failure is removed without dropping the multi-producer optimization.

This keeps the JAX upgrade and the dummy flake fix in one PR because the branch was already carrying both changes and the state buffer hardening is small.

Technical Details

• Approach: Raise the JAX and NumPy requirement bounds, keep the XLA import path aligned with the new JAX floor, shim np.bool8 for Gym 0.26 under NumPy 2, and only allocate a fresh state buffer synchronously when the background producer queue fails to hand one out.
• Code pointers:
  • third_party/pip_requirements/requirements-dev.txt: raises the JAX and NumPy version windows, with lockfile refreshes for dev and release requirements.
  • envpool/python/xla_template.py: removes the old JAX import fallback and tightens the runtime guard around legacy primitive / translation APIs.
  • envpool/__init__.py: bumps the package version to 0.9.0 and adds the NumPy 2 bool8 shim used by Gym 0.26.
  • envpool/core/state_buffer_queue.h: preserves the multi-producer background allocation path but falls back to a direct StateBuffer allocation if stock_buffer_.Get() returns nullptr.
• Notes: The dummy-env fix was validated on dev-0; I did not run a full MuJoCo sweep for this exact tree.

Test Plan

Automated

• dev-0: bazelisk test //envpool/dummy:dummy_envpool_test --config=test --test_output=errors --runs_per_test=20 --nocache_test_results: passed 20/20 after the state buffer fallback change.
• dev-0: bazelisk test //envpool/core:state_buffer_queue_test --config=test --test_output=errors --nocache_test_results: passed.
• dev-0: bazelisk test //envpool/atari:api_test //envpool/atari:atari_envpool_test //envpool/classic_control:classic_control_test //envpool/toy_text:toy_text_test --config=test --test_output=errors: passed 4/4 on the JAX 0.9 tree.

Suggested Manual

• dev-0: make bazel-test: broader regression sweep once we want to include MuJoCo coverage again.
• Install from the refreshed requirements and import envpool in a JAX 0.9 / NumPy 2 environment: confirms packaging and import-time compatibility outside Bazel.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: eea53b015c

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Replace the removed legacy JAX XLA path

With the new jax>=0.9.2 pins, any code path that calls env.xla() will now fail here because this implementation still relies on jax.core.Primitive and jax.interpreters.xla.backend_specific_translations, which are no longer present in the JAX versions this commit installs. That regresses the advertised XLA surface for users on the new dependency set, including examples/xla_step.py and the JAX-based Atari tests, from “supported” to a runtime error.

Useful? React with 👍 / 👎.