Paper deep dive

Automatic Generation of High-Performance RL Environments

Seth Karten, Rahul Dev Appapogu, Chi Jin

Year: 2026Venue: arXiv preprintArea: cs.LGType: PreprintEmbeddings: 69

Intelligence

Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 98%

Last extracted: 3/22/2026, 6:31:58 AM

Summary

The paper introduces a reusable, agent-assisted recipe for translating complex reinforcement learning (RL) environments into high-performance implementations (JAX/Rust) for under $10. The methodology utilizes a generic prompt template, hierarchical verification (property, interaction, rollout, and cross-backend policy transfer), and iterative repair to achieve significant speedups (up to 22,320x) while maintaining semantic equivalence.

Entities (7)

EmuRust · rl-environment · 100%HalfCheetah JAX · rl-environment · 100%Hierarchical Verification · methodology · 100%PokeJAX · rl-environment · 100%Puffer Pong · rl-environment · 100%TCGJax · rl-environment · 100%Gemini 3 Flash Preview · coding-agent · 95%

Relation Signals (4)

EmuRust → implementedin → Rust

confidence 100% · EmuRust (1.5x PPO speedup via Rust parallelism for a Game Boy emulator)

PokeJAX → implementedin → JAX

confidence 100% · PokeJAX, the first GPU-parallel Pokemon battle simulator

TCGJax → implementedin → JAX

confidence 100% · TCGJax, the first deployable JAX Pokemon TCG engine

Hierarchical Verification → verifies → RL Environment

confidence 95% · Hierarchical verification (property, interaction, and rollout tests) confirms semantic equivalence for all five environments

Cypher Suggestions (0)

No Cypher suggestions yet.

Abstract

Abstract:Translating complex reinforcement learning (RL) environments into high-performance implementations has traditionally required months of specialized engineering. We present a reusable recipe - a generic prompt template, hierarchical verification, and iterative agent-assisted repair - that produces semantically equivalent high-performance environments for <$10 in compute cost. We demonstrate three distinct workflows across five environments. Direct translation (no prior performance implementation exists): EmuRust (1.5x PPO speedup via Rust parallelism for a Game Boy emulator) and PokeJAX, the first GPU-parallel Pokemon battle simulator (500M SPS random action, 15.2M SPS PPO; 22,320x over the TypeScript reference). Translation verified against existing performance implementations: throughput parity with MJX (1.04x) and 5x over Brax at matched GPU batch sizes (HalfCheetah JAX); 42x PPO (Puffer Pong). New environment creation: TCGJax, the first deployable JAX Pokemon TCG engine (717K SPS random action, 153K SPS PPO; 6.6x over the Python reference), synthesized from a web-extracted specification. At 200M parameters, the environment overhead drops below 4% of training time. Hierarchical verification (property, interaction, and rollout tests) confirms semantic equivalence for all five environments; cross-backend policy transfer confirms zero sim-to-sim gap for all five environments. TCGJax, synthesized from a private reference absent from public repositories, serves as a contamination control for agent pretraining data concerns. The paper contains sufficient detail - including representative prompts, verification methodology, and complete results - that a coding agent could reproduce the translations directly from the manuscript.

PDF

Open source PDF →Open local PDF →

Trouble viewing inline? Open PDF directly →

Full Text

68,907 characters extracted from source content.

Expand or collapse full text

