Automatic Generation of High-Performance RL Environments
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Simulation throughput | HalfCheetah JAX | Steps Per Second (SPS)798 | 6 | |
| Throughput Measurement | cartpole | Speedup Factor (vs Gymnasium)4.33e+4 | 3 | |
| GRU Rollout | Puffer Pong | SPS (Simulations/s)140 | 2 | |
| PPO training | EmuRust | Samples Per Second (SPS)14.5 | 2 | |
| PPO training | TCGJax | Steps Per Second (SPS)1.53e+5 | 2 | |
| Random action rollout | EmuRust | SPS2.39e+5 | 2 | |
| Random action rollout | PokeJAX | Simulation Speed (SPS)500 | 2 | |
| Random action rollout | TCGJax | SPS7.17e+5 | 2 | |
| GRU PPO training | Puffer Pong | Samples Per Second (SPS)35.5 | 2 | |
| PPO training | PokeJAX | Steps Per Second (SPS)15.2 | 2 |