Automatic Generation of High-Performance RL Environments
Organizations: 1Princeton University · 2Independent Researcher
Abstract
Translating complex reinforcement learning (RL) environments into high-performance implementations has traditionally required months of specialized engineering. We present a closed-loop methodology that produces equivalent high-performance environments for minimal compute cost. Our method uses a generic prompt template, hierarchical verification (property, interaction, and rollout tests), iterative repair, and cross-backend policy transfer to verify no sim-to-sim gap. We demonstrate three distinct workflows across five environments: (1) Direct translation (no prior performance implementation exists) from Game Boy emulator PyBoy to our EmuRust (via Rust IPC) and from Pokemon Showdown to our PokeJAX (via JAX); (2) Translation verified against existing performance implementations via throughput parity with Puffer Pong, MJX and Brax at matched GPU batch sizes; and (3) New environment creation: TCGJax, the first Pokemon TCG Pocket environment, created from a web-extracted specification. At 200M parameters, the environment overhead drops below 4% of training time. Our closed-loop methodology confirms equivalence for all five environments. TCGJax, synthesized from a private reference absent from public repositories, serves as a contamination control for agent pretraining data concerns.
Figures & tables
| Env | Source | Target | Src LoC | Tgt LoC | Key Challenge |
|---|---|---|---|---|---|
| EmuRust | C/Python | Rust+PyO3 | 26K | 2.5k | Cycle-accurate emulation |
| PokeJAX | TypeScript | JAX | 100K | 55k | state complexity. |
| HalfCheetah | MuJoCo | JAX | 245 | 1.2k | Articulated body + contact |
| TCGJax | Web rules ⋆ | Py JAX | 29k | 4.2k | Rule extraction from web |
| Pong | C (PufferLib) | Rust+JAX | 225 | 235/318 | Already-optimized baseline |
| Environment | Benchmark | Reference (SPS) | Performance (SPS) | Speedup |
| Direct translation into newly performant environments (no prior performance implementation) | ||||
| EmuRust | Random action | 167K (PyBoy, 32p) | K (Rust, 64e) | |
| PPO training | 9.9K (PyBoy, 32p) | K (Rust, 128e) | ||
| PokeJAX | Random action | 21K (Showdown, 1p) | M (JAX, 65Kb) | |
| PPO training | 681 (Showdown) | M (JAX) | ||
| Translation verified against existing performance implementations | ||||
| Environment | Train Backend | Eval (Perf) | Eval (Ref) | Equiv. |
|---|---|---|---|---|
| Puffer Pong | C (ref) | |||
| JAX (perf) | ||||
| HalfCheetah | MJX (ref) | |||
| JAX (perf) | ||||
| EmuRust (Red) | PyBoy (ref) | |||
| Rust (perf) |
| Metric | EmuRust | PokeJAX | HalfCheetah | TCG | Pong |
|---|---|---|---|---|---|
| Target LoC | 2k | 55k | 1k | 4.2k | 235/318 |
| Modules | 5 | 30 | 5 | 11 | 1 |
| Total tests | 52 | 2k | 69 | 50 | 12 |
| Agent cost | $0.43 | $6 | $3.26 | $4.98 | $0.05 |
| Agent iterations | 72 | 63 | 20 | 51 | 13 |
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
| Environment | L1 | L2 | L3 ep. | Mode | Seeds | Xfer | Status |
|---|---|---|---|---|---|---|---|
| EmuRust | 32 | 12 | 100 | exact | 10 | ✓ | ✓ |
| PokeJAX | 1.89k | 670 | 100 | exact | 10 | ✓ | ✓ |
| HalfCheetah | 48 | 12 | 100 | ( ) | 10 | ✓ | ✓ |
| TCGJax | 20 | 24 | 100 | exact | 10 | ✓ | ✓ |
| Puffer Pong | 6 | 3 | 100 | exact | 10 | ✓ | ✓ |
| Environment | Agent | Iters | Tests | Cost |
|---|---|---|---|---|
| Pong | Gemini 3 Flash | 13 | 6/6 | $0.05 |
| Claude Sonnet 4.6 | 3 | 5/6 § | $0.08 | |
| HalfCheetah | Gemini 3 Flash | 20 | 69/69 | $3.26 |
| Claude Opus 4.6 | 6 | 69/69 | — † |
| Environment | Benchmark | PufferLib C | Rust | JAX | Speedup |
|---|---|---|---|---|---|
| Puffer Pong | Random (env only) | 60M | 122M | 275M | |
| PufferLib training | 2.4M (134K) | — | — | — | |
| GRU Rollout (2M) | 4.5M | 4.5M | 140M | ||
| GRU PPO (2M) | 854K | 855K | 35.5M |
| Environment | A6000 Ada SPS | vs. Reference |
|---|---|---|
| HalfCheetah JAX | 13.1M | (vs. Gymnasium 45K) |
| Pong JAX (scan) | 1.4B | (vs. C 60M env-only) |
| CartPole JAX (scan) | 7.1B | (vs. Gymnasium 164K) |
| Environment | Pass/Total | Stmts | Coverage | Notes |
|---|---|---|---|---|
| CartPole JAX | 9/9 | 107 | 60% | main(), rendering untested |
| Pong JAX | 5/6 | 158 | 73% | 1 statistical test ‡ |
| HalfCheetah JAX | 136/140 | 462 | 96% | 4 float32 tolerance ‡ |
| TCG Pocket JAX | 50/50 | 1k | 77% | L1/L2/L3 all passing |
| PokeJAX ∗ | 90/95 | 1k | 53% | JIT limits instrumentation |
| EmuRustGBA † | 218/218 | 527 | 98% | 110 Rust + 108 Python tests |
Reinforcement Learning
Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.
See all 1,717 papers