Automatic Generation of High-Performance RL Environments

Seth Karten; Rahul Dev Appapogu; Chi Jin

cs.LGMar 12, 2026

Automatic Generation of High-Performance RL Environments

Authors: Seth Karten, Rahul Dev Appapogu, Chi Jin

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

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. (1) Direct translation (no prior performance implementation exists): EmuRust (Game Boy emulator via Rust IPC); PokeJAX (Pokemon Showdown via JAX). (2) Translation verified against existing performance implementations : throughput parity with Puffer Pong, MJX, and Brax at matched GPU batch sizes. (3) New environment creation : TCGJax, the first Pokemon TCG Pocket environment, created from a web-extracted specification.

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.

Env	Source	Target	Src LoC	Tgt LoC	Key Challenge
EmuRust	C/Python	Rust+PyO3	${\sim}$ 26K	2.5k	Cycle-accurate emulation
PokeJAX	TypeScript	JAX	${\sim}$ 100K	55k	state complexity.
HalfCheetah	MuJoCo	JAX	245	1.2k	Articulated body + contact
TCGJax	Web rules ⋆	Py $\to$ JAX	29k	4.2k	Rule extraction from web
Pong	C (PufferLib)	Rust+JAX	225	235/318	Already-optimized baseline

Table 1: Environment overview. ⋆ Private reference (contamination control).

Environment	Benchmark	Reference (SPS)	Performance (SPS)	Speedup
Direct translation into newly performant environments (no prior performance implementation)
EmuRust	Random action	167K (PyBoy, 32p)	$239{\pm}6$ K (Rust, 64e)	$1.4\times$
EmuRust	PPO training	9.9K (PyBoy, 32p)	$14.5{\pm}0.4$ K (Rust, 128e)	$1.5\times$
PokeJAX	Random action	21K (Showdown, 1p)	$500{\pm}9$ M (JAX, 65Kb)	$23k\times$
PokeJAX	PPO training	681 (Showdown)	$15.2{\pm}0.2$ M (JAX)	$22k\times$
Translation verified against existing performance implementations

Table 2: Throughput comparison. Mean

\pm

std from

N{=}5

runs (CVs

{<}3\%

);

{\sim}

2M models; JAX excludes JIT warm-up.

Figure 3: PPO training time breakdown across model scales. Three bars per implementation show 2M, 20M, 200M parameter models. Performance implementations drop to

{\leq}4\%

env overhead at 200M. All on 1

\times

RTX 5090.

Environment	Train Backend	Eval (Perf)	Eval (Ref)	Equiv.
Puffer Pong	C (ref)	$28.01\pm 0.28$	$28.04\pm 0.29$	$\checkmark$
Puffer Pong	JAX (perf)	$28.23\pm 0.18$	$28.22\pm 0.20$	$\checkmark$
HalfCheetah	MJX (ref)	$1398\pm 497$	$1389\pm 511$	$\checkmark$
HalfCheetah	JAX (perf)	$1026\pm 636$	$1133\pm 562$	$\checkmark$
EmuRust (Red)	PyBoy (ref)	$12.01\pm 0.12$	$11.99\pm 0.15$	$\checkmark$
EmuRust (Red)	Rust (perf)	$12.06\pm 0.00$	$12.06\pm 0.01$	$\checkmark$

Table 3: Cross-backend policy transfer. Values are mean

\pm

std over 10 seeds. Equivalence confirmed via TOST (

\alpha{=}0.05

) with environment-specific margins (

\Delta

): Pong

\Delta{=}1.0

, HalfCheetah

\Delta{=}100

, EmuRust

\Delta{=}0.5

, PokeJAX

\Delta{=}0.02

, TCGJax

\Delta{=}0.05

. PokeJAX and TCGJax report win rate against a heuristic bot; others report return.

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 ).

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

Table 4: Translation cost. Costs include all iterations; base rates from Gemini 3 Flash Preview logs for environments requiring multiple revision cycles.

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	$\epsilon$ ( $10^{-3}$ )	10	✓	✓
TCGJax	20	24	100	exact	10	✓	✓
Puffer Pong	6	3	100	exact	10	✓	✓

Appendix

Table 5: Verification summary. L1 = component generation (property tests), L2 = interaction tests, L3 = rollout comparison, Xfer = cross-backend policy transfer.

Environment	Agent	Iters	Tests	Cost
Pong	Gemini 3 Flash	13	6/6	$0.05
	Claude Sonnet 4.6	3	5/6 §	${\sim}$ $0.08
