Examples¶
The pages below are short, runnable recipes. They mirror the scripts in examples/ and the patterns used in tests/.
Recipe map¶
| Recipe | What it shows | Underlying script |
|---|---|---|
| 01 — Single step | reset + one step on TransGridEnv |
examples/jax_02_grid_env.py |
| 02 — Batched rollout | vmap + lax.scan over many envs |
examples/jax_00_verify_device.py |
| 03 — Train PPO | single-agent PPO on BatteryEnv and TransGridEnv |
examples/train_ppo_battery.py, examples/train_ppo_transgrid.py |
| 04 — Train safe PPO | PPO-Lagrangian / CMDP on TransGridEnv |
examples/train_safe_ppo_transgrid.py |
| 05 — Train MARL | IPPO on GridMARLEnv |
examples/train_ippo_grid.py |
Notes before running¶
- Most recipes use small synthetic profiles for clarity and compile speed. Real-data paths are documented in Architecture → Data pipeline and the per-task pages under Benchmarks.
- Each recipe is self-contained: copy-paste into a Python file and run
python script.py. CPU works; GPU is faster after the first JIT compile. - For reproducible large-scale rollouts, prefer the helpers in Architecture → JAX Parallelization Architecture over hand-written Python loops.