Automatic Generation of High-Performance RL Environments Seth Karten 1 Rahul Dev Appapogu 2 Chi Jin 1 1 Princeton University 2 Independent Researcher Abstract. Translating complex reinforcement learning (RL) environments into high- performance implementations has traditionally required months of specialized engineering. We present a reusable recipe—a generic prompt template, hierarchical verification, and iterative agent-assisted repair—that produces semantically equivalent high-performance environments for <$10 in compute cost. We demonstrate three distinct workflows across five environments. Direct translation (no prior performance implementation exists): EmuRust (1.5×PPO speedup via Rust parallelism for a Game Boy emulator) and PokeJAX, the first GPU-parallel Pokemon battle simulator (500M SPS random action, 15.2M SPS PPO; 22,320×over the TypeScript reference). Translation verified against existing performance implementations: throughput parity with MJX (1.04×) and 5× over Brax at matched GPU batch sizes (HalfCheetah JAX); 42×PPO (Puffer Pong). New environment creation: TCGJax, the first deployable JAX Pokemon TCG engine (717K SPS random action, 153K SPS PPO; 6.6×over the Python reference), synthesized from a web-extracted specification. At 200M parameters, the environment overhead drops below 4% of training time. Hierarchical verification (property, interaction, and rollout tests) confirms semantic equivalence for all five environments; cross-backend policy transfer confirms zero sim-to-sim gap for all five environments. TCGJax, synthesized from a private reference absent from public repositories, serves as a contamination control for agent pretraining data concerns. The paper contains sufficient detail—including representative prompts, verification methodology, and complete results—that a coding agent could reproduce the translations directly from the manuscript. Date: March 13, 2026 Correspondence: sethkarten@princeton.edu 1 Introduction In typical reinforcement learning (RL) training, environment simulation consumes 50–90% of wall-clock time [11,20]. For complex simulators, such as Pokemon Showdown [3,5,6] at 100K+ lines of TypeScript, or cycle-accurate Game Boy emulators in C, this overhead is even more severe. The RL community has responded with award-winning hand-optimized rewrites: Brax [1], Gymnax [9], Pgx [7], JaxMARL [17], Craftax [14], and PureJaxRL [11]. Each required labor- intensive specialized engineering for a single domain. A method for producing performance environments cheaply and routinely, as a standard step in the RL workflow, would complement existing libraries. 1 arXiv:2603.12145v1 [cs.LG] 12 Mar 2026 Automatic Generation of High-Performance RL Environments 020406080100 % of Training Time 72% Env Typical RL Training 020406080100 % of Training Time 96% Model 4% Env With Performance Envs (Ours) Game Boy Emulator (EmuRust) 1.5× faster Pokémon Battles (PokeJAX) 22,320× faster TCG Pocket 6.6× faster HalfCheetah (MuJoCo) 1.04× vs MJX Puffer Pong 42× faster C/Python Rust | 14.5K SPSTypeScript JAX | 15.2M SPSPython JAX | 153K SPS Gymnasium JAX | 1.66M SPSC JAX | 35.5M SPS Env StepModel Fwd/BwdData Transfer Figure 1. Performance environments eliminate the environment bottleneck. (Top) Our methodology shifts training from environment-bound to model-bound. (Bottom) Five case studies, grouped by result type. Direct translation into newly performant environments (no prior performance implementation): EmuRust (1.5×CPU-to-CPU PPO); PokeJAX—the first GPU-parallel Pokemon battle simulator, 500M SPS at 65K batch. Translation verified against existing performance implementations: throughput parity with MJX (1.04×) and 5×over Brax at matched batch (HalfCheetah); 42×end-to-end PPO over expert-optimized C (Pong). New environment creation: TCGJax—the first deployable JAX Pokemon card-game engine, 717K SPS, created from a web-extracted specification. All produced for <$10 in agent compute. We show that the cost of producing high-performance RL environments has dropped by orders of magnitude. Two developments make this possible: coding agents with 1M+ token context windows, and per-token costs low enough that iterative translation costs only a few dollars. The human provides a generic translation prompt (Appendix B); the agent handles all source code generation and iterative repair for less than $10 in compute cost. We present three empirical findings: (1) modern coding agents can translate full RL envi- ronments across diverse domains including 100K+ LoC codebases with complex cross-system interactions; (2) the cost is low (<$10), orders of magnitude below what per-line extrapolation would suggest; and (3) hierarchical verification is critical—without it, agents fail to converge on complex environments (HalfCheetah ablation) and are measurably slower even on simple ones (Pong ablation); cross-backend policy transfer (L4) confirms zero sim-to-sim gap. We demonstrate this across five environments spanning discrete games, continuous physics, 2 Automatic Generation of High-Performance RL Environments hardware emulation, and multi-agent systems (Table 1). TCGJax is entirely new: no trainable RL Pokemon card-game engine existed prior to this work. PokeJAX is a direct translation of the existing Pokemon Showdown server into the first GPU-parallel Pokemon battle simulator. Key results (§4): end-to-end PPO speedups from 1.5×to 42×; throughput parity with MJX at matched batch; training curves across 10 seeds confirming policy equivalence; step-level rollout verification for all five environments; and cross-backend policy transfer confirming zero sim-to-sim gap for all five environments. We contribute: (1) empirical evidence that high-performance RL environments can be produced cheaply, validated with two additional coding agents on representative environments (Table 6); (2) five high-performance environments with complete verification suites; and (3) a reusable translation recipe with ablation evidence that hierarchical verification drives convergence. The paper contains sufficient detail—including representative prompts (Appendix B), verification methodology, and complete results—that a coding agent could reproduce the translations directly from the manuscript. 2 Related Work Hardware-accelerated environments. A growing body of work manually reimplements RL environments in JAX or on GPU. Brax [1] rewrites rigid-body physics; MJX [21] ports MuJoCo to XLA; Gymnax [9] reimplements classic control; Pgx [7] covers board games; JaxMARL [17] provides multi-agent environments; Craftax [14] reimplements Crafter (250×speedup); and PureJaxRL [11] demonstrated 4,000×speedup from end-to-end JAX compilation. Each required significant specialized engineering for a single domain. Our methodology produces five translations using the same generic prompt template for<$10 in agent compute cost, complementing these libraries by enabling fast versions of environments they do not cover. High-throughput RL systems. Gymnasium [22] standardizes the environment interface used by most RL libraries. EnvPool [23] achieves high throughput via C++ async batching; PufferLib [20] provides a unified interface for C environments; Sample Factory [15] maximizes GPU utilization. Our work is complementary: reducing per-step time lets these systems fully exploit their parallelism. Even PufferLib’s optimized C Pong achieves a 42×PPO speedup when translated to a JAX pipeline. LLM-assisted code generation. Neural code translation [8], AlphaCode [10], and SWE-bench [4] address function- or API-level tasks. Ziftci et al.[26]report∼50% effort reduction from LLM- assisted migration at Google. Eureka [13] and Text2Reward [24] use LLMs to generate reward functions. Our setting differs: we preserve exact semantics across interacting subsystems over thousands of timesteps, where silent errors corrupt training signals. Our hierarchical verification addresses this by providing structured, localized error signals. Scaling RL. Foundation RL architectures [2,16] that train across many environments amplify the cost of slow simulation, motivating scalable methods for producing performance environments. 3 Translation Recipe We translate reference RL environments into high-performance equivalents via coding agents guided by hierarchical verification and sim-to-sim gap detection after training. Figure 2 summa- rizes the pipeline. 3.1 Problem Statement For a reference environmentE ref in source programming languageL src , we produce a high performance (fast training throughput)environmentE perf in target programming languageL tgt 3 Automatic Generation of High-Performance RL Environments Reference Implementation Python / TS / C Decompose into Modules Agent-Assisted Translation + Verification L1 Property L2 Interaction L3 Rollout Individual components Cross-module interactions Full episodes fix & retry Performance Environment JAX / Rust L4: Cross-Backend Transfer Train(perf) Eval(ref) Sim-to-sim gap? gap found? add tests Rule Extraction & Scraping Web Reference TCG Pocket path (optional) Target: JAX (GPU, vmap) or Rust (CPU, Rayon) Figure 2. Translation and verification pipeline. A reference environment is decomposed into modules, translated by a coding agent, and verified through four levels of increasing scope. Failures at any level trigger targeted repair and re-verification; Level 4 cross-backend policy transfer closes the outer loop. satisfying semantic equivalence: for any seed and action sequence, both environments produce identical observations, rewards, and termination signals at every timestep. For continuous-valued environments (e.g., physics simulations), we relax this toε-equivalence: per-step outputs agree within per-componentL ∞ toleranceε, verified to produce equivalent training dynamics across seeds (see §4.4 for environment-specific tolerances). We additionally verify cross-backend policy equivalence (Level 4): a policy trained inE perf should achieve statistically indistinguishable reward when evaluated inE ref , confirming that there is no sim-to-sim gap. These are empirical behavioral equivalence properties verified over tested inputs (100 episodes, diverse RNG paths), not formal semantic equivalence over all possible inputs. Formal guarantees (e.g., bisimulation or bounded-error compilation verification) would require reasoning about all reachable states, which is intractable for the environments in this study. Instead, we rely on layered testing to provide high confidence without formal proof. We additionally require thatE perf achieves sufficient throughput to minimize environment overhead relative to training time. We select between JAX and Rust based on environment structure: JAX compiles pure- function environments to GPU via XLA (vmap,jax.lax.scan); Rust suits stateful, memory- intensive environments with CPU parallelism. See Appendix A for detailed selection criteria. 3.2 Hierarchical Verification Checking semantic equivalence via exhaustive rollout comparison alone is insufficient: when a discrepancy is detected, localizing the root cause in a large codebase is intractable. We decompose verification into four levels forming a closed feedback loop: failures at any level trigger targeted repair and re-verification at lower levels, and Level 4 cross-backend policy transfer closes the loop by feeding back into earlier stages when a sim-to-sim gap is detected. Level 1: Property tests (L1) verify individual components in isolation by asserting input- output pairs fromE ref matchE perf . Level 2: Interaction tests (L2) verify cross-module state dependencies and event ordering by exercising multi-subsystem operation sequences. Level 3: Rollout comparison (L3) executes full episodes in both environments under matched seeds and identical action sequences, comparing all outputs at every timestep. Level 4: Cross-backend policy transfer (L4). A policy trained inE perf is evaluated inE ref (and vice versa), testing the environment under the stochastic state distribution induced by a learned policy. Each level catches a distinct bug class: L1 catches arithmetic/boundary errors, L2 catches 4 Automatic Generation of High-Performance RL Environments Table 1. Environment overview. ⋆ Private reference (contamination control). EnvSourceTargetSrc LoC Tgt LoC Key Challenge EmuRustC/PythonRust+PyO3 ∼26K2,511 Cycle-accurate emulation PokeJAXTypeScriptJAX∼100K55,629 2,834 species, 1,370 moves HalfCheetah MuJoCoJAX2451,202 Articulated body + contact TCGJaxWeb rules ⋆ Py→JAX29,5264,235 Rule extraction from web PongC (PufferLib) Rust+JAX225 235/318 Already-optimized baseline ordering/propagation errors, L3 catches accumulating drift and reset logic errors, and L4 catches any sim-to-sim gap affecting policy quality. Failures at any level trigger repair; L4 feeds sim-to- sim gaps back into targeted L1/L2 tests. The iterative cycle, not any single verification pass, drives convergence to a correct translation. 3.3 Agent-Assisted Translation Process Translation proceeds in a closed verification loop across four phases: (1) Module translation: the agent translates each module independently (ordered by dependency) and verifies with Level 1 tests before proceeding. (2) Integration: Level 2 tests verify composed modules; failures trigger repair while preserving Level 1 correctness. (3) Validation: Level 3 rollout comparison provides end-to-end verification; discrepancies trigger root-cause analysis with new targeted lower-level tests. (4) Cross-backend validation: a policy trained inE perf is evaluated inE ref (Level 4); a detected sim-to-sim gap feeds back into phases (1)–(3) until the gap closes. If the agent fails to make progress afterT =50 iterations at any level, human intervention adds targeted tests or refines the prompt. All translations used Gemini 3 Flash Preview, invoked via the Gemini CLI in non-interactive mode (gemini –yolo); however, the methodology is agent-agnostic. The agent receives module source code, target language specification, and test requirements in a single prompt (see §4.4 for measured costs). Human involvement is limited to writing translation prompts, specifying module decomposition and target architecture, and designing verification test structures. Appendix B provides representative prompt types, and Appendix C provides backend-specific optimization checklists. Algorithm 1 in Appendix A.9 formalizes this process. 4 Experiments We test whether coding agents can translate diverse RL environments into semantically equiva- lent high-performance implementations, and whether hierarchical verification is necessary for convergence. Results confirm both across five environments spanning discrete games, continuous physics, and multi-agent systems. We apply our methodology to five environments (Table 1). All benchmarks use 1×RTX 5090, 32 AMD Ryzen cores, CUDA 12.8, JAX 0.4.39. Training curves useN =10 seeds with matched PPO [19] hyperparameters. Additional details in Appendix A.12. 4.1 Throughput Results Table 2 consolidates throughput for all five environments. Results span three categories. Direct translations (EmuRust, PokeJAX) produce newly performant versions where none existed; PokeJAX’s 23,810× speedup reflects a paradigm shift from sequential CPU server to GPU-parallel pure functions, enabling convergent training previ- ously impractical at 681 SPS. Verified translations (Pong, HalfCheetah) achieve speedups over already-optimized baselines: Pong achieves 42×PPO via scan-fused GPU training; HalfCheetah reaches throughput parity with Google’s MJX (1.04×), demonstrating that agent-generated code matches hand-optimized engines. New environment creation (TCGJax) translates a web- extracted specification into a trainable JAX environment. Per-environment details are in Appendix A.1. 5 Automatic Generation of High-Performance RL Environments Table 2. Throughput comparison. Mean±std fromN =5 runs (CVs<3%);∼2M models; JAX excludes JIT warm-up. EnvironmentBenchmarkReference (SPS)Performance (SPS) Speedup Direct translation into newly performant environments (no prior performance implementation) EmuRust Random action167K (PyBoy, 32p)239±6K (Rust, 64e)1.4× PPO training9.9K (PyBoy, 32p)14.5±0.4K (Rust, 128e)1.5× PokeJAX Random action21K (Showdown, 1p)500±9M (JAX, 65Kb)23,810× PPO training681 (Showdown)15.2±0.2M (JAX)22,320× Translation verified against existing performance implementations Puffer Pong GRU Rollout (2M)4.5±0.008M (C, CPU)140±1.5M (JAX, GPU)31× GRU PPO (2M)854±4K (C, CPU)35.5±0.3M (JAX, GPU)42× HalfCheetah JAX vs Gymnasium45K (1 proc)1.66M (JAX, 32Kb)37× vs Brax160K (Brax, 4Kb)798K (JAX, 4Kb)5.0× vs MJX (Google)1.6M (MJX, 32Kb)1.66M (JAX, 32Kb)1.04× New environment creation (no prior trainable RL env) TCGJax Random action140K (Python, 16p)717±0.6K (JAX, 16Kb)5.1× PPO training23K (Python, 16p)153±5K (JAX, 4Kb)6.6× p = processes, e = env instances, b/Kb = JAX batch (thousands). EmuRust comparison at matched 32 CPU cores (Appendix A.13). 0 20 40 60 80 100 % of PPO Training Time 51% 71% 35% 25% 39% 34% 2M2M2M2M2M2M 30% 71% 26% 9% 11% 18% 20M20M20M20M20M20M 12% 26% 7% 2.9% 1.1% 3% 200M200M200M200M200M200M PyBoy+SB3 (ref) Showdown (ref) TCG Pocket (ref) EmuRust (perf) PokeJAX (perf) TCGJax (perf) ref | perf PPO Training Breakdown Across Model Scales Env StepData TransferModel Fwd/Bwd Figure 3. PPO training time breakdown across model scales. Three bars per implementation show 2M, 20M, 200M parameter models. Performance implementations drop to≤4% env overhead at 200M. All on 1× RTX 5090. 4.2 Training Time Breakdown Figure 3 profiles PPO iteration time across model scales (2M, 20M, 200M parameters). At 200M, all single-agent performance implementations contribute≤4% of training time (down from 50–90% for references). 4.3 Policy Equivalence All five environments pass L3 rollout comparison (100 episodes, matched RNG seeds, step-level output comparison; exact for discrete envs,ε=10 −3 for HalfCheetah). Figure 4 shows matched training curves: Pong (10 seeds), HalfCheetah (10 seeds), and EmuRust (10 seeds) confirm 6 Automatic Generation of High-Performance RL Environments Table 3. Cross-backend policy transfer. Values are mean±std over 10 seeds. Equivalence confirmed via TOST (α=0.05) with environment-specific margins (∆): Pong ∆=1.0, HalfCheetah ∆=100, EmuRust ∆=0.5, PokeJAX ∆=0.02, TCGJax ∆=0.05. PokeJAX and TCGJax report win rate against a heuristic bot; others report episode return. EnvironmentTrain BackendEval (Perf)Eval (Ref) Equiv. Puffer Pong C (ref)28.01± 0.2828.04± 0.29✓ JAX (perf)28.23± 0.1828.22± 0.20✓ HalfCheetah MJX (ref)1398± 4971389± 511✓ JAX (perf)1026± 6361133± 562✓ EmuRust (Red) PyBoy (ref)12.01± 0.1211.99± 0.15✓ Rust (perf)12.06± 0.0012.06± 0.01✓ PokeJAX Showdown (ref)0.313± 0.0070.313± 0.007✓ JAX (perf)0.406± 0.0030.406± 0.003✓ TCGJax Python (ref)0.575± 0.0540.543± 0.045✓ JAX (perf)0.583± 0.0620.558± 0.042✓ Table 4. Translation cost. Costs include all iterations; base rates from Gemini 3 Flash Preview logs for environments requiring multiple revision cycles. MetricEmuRust PokeJAX HalfCheetah TCGPong Target LoC2,51155,6291,202 4,235 235/318 Modules5305111 Total tests522,783695012 Agent cost$0.43$6$3.26 $4.98$0.05 Agent iterations7263205113 consistent learning dynamics across backends. All five environments achieve L4 cross-backend policy transfer (Table 3). Details in Appendix A.2. Cross-backend policy transfer (L4). Table 3 evaluates policies trained in one backend on both backends (10 seeds each). We use the TOST (Two One-Sided Tests) equivalence procedure [18] with environment-specific margins ∆ (caption of Table 3); a significant TOST result (p<0.05) confirms that the two backends produce equivalent performance within±∆. All five environments pass: Pong shows zero sim-to-sim gap, HalfCheetah confirms equivalent transfer despite high variance, and EmuRust-trained Pokemon Red policies transfer to PyBoy with near-identical reward. PokeJAX achieves exact transfer: win rates are bit-identical across backends. TCGJax likewise confirms equivalence in both directions. 4.4 Translation Effort and Verification Table 4 summarizes translation cost. All environment code is agent-generated; no lines were written by hand. Costs include all translation iterations (e.g., HalfCheetah required four solver revisions; EmuRust required three fix cycles). Verification scope is summarized in Table 5 (Appendix A.3); all five environments pass all levels. Hierarchical verification is critical: on HalfCheetah, L3-only verification failed to converge in 42 iterations, while the full hierarchy converged in 5 (Appendix A.5). The methodology is agent- agnostic: re-translating Pong with Claude Sonnet 4.6 and HalfCheetah with Claude Opus 4.6 produces functionally correct translations using identical prompts (Table 6 in Appendix A.4). 7 Automatic Generation of High-Performance RL Environments 010002000 Gradient steps 0 500 1000 1500 Episode reward HalfCheetah Gymnasium JAX 0200400600 Gradient steps 20 10 0 10 20 30 Episode reward Pong C PufferLib JAX 02000040000 Gradient steps 0 1 2 3 4 Gym badges Pokemon Red PyBoy (mean) EmuRust (mean) EmuRust (max) 02000400060008000 Gradient steps 1600 1800 2000 2200 Elo TCG Pocket JAX Python Minimax Random 05001000 Gradient steps 1400 1450 1500 1550 Elo PokeJAX PokeJAX Showdown Random Kaizo Bot Figure 4. Policy equivalence. Pong (10 seeds), HalfCheetah (10 seeds), EmuRust (10 seeds): matched reward curves across backends. TCGJax and PokeJAX: matched Elo curves (JAX vs reference). All five environments achieve L4 cross-backend transfer (Table 3). 5 Conclusion Coding agents guided by a closed-loop hierarchical verification recipe translate reference RL environments into semantically equivalent high-performance implementations for <$10. Across three workflows—direct translation (EmuRust, PokeJAX), translation verified against exist- ing performance implementations (Puffer Pong, HalfCheetah), and new environment creation (TCGJax)—results include throughput parity with MJX and 5×over Brax at matched batch sizes, 1.5–42×end-to-end PPO speedups, and training enablement for environments previously too slow to train. Four-level verification (L1–L4) confirms semantic equivalence; cross-backend policy transfer (L4) confirms zero sim-to-sim gap, with failures feeding back into targeted L1/L2 repair. The methodology is agent-agnostic, and the hierarchical test structure is essential: without L1/L2 tests, agents fail to converge on complex physics. The approach is most effective for environments with reproducible transitions, clear module boundaries, and fixed-size state rep- resentations; environments with non-deterministic external dependencies or unbounded dynamic allocation may require additional engineering beyond what the generic recipe provides. The methodology decouples environment complexity from training cost: researchers can produce performance versions of the environments they need, rather than being limited to existing JAX ports. Re-translating when a reference updates costs under $1, with the test suite serving as a regression guard. As coding agents improve and per-token costs fall, fast verified simulation becomes a default step in the RL workflow rather than a bottleneck requiring months of specialized engineering—closing the gap between the environments researchers want to study and the environments they can afford to train on. 8 Automatic Generation of High-Performance RL Environments Acknowledgements This work was supported by the National Science Foundation Graduate Research Fellowship under Grant No. DGE-2444107. References [1]C. D. Freeman, E. Frey, A. Raichuk, S. Girgin, I. Mordatch, and O. Bachem. Brax–a differen- tiable physics engine for large scale rigid body simulation. arXiv preprint arXiv:2106.13281, 2021. 1, 2 [2]J. Grigsby, L. Fan, and Y. Zhu. Amago: Scalable in-context reinforcement learning for adaptive agents. arXiv preprint arXiv:2310.09971, 2023. 2 [3]J. Grigsby, Y. Xie, J. Sasek, S. Zheng, and Y. Zhu. Human-level competitive pokémon via scalable offline reinforcement learning with transformers. In Reinforcement Learning Conference (RLC), 2025. arXiv:2504.04395. 1, A.1 [4]C. E. Jimenez, J. Yang, A. Wettig, S. Yao, K. Pei, O. Press, and K. Narasimhan. Swe-bench: Can language models resolve real-world github issues? arXiv preprint arXiv:2310.06770, 2023. 2 [5]S. Karten, J. Grigsby, S. Milani, K. Vodrahalli, A. Zhang, F. Fang, Y. Zhu, and C. Jin. The pokéagent challenge: Competitive and long-context learning at scale. NeurIPS Competition Track, 2025. 1, A.1 [6]S. Karten, A. L. Nguyen, and C. Jin. Pokéchamp: an expert-level minimax language agent. arXiv preprint arXiv:2503.04094, 2025. 1, A.1 [7]S. Koyamada, S. Okano, S. Nishimori, Y. Murata, K. Habara, H. Kita, and S. Ishii. Pgx: Hardware-accelerated parallel game simulators for reinforcement learning. Advances in Neural Information Processing Systems, 36:45716–45743, 2023. 1, 2 [8]M.-A. Lachaux, B. Roziere, L. Chanussot, and G. Lample. Unsupervised translation of programming languages. arXiv preprint arXiv:2006.03511, 2020. 2 [9]R. T. Lange. gymnax: A jax-based reinforcement learning environment library. Version 0.0, 4, 2022. 1, 2, A.20 [10]Y. Li, D. Choi, J. Chung, N. Kushman, J. Schrittwieser, R. Leblond, T. Eccles, J. Keeling, F. Gimeno, A. Dal Lago, et al. Competition-level code generation with alphacode. Science, 378(6624):1092–1097, 2022. 2 [11]C. Lu, J. Kuba, A. Letcher, L. Metz, C. Schroeder de Witt, and J. Foerster. Discovered policy optimisation. Advances in Neural Information Processing Systems, 35:16455–16468, 2022. 1, 1, 2 [12]G. Luo. Pokémon showdown.https://github.com/smogon/pokemon-showdown, 2011. A.1 [13]Y. J. Ma, W. Liang, G. Wang, D.-A. Huang, O. Bastani, D. Jayaraman, Y. Zhu, L. Fan, and A. Anandkumar. Eureka: Human-level reward design via coding large language models. arXiv preprint arXiv:2310.12931, 2023. 2 [14]M. Matthews, M. Beukman, B. Ellis, M. Samvelyan, M. Jackson, S. Coward, and J. Foerster. Craftax: A lightning-fast benchmark for open-ended reinforcement learning. arXiv preprint arXiv:2402.16801, 2024. 1, 2 9 Automatic Generation of High-Performance RL Environments [15]A. Petrenko, Z. Huang, T. Kumar, G. Sukhatme, and V. Koltun. Sample factory: Ego- centric 3d control from pixels at 100000 fps with asynchronous reinforcement learning. In International Conference on Machine Learning, pages 7652–7662. PMLR, 2020. 2 [16]S. Reed, K. Zolna, E. Parisotto, S. G. Colmenarejo, A. Novikov, G. Barth-Maron, M. Gimenez, Y. Sulsky, J. Kay, J. T. Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.06175, 2022. 2 [17] A. Rutherford, B. Ellis, M. Gallici, J. Cook, A. Lupu, G. Ingvarsson Juto, T. Willi, R. Hammond, A. Khan, C. Schroeder de Witt, et al. Jaxmarl: Multi-agent rl environments and algorithms in jax. Advances in Neural Information Processing Systems, 37:50925–50951, 2024. 1, 2 [18]D. J. Schuirmann. A comparison of the two one-sided tests procedure and the power approach for assessing the equivalence of average bioavailability. Journal of pharmacokinetics and biopharmaceutics, 15(6):657–680, 1987. 4.3 [19]J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy opti- mization algorithms. arXiv preprint arXiv:1707.06347, 2017. 4 [20] J. Suarez. Pufferlib: Making reinforcement learning libraries and environments play nice. arXiv preprint arXiv:2406.12905, 2024. 1, 2, A.1 [21]E. Todorov, T. Erez, and Y. Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, pages 5026–5033. IEEE, 2012. 2 [22] M. Towers, A. Kwiatkowski, J. Terry, J. U. Balis, G. De Cola, T. Deleu, M. Goulão, A. Kallinteris, M. Krimmel, A. KG, et al. Gymnasium: A standard interface for reinforcement learning environments. arXiv preprint arXiv:2407.17032, 2024. 2 [23]J. Weng, M. Lin, S. Huang, B. Liu, D. Makoviichuk, V. Makoviychuk, Z. Liu, Y. Song, T. Luo, Y. Jiang, et al. Envpool: A highly parallel reinforcement learning environment execution engine. Advances in Neural Information Processing Systems, 35:22409–22421, 2022. 2 [24]T. Xie, S. Zhao, C. H. Wu, Y. Liu, Q. Luo, V. Zhong, Y. Yang, and T. Yu. Text2reward: Reward shaping with language models for reinforcement learning. arXiv preprint arXiv:2309.11489, 2023. 2 [25]M. Ynddal. Pyboy: Game boy emulator written in python.https://github.com/ Baekalfen/PyBoy, 2018. A.1 [26]C. Ziftci, S. Nikolov, A. Sjövall, B. Kim, D. Codecasa, and M. Kim. Migrating code at scale with llms at google. In Proceedings of the 33rd ACM International Conference on the Foundations of Software Engineering, pages 162–173, 2025. 2 10 Automatic Generation of High-Performance RL Environments A Supplementary Details This appendix provides additional tables, figures, and detailed descriptions that support the main text. A.1 Per-Environment Details The following paragraphs provide detailed descriptions for each environment, complementing the summary in §4.1. EmuRust (C/Python→Rust). The Game Boy emulator decomposes into five modules (CPU, memory, PPU, core, bindings). Both reference and translation run on CPU: PyBoy [25] uses Python multiprocessing (one process per instance), while EmuRust uses Rayon’s work-stealing thread pool within a single process. The 1.5×comparison is at matched CPU resources: both backends use the same 32 cores, but PyBoy saturates at 32 processes (one per core) while EmuRust packs 128 environments into a single process via Rayon’s shared-memory thread pool, achieving higher per-core utilization through cooperative scheduling with zero IPC overhead (Table 8). PokeJAX (TypeScript→JAX). PokeJAX is the first GPU-parallel Pokemon battle simulator. The standard tool for RL researchers was Pokemon Showdown [12], a TypeScript server designed for human online play that has since become the primary testbed for competitive Pokemon AI [3,5,6], though not originally designed for RL training. Translating it (100K+ lines) required server/client flattening, fixed-size state arrays, and branch-parallel dispatch viajax.lax.switch. The full 55,629-line translation is complete across∼30 modules; only minor rule-edge-case modules are excluded. The∼$6 cost in Table 4 is extrapolated from a 5-module subset for which session-level cost logs were available. The reference baseline (21K SPS) reflects a single-threaded server not designed for throughput; running multiple instances via PokeEnv yields only 681 SPS due to WebSocket overhead. The 23,810×is an enabling number—without this translation, training a Pokemon battling agent is impractical (>4 days for basic curriculum learning vs. 15 minutes with PokeJAX). The speedup decomposes into JAX compilation + GPU batching at 1K instances (560×) and batch scaling from 1K to 65K (42.5×), reflecting an architectural change (sequential CPU server to GPU-parallel pure functions), not per-instruction optimization. The 1,370 move functions dispatched vialax.switchproduce a large XLA HLO graph, reflected in the 45 s JIT time; every step pays the cost of all 1,370 move computations regardless of which move is used—a known overhead of branchless GPU execution. Verification comprises 2,783 tests across all three levels; 68% of bugs were caught by L1, 24% by L2, and 8% by L3. HalfCheetah JAX (Gymnasium/MuJoCo→JAX). The hardest translation: MuJoCo’s HalfChee- tah requires articulated-body dynamics (9 DOFs, 7 rigid bodies, 6 actuators) with ground contact. The agent translated forward kinematics, the Composite Rigid Body Algorithm for mass ma- trices, analytical RNEA for bias forces, and contact Jacobians, all as pure JAX (1,202 lines, 5 modules). Total translation cost $3.26 across four solver revisions (penalty-spring, PGS, Jacobi, Newton/LCP), with all 69 tests passing. At matched batch size (32,768), our translation achieves throughput parity with MJX (1.66M vs. 1.6M SPS) and 5×over Brax at batch 4,096. Both our translation and MJX use the same Newton contact solver formulation (acceleration-space QP with Cholesky factorization) and float32 precision; the throughput parity demonstrates that agent-generated, environment-specific code matches the performance of Google’s hand-optimized general-purpose engine. The 37×speedup over Gymnasium’s single-process CPU execution remains the practically relevant number for training workflows. 11 Automatic Generation of High-Performance RL Environments TCGJax (Web rules→Python→JAX). TCG Pocket demonstrates specification-to-implementation translation. We extracted rules from official web sources, built a Python reference (29,526 lines), then translated to JAX (4,235 lines). The entire translation cost $4.98 across 11 modules (including an early attempt that erroneously applied rules from a different trading card game; L1 tests and rule verification caught the errors, and additional iterations corrected them). TCG Pocket serves as a contamination control: the Python reference is private (no public repository), so the agent cannot rely on pretraining memorization. The Python reference at 23K SPS (16 processes) is too slow for practical training; JAX at 153K SPS (batch 4K) converges to reward 1.0 in ∼12 minutes. Puffer Pong (C→Rust + JAX). PufferLib’s [20] C Pong is already optimized (60M SPS random). Translating to JAX enablesjax.lax.scan-fused rollouts where the entire rollout compiles into a single GPU kernel with zero CPU↔GPU transfer. C environments cannot exploit this fusion. The 42×PPO speedup reflects the CPU-to-GPU architectural change, not like-for-like optimization. This is the core argument for JAX as a target language. A.2 Detailed Policy Equivalence Discussion Our verification hierarchy comprises four levels, each providing progressively stronger evidence of semantic equivalence. L1–L2 (property and interaction tests) verify individual functions and module interactions. L3 verifies that the environment transition function is identical (discrete) orε-close (continuous) for 100 tested action sequences under matched RNG seeds. Training curves provide complementary evidence under stochastic exploration. Neither L3 nor training curves alone is sufficient: L3 cannot cover all possible action sequences, and training curves cannot isolate which step introduced an error. Together, they provide strong evidence. We note that overlapping±1σerror bands across seeds is a necessary but not sufficient condition for formal statistical equivalence. L4: Cross-backend policy transfer. L4 is the strongest verification level: a policy trained entirely in one backend is evaluated in the other (Table 3). Unlike L3, which tests the environment under scripted action sequences, L4 exercises the environment under the stochastic state distribution induced by a learned policy—states the agent actually visits during training, which may differ substantially from those reached by random or scripted actions. This makes L4 sensitive to subtle semantic differences that L3 may miss. The results confirm zero sim-to-sim gap for all five environments. We assess equivalence using the TOST (Two One-Sided Tests) equivalence procedure with pre-specified equivalence margins ∆ (Table 3 caption); unlike a standardt-test (which tests whether means differ), TOST tests whether the difference falls within±∆, making it the appropriate test for equivalence claims. A significant TOST result (p<0.05) rejects the null hypothesis of non-equivalence. For Puffer Pong (∆=1.0), L4 confirms a C-trained policy achieves 28.01±0.28 on C and 28.04±0.29 on JAX, with equivalence confirmed at α=0.05. For HalfCheetah (∆=100), the JAX translation replicates MJX’s complete physics pipeline: Newton contact solver (acceleration-space QP with SOLIMP impedance, pyramidal friction cone, and Cholesky factorization), joint limit constraints (with per-DOF impedance), and implicit Euler integration. Cross-backend evaluation confirms equivalence: MJX-trained policies retain 101% on JAX (1398±497 vs. 1389±511) and JAX-trained policies retain 110% on MJX (1133± 562 vs. 1026± 636), both confirmed equivalent via TOST at α=0.05. For EmuRust (∆=0.5), Rust-trained Pokemon Red policies transfer to PyBoy with near- identical reward (12.06±0.00 vs. 12.06±0.01), and PyBoy-trained policies likewise transfer to EmuRust (12.01±0.12 vs. 11.99±0.15), both confirmed equivalent via TOST atα=0.05. This confirms pixel-level fidelity of the emulator translation despite the complexity of full Game Boy hardware emulation (CPU, PPU, memory, interrupts, timers). 12 Automatic Generation of High-Performance RL Environments For PokeJAX (∆=0.02), cross-backend transfer is exact: JAX-trained policies achieve identical win rates when evaluated on Showdown (0.406±0.003 in both directions), and Showdown-trained policies likewise transfer perfectly to JAX (0.313±0.007). The bit-identical results trivially satisfy TOST and reflect the deterministic nature of the battle simulator—given the same RNG seed and action sequence, both backends produce identical game outcomes. For TCGJax (∆=0.05), cross-backend transfer confirms equivalence in both directions: JAX- trained policies achieve 0.583±0.062 win rate on JAX and 0.558±0.042 on the Python reference, while Python-trained policies achieve 0.575±0.054 on Python and 0.543±0.045 on JAX, both confirmed equivalent via TOST atα=0.05. This confirms faithful translation despite the card game’s complex branching logic (1,000+ card effects dispatched via lax.switch). Per-environment details. For HalfCheetah, the per-step kinematics tolerance isε=10 −3 (Ta- ble 5); the accumulated dynamics tolerance across 5 frame-skip substeps reaches 1.0 in absolute terms but this reflects Euler integration error compounding, not semantic translation errors. Gymnasium and JAX training curves converge to comparable rewards on their respective back- ends (Figure 4); cross-backend policy transfer (JAX vs. MJX) confirms equivalent performance via TOST (∆=100,α=0.05): MJX-trained policies retain 101% on JAX and JAX-trained policies retain 110% on MJX, demonstrating zero sim-to-sim gap. Beyond cross-backend equivalence, two environments also demonstrate training enablement— where reference implementations are too slow for practical training. TCGJax (a new environment creation from a web-extracted specification) requires∼65M environment steps to converge; at the Python reference’s 23K SPS, the actual training loop extends to several hours, during which training instabilities compound. JAX converges in∼12 minutes. PokeJAX similarly enables practical training: a GRU PPO agent trained at 145K SPS with curriculum learning across 4 heuristic opponents completes all 4 stages in under 15 minutes vs. over 4 days at Showdown’s 681 SPS. Both environments achieve L4 cross-backend transfer (Table 3): PokeJAX with bit-identical win rates and TCGJax with equivalence confirmed via TOST in both directions. A.3 Verification Summary Table 5. Verification summary. L1 = property tests, L2 = interaction tests, L3 = rollout comparison, Xfer = cross-backend policy transfer. EnvironmentL1 L2 L3 ep. ModeSeeds Xfer Status EmuRust32 12100 exact10✓ PokeJAX1,890 670100 exact10✓ HalfCheetah48 12100 ε (10 −3 ) 10✓ TCGJax20 24100 exact10✓ Puffer Pong63100 exact10✓ A.4 Multi-Agent Validation We re-translated Pong with Claude Sonnet 4.6 and HalfCheetah with Claude Opus 4.6, using identical prompts and test suites. Both agents converge to functionally correct translations (Table 6), confirming the methodology is agent-agnostic. A.5 Verification Ablation Details HalfCheetah (6-DOF, complex physics). The L3-only run used 8 end-to-end tests and consumed 42 agent iterations over 35 minutes ($0.17) without converging. The agent could not isolate dynamics bugs (Coriolis force sign errors, contact Jacobian issues) from end-to-end rollout failures, cycling through vectorization rewrites and stability patches. In contrast, the hierarchical translation converged in 5 iterations ($0.82, all 69 tests passing), 8.4×faster in iteration count. 13 Automatic Generation of High-Performance RL Environments Table 6. Multi-agent comparison. Identical inputs; functionally equivalent outputs. Environment AgentIters TestsCost PongGemini 3 Flash136/6 $0.05 Claude Sonnet 4.63 5/6 § ∼$0.08 HalfCheetah Gemini 3 Flash20 69/69 $3.26 Claude Opus 4.66 69/69— † § Same statistical test applied to both agents. † Cost not separately tracked. The L3-only agent’s failure mode—code that passes shape and API tests but produces unstable dynamics—is precisely what L1 property tests catch immediately (e.g., mass matrix symmetry, bias force magnitude bounds). Pong (simple game logic). The L3-only run converged in 15 iterations over 8.4 minutes ($0.047). The hierarchical translation converged in 13 iterations over 3.5 minutes ($0.050) with all 6 tests passing. L3-only succeeded but required 15% more iterations and 2.4×longer wall-clock time due to reliance on coarse statistical feedback rather than fine-grained L1 signals. Two data points spanning simple logic (Pong) and moderate physics (HalfCheetah, 6-DOF) consistently show that L3-only fails when contact dynamics and multi-body kinematics chains are involved. The complexity threshold appears to lie between simple game logic and rigid-body physics with ≥6 degrees of freedom. A.6 Test Adequacy Test adequacy is supported by three complementary signals: (1) L1 tests target every exported function’s boundary conditions; measured line coverage ranges from 60% (CartPole JAX) to 98% (EmuRustGBA, 110 Rust unit tests), with physics modules at 86–96% and TCG Pocket JAX at 77% (Table 11); (2) L3 rollout comparison exercises the full composed system under diverse RNG paths (100 episodes); and (3) training curves test the environment under stochastic exploration of a learned policy. Branch coverage is not measured; rare-event paths may be undertested by 100 L3 episodes, and distributional testing over larger episode counts is future work. A.7 Target Language Selection Criteria Table 7 summarizes the criteria used to select between JAX and Rust (§3.1). Table 7. Target language selection criteria. PropertyJAXRust Branching Many conditionals (lax.switch) Sequential, data-dependent State repr. Fixed-size arraysVariable-size, pointer-based Parallelism GPU SIMD (vmap)CPU threads (Rayon) Best forTurn-based/card gamesHardware emulation A.8 Methodology Details Stopping criteria. Each verification level has explicit completion criteria. Level 1 requires all property tests to pass with 100% module coverage. Level 2 requires all interaction test scenarios to pass. Level 3 requires full rollout comparison to match forN =100 episodes under controlled RNG. We choseN =100 because: (1) 100 episodes cover all primary game mechanics under diverse RNG paths, (2) coverage plateaus (in PokeJAX, no new bug class was discovered after episode 47), and (3) exact step-level comparison within each episode is strictly stronger than statistical comparison. 14 Automatic Generation of High-Performance RL Environments Module decomposition. We decompose along natural abstraction boundaries: each module should have a clear interface and minimal coupling. For game environments, natural modules include: core state transitions, entity logic, observation generation, reward computation, and I/O bindings. Smaller modules (100–500 lines) translate more reliably. Coding agent specification. All translations used Gemini 3 Flash Preview via the Gemini CLI in non-interactive mode (gemini –yolo). Human involvement is limited to writing translation prompts and designing verification test structures; all code is agent-generated. A.9 Translation Algorithm Algorithm 1 formalizes the closed-loop translation process described in §3.3. A.10 Scope and Limitations Two translations experienced significant difficulty: PokeJAX required 63 agent iterations (5- module subset), and the HalfCheetah L3-only ablation failed to converge. Several environment classes challenge or break the methodology: non-reproducible environments (race conditions, async I/O) break L3 verification; external dependencies (databases, APIs, hardware-in-the-loop) cannot be fully captured; very large codebases (>100K LoC) strain agent context windows; and private codebases not in LLM pretraining data may require more iterations, though verification ensures correctness regardless. Speedup magnitude varies widely: from 1.5× to 23,810×. A.11 Extended Practical Guidance PokeJAX as boundary case. PokeJAX (55,629 lines) represents the methodology’s boundary in codebase complexity. The human input consists of filling in the generic prompt template (Appendix B) with module-specific source code paths and interface contracts—a single specifica- tion prompt using the same generic template as all other environments. No environment code was written by hand. The boundary is the agent’s iteration count (63 iterations for a 5-module subset), not human effort. Environment speed vs. sample efficiency. Model-based methods and offline RL reduce sample requirements, partially alleviating the environment bottleneck. However, fast environments remain critical for on-policy methods requiring billions of samples, foundation RL systems training across many environments, and training enablement where the reference is too slow for any algorithm. Code quality and maintenance. Agent-generated code passes all verification tests but varies in readability. We did not hand-edit any generated code post-verification. Long-term maintainability remains open, though the test suite provides a safety net. When the reference updates, re- translating costs under $1 with the test suite as regression guard. Framework compatibility. JAX environments expose a standardstep(state, action) -> (state, reward, done)interface compatible with PureJaxRL-style scan-fused training. They can also be wrapped for Gymnasium-based frameworks via a NumPy bridge. Rust translations expose a PufferLib-compatible Gymnasium interface via PyO3. Reproducibility. Appendix B provides representative prompts with sufficient detail to reproduce the translations. The generic template structure stays constant across all environments; only the module source code, target constraints, and interface contracts vary. 15 Automatic Generation of High-Performance RL Environments A.12 Experimental Details Throughput measurement. All JAX benchmarks exclude one-time JIT compilation from steady- state timing (warm-up call before measurement). JIT compilation ranges from∼3 s (HalfCheetah) to∼45 s (PokeJAX); for a 10-minute training run, this amortizes to<1% for all environments. For PokeJAX (45 s JIT), amortization over a typical 30-minute run adds ∼2.5%. GPU memory. HalfCheetah∼4 GB (65K batch), Pong∼2 GB, PokeJAX∼28 GB (65K), TCGJax ∼8 GB (16K). Training hyperparameters. Learning rate 2.5×10 −4 , clip ratio 0.2, 4 epochs, GAEλ= 0.95, γ = 0.99, with environment-specific batch sizes matched between backends. A.13 EmuRust Scaling Ablation Table 8 reveals the EmuRust scaling advantage. The fair comparison is at matched CPU resources: PyBoy peaks at 32 processes (one per core, its architectural limit), while EmuRust scales to 128+ environments on the same 32 cores via Rayon’s work-stealing thread pool, achieving 1.5× higher throughput through efficient shared-memory parallelism with zero IPC overhead. Table 8. EmuRust scaling ablation. PPO training SPS at various environment counts. Backend8 env16 env32 env64 env EmuRust— 10,263 ± 134 13,128 ± 386 PyBoy4,197 ± 118 6,236 ± 356 9,852 ± 1,268— Backend128 env256 env EmuRust 14,482 ± 40 14,387 ± 904 A.14 PufferLib Detailed Comparisons Table 9. PufferLib comparisons. “PufferLib training” reports their full pipeline; matched rows use∼2M GRU. All on 1× RTX 5090. Environment BenchmarkPufferLib C Rust JAX Speedup Puffer Pong Random (env only)60M 122M 275M4.6× PufferLib training 2.4M (134K)— GRU Rollout (2M)4.5M 4.5M 140M31× GRU PPO (2M)854K 855K 35.5M42× A.15 Cross-Hardware Validation Table 10. A6000 Ada throughput. Peak SPS at batch 65,536. No code changes required. EnvironmentA6000 Ada SPSvs. Reference HalfCheetah JAX13.1M290× (vs. Gymnasium 45K) Pong JAX (scan)1.4B23,333× (vs. C 60M env-only) CartPole JAX (scan)7.1B43,279× (vs. Gymnasium 164K) 16 Automatic Generation of High-Performance RL Environments 141664128 Parallel Environments 0 10 20 30 40 50 60 SPS (thousands) PyBoy EmuRust Random CNN rollout PPO 1K4K16K66K Batch Size 100K 1M 10M 100M SPS Showdown (21K) PokeJAX Random RNN rollout PPO 2561K4K16K Batch Size 100 200 300 400 500 600 700 SPS (thousands) Python 16p TCG Pocket Random GRU rollout PPO Figure 5. Throughput scaling. EmuRust (left) saturates at 128 CPU envs. PokeJAX (center) scales linearly with GPU batch size. TCG Pocket (right): Python peaks at 16 processes; JAX scales with batch size. RandomGRU RolloutGRU PPO 1M 10M 100M 1000M SPS (log scale) 4.6x 47x 35x Puffer Pong C (PufferLib) Rust+PyO3 JAX scan Figure 6. PufferLib comparisons. Puffer Pong: JAX achieves 31× rollout and 42× PPO over C. A.16 Throughput Scaling Figures A.17 Training Time Breakdown Figures A.18 TCG Pocket Agent Translation Metrics The TCG Pocket translation was conducted entirely through logged sessions with Gemini 3 Flash Preview. Phase 1 translated five core modules (1,452 source lines) via programmatic API calls: 20 iterations consuming 83K tokens ($0.02). Phase 2 used the Gemini CLI to translate six logic-heavy modules: 29.3M input tokens across 256 messages ($4.96), with 79–95% cache 17 Automatic Generation of High-Performance RL Environments 0 20 40 60 80 100 % of PPO Training Time 1.8% 1.5% 1.0% 2M2M2M 0.4% 0.3% 0.2% 20M20M20M 0.0%0.0%0.0% 200M200M200M Pong C (PufferLib) Pong Rust (PyO3+Rayon) Pong JAX (scan-fused) Puffer Pong: PPO Training Breakdown Env StepData TransferModel Fwd/Bwd Figure 7. PufferLib PPO training breakdown. C and Rust backends incur CPU→GPU data transfer overhead; the all-JAX stack eliminates this entirely. hit rates. 051015202530 Cumulative tokens (M) 0 10 20 30 40 50 Count Phase 1 (SDK) Phase 2 (Gemini CLI) ab_utils+cb ab_eff mv_eff supporter engine Agent translation progress (Gemini 3 Flash Preview) Cumulative tests passing Cumulative errors 0246 Iterations utils card deck player game_state ab_utils+card_base ability_eff move_eff supporter ptcg_engine PASS PASS PASS PASS PASS PASS PASS PASS PASS PASS Module status Phase 1 (SDK) Phase 2 (CLI) Figure 8. Agent translation metrics for TCG Pocket. Cumulative L1 tests passing vs. tokens consumed. Total: $4.98 for 4,235 lines. A.19 Detailed Per-Environment Architecture EmuRust module structure. Five modules: CPU (SM83, 161 lines), memory (MBC1/3/5, 315 lines), PPU (scanline rendering, 400 lines), emulator core (1,008 lines), and PyO3 bindings (318 lines). Rayon’spar_iter_mut()parallelizes acrossNinstances with zero-copy NumPy buffers. 18 Automatic Generation of High-Performance RL Environments PokeJAX architectural changes. Three changes: (1) server/client flattening into pure functions on state pytrees, (2) fixed-size state representation, and (3) branch-parallel effect dispatch via jax.lax.switch. The complete translation is 55,629 lines across 10 module groups. HalfCheetah JAX architecture. Five modules: model constants (242 lines), forward kinematics (183 lines), forward dynamics (348 lines, analytical RNEA), contact solver (230 lines, analytical Jacobians), and environment wrapper (199 lines). Reference baselines. PyBoy: 167K FPS at 32 processes. Pokemon Showdown: 21K SPS native, 681 SPS via PokeEnv. TCG Pocket Python: 140K SPS at 16 processes. Gymnasium HalfCheetah: 45K SPS. MJX: 1.6M SPS at batch 32K. PufferLib C Pong: 60M SPS (2.4M training). A.20 CartPole JAX CartPole JAX (229 lines) achieves 838M SPS at batch 65,536 (5,112×over Gymnasium) and 187M SPS for scan-fused PPO. Our agent-generated implementation is 2.7×faster than Gymnax’s [9] hand-authored CartPole. Training curves across 10 seeds confirm policy equivalence (Figure 9). 0100200300400 Gradient steps 100 200 300 400 500 Episode reward CartPole Gymnasium JAX Figure 9. Appendix training curves. CartPole (10 seeds,±1σ), showing JAX and Gymnasium converge to the same maximum reward of 500, confirming policy equivalence for this environment. A.21 Test Coverage 19 Automatic Generation of High-Performance RL Environments Table 11. Test coverage by environment. Pass/Total counts allpytest-collected test functions, including parametrized variants and regression tests added after translation (hence larger than Table 4’s translation- time counts). Line coverage measured withpytest-cov(Python) orcargo test(Rust). ∗ PokeJAX: JIT-compiled dispatch prevents line-level instrumentation; figure reflects only the directly instrumented mechanics/core modules (1,097 stmts). † EmuRustGBA: 110 unit tests in Rust source; 85 integration tests and 23 hardware-feature tests exercise the PyO3 Python bindings. ‡ Failing tests: Pong’s 1 failure is a statistical distribution test sensitive to sample size (same test fails for both Gemini and Claude translations, Table 6); HalfCheetah’s 4 failures are tight-tolerance parametrized tests affected by float32 vs. float64 differences (all core L1/L2/L3 tests pass). EnvironmentPass/Total Stmts Coverage Notes CartPole JAX9/910760% main(), rendering untested Pong JAX5/615873% 1 statistical test ‡ HalfCheetah JAX136/14046296% 4 float32 tolerance ‡ TCG Pocket JAX50/50 1,85877% L1/L2/L3 all passing PokeJAX ∗ 90/95 1,09753% JIT limits instrumentation EmuRustGBA † 218/21852798% 110 Rust + 108 Python tests Total508/518 B Representative Agent Prompts This appendix presents representative prompts used during agent-assisted translation (§3.3), condensed for space. Each prompt follows a generic template structure that stays constant across all environments: (1) source module specification with line count, (2) target language constraints, (3) interface contract (function signatures and return types), (4) reference behavior (source code pasted verbatim), and (5) instruction to generate Level 1 property tests. The parts that vary per-environment are the module source code, target constraints, and interface contracts—filled in by the human for each module. The examples below are instantiated for EmuRust and contain sufficient detail to reproduce the translation methodology for any environment. B.1 Module Translation Prompt The following prompt initiates translation of a single module. The agent receives the source code, target language constraints, and the module’s interface contract. Translate the following Game Boy CPU module from C/Python (PyBoy) to Rust. Source module: cpu.py (161 lines) –- SM83 instruction set implementation. Target constraints: - Pure Rust, no unsafe except for FFI boundaries - All registers as a struct with public fields (for save/load state) - Instruction dispatch via match on opcode byte - Return cycle count from each instruction for PPU synchronization Interface contract: - fn step(&mut self, mem: &mut Memory) -> u32 –- execute one instruction, return T-cycles - fn handle_interrupts(&mut self, mem: &mut Memory) –- check and dispatch IF/IE Reference behavior (from PyBoy source): [Source code of cpu.py pasted here, 161 lines] Begin translation. After completing, write Level 1 property tests that verify each instruction against reference input/output pairs. 20 Automatic Generation of High-Performance RL Environments B.2 Level 1 Test Generation Prompt After module translation, this prompt generates property tests from reference I/O pairs. Generate Level 1 property tests for the Rust CPU module. Test strategy: For each instruction category (arithmetic, load, jump, bitwise), create test cases using known input/output pairs extracted from the reference implementation. Example reference pairs (from running PyBoy with instrumentation): - ADD A,B: A=0x3C, B=0x12 → A=0x4E, F.Z=0, F.N=0, F.H=0, F.C=0 - ADD A,B: A=0xFF, B=0x01 → A=0x00, F.Z=1, F.N=0, F.H=1, F.C=1 - SUB A,B: A=0x3E, B=0x3E → A=0x00, F.Z=1, F.N=1, F.H=0, F.C=0 - RL A: A=0x80, F.C=0 → A=0x00, F.Z=1, F.C=1 Requirements: - Cover all flag-affecting instructions (at least 2 cases each: normal + edge) - Test boundary conditions: zero, overflow, half-carry - Each test constructs a CPU state, executes one instruction, asserts post-state Write #[cfg(test)] Rust tests. Each test should be self-contained and clearly named (e.g., test_add_a_b_overflow). B.3 Level 2 Interaction Test Prompt Interaction tests verify cross-module behavior that individual property tests cannot catch. Generate Level 2 interaction tests for the Game Boy emulator. Module interactions to test: 1. CPU–PPU timing: CPU instructions must advance the PPU by the correct number of T-cycles. After executing a 12-cycle instruction, PPU.dot should advance by 12. 2. Memory banking: Writing 0x01 to 0x2000 should switch ROM bank, and subsequent reads from 0x4000–0x7F should return bank 1 data. 3. Interrupt–CPU: Setting IF bit 0 (VBlank) and enabling IE bit 0 with IME=1 should cause the CPU to push PC, jump to 0x0040, and clear the IF bit. Test pattern: 1. Set up initial state across multiple modules 2. Execute a sequence of operations 3. Assert state changes in all affected modules Focus on interactions where bugs were found during initial translation (timing drift between CPU and PPU was the most common failure mode). B.4 Bug Repair Prompt When Level 3 rollout comparison detects a divergence, this prompt structure feeds the failure back to the agent for root-cause analysis. Level 3 rollout comparison failed at step 847. Divergence: - Frame 847: pixel (23, 91) = 0x1C (EmuRust) vs 0x9C (PyBoy) - First diverging byte in VRAM at offset 0x0340 21 Automatic Generation of High-Performance RL Environments - All frames 0–846 matched exactly State at step 846 (last matching): - PC = 0x0267, SP = 0xFFF8 - LCDC = 0x91, STAT = 0x85, LY = 0x90 - ROM bank = 1, RAM bank = 0 Action taken at step 847: button = 0x00 (no input) Diagnose the root cause. Start by: 1. Check what instruction executes at PC=0x0267 in both implementations 2. Compare memory writes in the PPU scanline that produces line 91 3. Check if the VRAM divergence affects tile data or the background map After identifying the bug, fix it and add a targeted Level 1 or Level 2 test that would have caught this failure. C Performance Optimization Guide After a translated environment passes all three verification levels, the next step is performance optimization. This appendix provides concrete techniques and a reusable agent prompt for maximizing environment throughput. The techniques are organized by target backend: JAX (GPU) and Rust (CPU). C.1 JAX Optimization Checklist The following patterns, distilled from our case studies, consistently improve JAX environment throughput. They are ordered by typical impact. 1. Fixed-size state arrays. JAX requires array shapes to be known at compile time. Replace all dynamic-length data structures (lists, dicts with varying keys, variable-length arrays) with fixed-sizejnp.ndarrayfields padded to maximum capacity. Use a sentinel value (e.g.,-1or NO_CARD_ID) for unused slots. In TCG Pocket, this reduced card zone storage from Python lists to fixed (MAX_HAND_SIZE,) arrays, enabling JIT compilation of the entire game engine. 2. Branchless conditionals withjnp.where. Replace Pythonif/elsewithjnp.where(condition, true_val, false_val). Both branches are computed and the result is selected by mask—this is faster on GPU because it avoids warp divergence. For multi-way branches, use nestedjnp.where orjax.lax.switch. Reservejax.lax.condfor unbatched cases where one branch is signifi- cantly more expensive (it evaluates only the selected branch). Note that undervmap,lax.cond evaluates both branches regardless, because different batch elements may take different paths; in batched contexts,jnp.whereis preferred. In Puffer Pong, all ball-paddle collision logic uses jnp.where: ball_vy = jnp.where(wall_hit, -ball_vy, ball_vy) ball_vx = jnp.where(paddle_hit, -ball_vx, ball_vx) 3.vmapfor batch parallelism. Write environment logic for a single instance, then apply jax.vmapto vectorize across the batch dimension. This generates fused GPU kernels that process all environments in one call. Mark shared constants (terrain maps, card databases) with in_axes=None so they are broadcast rather than duplicated: 22 Automatic Generation of High-Performance RL Environments step_batch = jax.vmap(step_single, in_axes=(0, 0)) step_with_terrain = jax.vmap( partial(step, terrain=terrain), in_axes=(0, 0) # terrain not batched ) 4. JIT the outer interface. Applyjax.jitto thevmappedstep and reset functions so the entire batch operation compiles to a single GPU kernel. Pre-compile during initialization to avoid first-call latency during training: self._step_jit = jax.jit(step_batch) self._reset_jit = jax.jit(reset_batch) # Warmup: call once with dummy data _ = self._step_jit(dummy_states, dummy_actions) 5.lax.scanfor multi-step fusion. When the training loop callsenv.stepinside a rollout loop, fuse the loop withjax.lax.scanto compile the entire rollout into one kernel. This eliminates per-step CPU→GPU dispatch overhead. In CartPole, this improved throughput by 3.2×over a Python loop calling jitted steps: def scan_body(states, actions_t): states, rewards, terminals = step_batch(states, actions_t) return states, (rewards, terminals) rollout = jax.jit(partial(jax.lax.scan, scan_body)) 6. Minimize data types. Useint8for categorical state (entity types, directions, flags) and float32only for values requiring arithmetic. For example, usingint8for categorical entity fields can reduce per-environment state significantly, improving memory bandwidth utilization. 7. Pre-allocate reward and observation buffers. Initialize all output arrays (rewards, terminals, observations) as zeros in the state. Update in-place with.at[].set()rather than creating new arrays. Avoid jnp.concatenate or jnp.stack in the hot path. 8. Normalize observations at the source. Compute normalized observations inside the JIT- compiled step function rather than in a separate Python post-processing step. Pre-compute constant denominators: PADDLE_RANGE = MAX_PADDLE_Y - MIN_PADDLE_Y # constant obs_paddle = (state.paddle_y - MIN_PADDLE_Y) / PADDLE_RANGE C.2 Rust Optimization Checklist 1. Rayonpar_iterfor environment parallelism. Userayon::prelude::par_iter_mutto step all environments in parallel across CPU cores. Each environment is independent, making this embarrassingly parallel: self.emulators.par_iter_mut() .zip(actions.iter()) .for_each(|(emu, &action)| emu.step(action)); This typically provides near-linear scaling up to the number of physical cores (8−16×). 23 Automatic Generation of High-Performance RL Environments 2. Pre-allocate observation buffers. Allocate observation, reward, and terminal buffers once at initialization, then reuse every step via slice copies. AvoidVec::pushor allocation in the step loop: let obs_buffer = vec![0u8; num_envs * OBS_SIZE]; // In step(): copy directly into pre-allocated slice obs_buffer[i*OBS_SIZE..(i+1)*OBS_SIZE] .copy_from_slice(&emu.get_obs()); 3. Frame skip without rendering. For emulator environments, implement a fast path that skips PPU/rendering for intermediate frames. Only render the final frame that produces the observation. In EmuRust, this saved ∼60% of per-step time at frame skip 24: emu.run_frames_no_render(frame_skip - 1); // fast path emu.run_frame(); // render last frame 4. Lookup tables for game mechanics. Replace computed game logic with pre-computedconst arrays. For example, element-type effectiveness matrices, passability checks, and noise gradients can all be pre-computed as static lookup tables: const EFFECT_MATRIX: [[i32; 5]; 5] = [[1,1,1,1,1], ...]; let damage_mult = EFFECT_MATRIX[atk_type][def_type]; 5.#[inline(always)]on hot functions. Mark observation writing, single-step physics, and reward computation as#[inline(always)]to eliminate function call overhead in tight loops. Profile first—only inline functions called millions of times per second. 6.Arc<Vec<»for shared immutable data. When each environment instance needs access to large immutable data (ROM images, card databases, terrain maps), wrap it inArcand clone the reference: let rom = Arc::new(rom_data); let emulators: Vec<_> = (0..num_envs) .map(|_| Emulator::new(rom.clone())) .collect(); One copy in memory regardless of batch size. 7. Compact struct layout. Separate hot data (accessed every step) from cold data (accessed occasionally). Keep entity structs small—usei32instead ofi64, pack booleans into bitfields or i32 flags. This improves L1/L2 cache utilization. 8. Efficient PyO3 bindings. For the Python↔Rust boundary: accept NumPy arrays via PyReadonlyArrayN(zero-copy read), return observations by writing directly into a pre-allocated NumPy array viaPyArrayN::as_slice_mut(). Minimize the number of Python→Rust calls per step (one call for all environments, not one per environment). C.3 Optimization Agent Prompt The following prompt is used after the environment passes Level 1–3 verification. It instructs the coding agent to optimize throughput without changing semantics. 24 Automatic Generation of High-Performance RL Environments The [JAX/Rust] environment implementation has passed all verification tests (Level 1 property tests, Level 2 interaction tests, Level 3 rollout comparison). Now optimize it for maximum steps-per-second (SPS) throughput. Current performance: [X] SPS at batch size [B] on [hardware]. Target: Maximize SPS while maintaining all existing tests passing. Constraints: - All Level 1, 2, and 3 tests must continue to pass after optimization - Do not change the environment’s external API (step, reset, observation/reward shapes) - Do not change game semantics or reward logic [For JAX environments] Apply these optimizations in order: 1. Replace any remaining Python if/else on JAX values with jnp.where or jax.lax.cond 2. Ensure all state arrays have static shapes (no dynamic allocation) 3. Apply jax.vmap for batch parallelism over a single-instance step function 4. Wrap the vmapped function with jax.jit 5. Reduce data types: use int8 for categorical fields, float32 only for arithmetic 6. Pre-compute observation normalization constants 7. Profile with jax.profiler and eliminate remaining bottlenecks [For Rust environments] Apply these optimizations in order: 1. Add rayon dependency and parallelize step/reset with par_iter_mut 2. Pre-allocate all output buffers (obs, rewards, terminals) at initialization 3. Add #[inline(always)] to step, observation, and reward functions 4. Replace computed game logic with const lookup tables where applicable 5. Implement frame-skip fast path (skip rendering for intermediate frames) 6. Use Arc<Vec<» for shared immutable data across environments 7. Profile with cargo flamegraph and eliminate remaining bottlenecks After each optimization: 1. Run the full test suite to verify correctness 2. Measure SPS at batch sizes [32, 128, 512, 2048, 8192] 3. Report the speedup from each change Begin with a profiling analysis to identify the current bottleneck, then apply optimizations targeting that bottleneck first. 25 Automatic Generation of High-Performance RL Environments Algorithm 1 Hierarchical translation and verification. Require:Reference environmentE ref , modulesm 1 ,...,m K in dependency order, test specifi- cations T 1 ,T 2 ,T 3 , max iterations T, episode count N Ensure: Performance environment E perf satisfying semantic equivalence 1: Phase 1: Module translation (Level 1) 2: for k = 1 to K do 3: m ′ k ← Agent(m k ,L tgt ) Translate module m k to target language 4: for t = 1 to T do 5:if RunTests(T 1 ,m ′ k ) = Pass then 6:break 7:else 8:m ′ k ← Agent(failures,m ′ k ) Repair using L1 diagnostics 9:end if 10: end for 11: if t = T then 12:Request human intervention for module m k 13: end if 14: end for 15: Phase 2: Integration (Level 2) 16: E perf ← Compose(m ′ 1 ,...,m ′ K ) 17: for t = 1 to T do 18: if RunTests(T 2 ,E perf ) = Pass then 19:break 20: else 21:Identify failing module(s); repair while preserving L1 correctness 22: end if 23: end for 24: Phase 3: Validation (Level 3) 25: for t = 1 to T do 26: Run N episodes in E ref and E perf with matched seeds and actions 27: if all per-step outputs match (exact or within ε) then 28:break 29: else 30:Root-cause analysis: add targeted L1/L2 tests; repair and re-verify L1, L2 31: end if 32: end for 33: Phase 4: Cross-backend validation (Level 4) 34: repeat 35: Train policy π in E perf 36: Evaluate π in E ref ; compute reward gap ∆ 37: if ∆ is statistically significant then 38:Diagnose sim-to-sim gap; add targeted L1/L2 tests 39:go to Phase 1 with new tests 40: end if 41: until ∆ is not statistically significant 42: return E perf 26

Instant research discovery