HalfCheetah	Gemini 3 Flash	20	69/69	$3.26
	Claude Opus 4.6	6	69/69	— †

Appendix

Table 6: Multi-agent comparison. Identical inputs; functionally equivalent outputs.

Environment	Benchmark	PufferLib C	Rust	JAX	Speedup
Puffer Pong	Random (env only)	60M	122M	275M	$4.6\times$
	PufferLib training	2.4M (134K)	—	—	—
	GRU Rollout (2M)	4.5M	4.5M	140M	$31\times$
	GRU PPO (2M)	854K	855K	35.5M	$42\times$

Appendix

Table 7: PufferLib comparisons. “PufferLib training” reports their full pipeline; matched rows use

{\sim}

2M GRU. All on 1

\times

RTX 5090.

Environment	A6000 Ada SPS	vs. Reference
HalfCheetah JAX	13.1M	$290\times$ (vs. Gymnasium 45K)
Pong JAX (scan)	1.4B	$23k\times$ (vs. C 60M env-only)
CartPole JAX (scan)	7.1B	$43k\times$ (vs. Gymnasium 164K)

Appendix

Table 8: A6000 Ada throughput. Peak SPS at batch 65k. No code changes required.

Figure 5: Agent translation metrics for TCG Pocket. Cumulative L1 tests passing vs. tokens consumed. Total: $4.98 for 4k lines.

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

Appendix

Table 9: Test coverage by environment. Pass/Total counts all pytest -collected test functions, including parametrized variants and regression tests added after translation (hence larger than Table 4 ’s translation-time counts). Line coverage measured with pytest-cov (Python) or cargo test (Rust). ∗ PokeJAX: JIT-compiled dispatch prevents line-level instrumentation; figure reflects only the directly instrumented mechanics/core modules. † 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).

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

Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAX
Oct 2, 2025Waris Radji, Thomas Michel, Hector PiteauRL Benchmarks Reinforcement Learning
Univ. Lille, Inria, CNRS, Centrale Lille, UMR 9189-CRIStAL, France · Independent Researcher
Envs-FORGE: Frontier-Optimized Reward-Grounded Environment Synthesis for Agent RL
Aug 14, 2026Xiaojun Wu, Cehao Yang, Honghao Liu +7RL for LLM Agents Reinforcement Learning
IDEA Research · The Hong Kong University of Science and Technology (Guangzhou) · DataArcTech Ltd.
EnvFactory: Scaling Tool-Use Agents via Executable Environments Synthesis and Robust RL
May 18, 2026Minrui Xu, Zilin Wang, Mengyi DENG +12Tool-Augmented Language Model Agents Agentic RL
LARK, HKUST (GZ) · University of Cambridge · UCL +2

Oct 2, 2025J/K move · Enter open · S save

Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAX

Waris Radji, Thomas Michel, Hector Piteau

Univ. Lille, Inria, CNRS, Centrale Lille, UMR 9189-CRIStAL, France · Independent Researcher

RL Benchmarks Reinforcement Learning GPU-Accelerated RL RL for Game AI

Open paper AskPDF arXiv

Reinforcement learning (RL) research requires diverse, challenging environments that are both tractable and scalable. While modern video games may offer rich dynamics, they are computationally expensive and poorly suited for large-scale experimentation due to their CPU-bound execution. We introduce Octax, a high-performance suite of classic arcade game environments implemented in JAX, based on CHIP-8 emulation, a predecessor to Atari, which is widely adopted as a benchmark in RL research. Octax provides the JAX community with a long-awaited end-to-end GPU alternative to Atari games, offering image-based environments, spanning puzzle, action, and strategy genres, all executable at massive scale on modern GPUs. Our JAX-based implementation achieves orders-of-magnitude speedups over traditional CPU emulators. We demonstrate Octax's capabilities by training RL agents across multiple games, showing significant improvements in training speed and scalability compared to existing solutions. The environment's modular design enables researchers to easily extend the suite with new games or generate novel environments using large language models, making it an ideal platform for large-scale RL experimentation. Our open-source framework is available at https://github.com/riiswa/octax/.

Abstract

Reinforcement Learning

Explore similar work

Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAX

Envs-FORGE: Frontier-Optimized Reward-Grounded Environment Synthesis for Agent RL

EnvFactory: Scaling Tool-Use Agents via Executable Environments Synthesis and Robust RL