All Tools
J
AgentsFreeOpen Source
JUMANJI
JAX environments for scalable reinforcement learning research
Apache-2.0
ABOUT
CPU-bound RL environments make large-scale sweeps slow and hard to vectorize on accelerators. Jumanji provides a JAX Env API and a suite of games plus combinatorial tasks whose difficulty can scale, so training loops can JIT compile end to end on GPU or TPU.
INSTALL
pip install jumanjiINTEGRATION GUIDE
1. Train a JAX policy on Snake, PacMan, or other game environments at hardware speed
2. Benchmark combinatorial RL on routing, packing, or job-shop style Jumanji tasks
3. Wrap Jumanji environments for Acme, RLlib, Stable-Baselines3, or Gymnasium
4. Scale environment difficulty to stress-test industrial planning agents
TAGS
pythonjaxreinforcement-learningenvironmentsinstadeepgpu