PaperlayerView exact revision on arXiv

Source-linked AI summary

Automatic Generation of High-Performance RL Environments

Seth Karten, Rahul Dev Appapogu, Chi Jin

arXiv:2603.12145v2cs.LGcs.AIcs.SE

TL;DR

Complex RL environments often make simulation the dominant training cost, while producing faster implementations has required specialized engineering. The paper presents an agent-driven closed-loop translation and verification methodology, validating equivalent environments across five cases with end-to-end PPO speedups of 1.5× to 42×.

  • Problem

    RL simulation typically consumes 50–90% of wall-clock training time, while high-performance rewrites have required labor-intensive specialized engineering for individual domains.

  • Method

    The methodology combines hierarchical component, interaction, rollout, and cross-backend policy-transfer verification with iterative repair and an agent-agnostic translation workflow.

  • Results

    Across five environments, the methodology confirmed no sim-to-sim gap and achieved end-to-end PPO speedups ranging from 1.5× to 42×, including parity with Google’s MJX.

  • Takeaways & Limitations

    The approach can make fast verified simulation a routine RL workflow step, with re-translation costing under $1 when a reference updates.

  • Takeaways & Limitations

    The method is best suited to reproducible, bounded environments and may require additional engineering for nondeterministic dependencies or unbounded allocation; mature hand-optimized engines may yield parity rather than speedup.

Abstract

from arXiv · show

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.

1 Introduction

RL environment simulation can dominate training cost, while existing high-performance rewrites require specialized engineering. The paper proposes a closed-loop, verifiable approach using coding agents and validates it across direct translation, optimized-baseline translation, and new environment creation.

  • Motivation: Environment simulation consumes 50–90% of typical RL training wall-clock time, with greater overhead for complex simulators.Pokemon Showdown exceeds 100K lines of TypeScript, and cycle-accurate Game Boy emulators are also costly to simulate.
  • Motivation: Existing high-throughput RL environments were produced through labor-intensive, domain-specific engineering for individual environments.The cited examples include Brax, Gymnax, Pgx, JaxMARL, Craftax, and PureJaxRL.
  • Problem: Generation alone is insufficient because silent errors in game mechanics or physics accumulate over long rollouts and corrupt training signals.The paper identifies structured verification as the missing feedback needed for agents to iterate toward correct translations.
  • Contribution: The proposed contribution pairs hierarchical verification with cross-backend policy transfer in a closed loop, providing structured error signals for iterative repair.The verification hierarchy includes property, interaction, and rollout tests, with policy transfer closing the loop.
  • Evaluation scope: The methodology is validated across direct translation, translation against existing optimized implementations, and new environment creation from a web-extracted specification.The case studies include EmuRust, PokeJAX, Pong and HalfCheetah, and TCGJax.

2 Related Work

Prior work has established high-throughput RL environments through manual domain-specific optimization and systems-level batching. This paper instead uses coding agents to generate parallelizable environments and hierarchical verification to address silent translation errors, complementing existing infrastructure.

  • Hardware-accelerated environments: Hardware-accelerated RL libraries manually reimplement environments in JAX or on GPUs, typically requiring substantial specialized engineering per domain.Examples include Brax, MJX, Gymnax, Pgx, JaxMARL, Craftax, and PureJaxRL.
  • High-throughput RL systems: High-throughput systems use standardized interfaces, asynchronous batching, unified C-environment access, or GPU-utilization strategies to accelerate RL.The paper positions its method as complementary because reducing per-step time lets these systems exploit parallelism.
  • Sim-to-real transfer: Sim-to-sim verification transfers policies between reference and translated simulators to test behavioral equivalence rather than robustness to physical mismatch.The paper distinguishes this from sim-to-real transfer, which addresses differences between simulation and physical systems.
  • LLM-assisted code generation: Unlike function- or API-level code generation, this setting translates full RL environments where silent mechanics or physics errors can accumulate across timesteps.Hierarchical verification localizes errors at the component level before they propagate to rollouts.

3 Methodology

The methodology translates a reference environment into a high-performance equivalent through agent-generated modules, hierarchical tests, iterative repair, and cross-backend policy transfer. Equivalence is assessed behaviorally and, for continuous environments, within specified tolerances, while throughput must move training away from environment-bound execution.

  • Problem statement: Behavioral equivalence requires matching observations, rewards, and termination signals for every seed and action sequence; continuous environments use per-component L∞ tolerance ϵ.The definition compares the reference environment Eref with the performance environment Eperf across timesteps.
  • Sim-to-sim verification: Cross-backend policy transfer evaluates policies trained in each environment on the other using identical hyperparameters and training algorithms.Reward indistinguishability is tested over 100 episodes with diverse random seeds rather than formal proof.
  • Pipeline: The pipeline decomposes a reference environment into modules, translates them with a coding agent, and verifies progressively broader behavior with targeted repair after failures.Level 4 cross-backend policy transfer closes the outer verification loop.
  • Backend selection: JAX is selected for small-state environments benefiting from GPU parallelism, whereas Rust is selected for sequential or memory-intensive environments.Backend choice is based on environment structure rather than a single universal implementation strategy.
  • Verification hierarchy: Exhaustive rollout comparison is intractable, so four increasing verification levels test components, interactions, matched rollouts, and learned-policy transfer.Detected gaps feed back into Levels 1–3 until they close.
  • Agent translation: The agent receives module source code, target-language specifications, and test requirements in one prompt, while human involvement is limited to prompt and verification-structure design.The methodology is described as agent-agnostic despite the reported use of Gemini 3 Flash Preview.

4 Experiments

Across five diverse environments and three translation settings, agent-generated implementations achieve high throughput and behavioral equivalence. Hierarchical verification supports convergence, while training-time profiles show the environment ceases to be the bottleneck at larger model scales.

  • Experimental design: Five environments span discrete games, continuous physics, hardware emulation, and multi-agent systems across direct translation, verified translation, and specification-based creation.The set includes EmuRust, PokeJAX, Pong, HalfCheetah, and TCGJax.
  • Throughput results: 23k× speedup for PokeJAX enables convergent training, while HalfCheetah reaches 1.04× throughput parity with Google’s MJX.Pong achieves 42× end-to-end PPO speedup, and TCGJax becomes a trainable JAX environment from a web-extracted specification.
  • Training-time breakdown: At 200M parameters, all single-agent performance implementations contribute ≤4% of training time, shifting training from environment-bound to model-bound.The reference environments consume 50–90% of wall-clock time, whereas the performance implementations reduce this overhead substantially.
  • Policy equivalence: All five environments pass L3 rollout comparison over 100 episodes with matched seeds and step-level output checks, using exact matching for discrete environments and ϵ=10−3 for HalfCheetah.Matched training curves across Pong, HalfCheetah, and EmuRust show consistent learning dynamics across backends.
  • Policy equivalence: All five environments pass cross-backend policy transfer equivalence tests using TOST with environment-specific margins, including exact win-rate transfer for PokeJAX.EmuRust policies transfer to PyBoy with near-identical reward, while TCGJax confirms equivalence in both directions.
  • Verification and effort: Hierarchical verification is necessary for complex environments: L3-only verification failed after 42 HalfCheetah iterations, whereas L1 property tests immediately caught dynamics errors.The reported errors included a Coriolis force sign error, mass matrix asymmetry, and bias force magnitude violations.

5 Conclusion

The paper presents a closed-loop methodology that translates reference RL environments into equivalent high-performance implementations using coding agents and hierarchical verification. Across five environments, it confirms equivalence and reports broad training-speed improvements, while identifying scope limits for nondeterministic or already optimized environments.

  • Method: The methodology combines component generation, interaction tests, rollout comparison, cross-backend policy transfer, and iterative repair.These verification levels provide feedback for closing the sim-to-sim gap.
  • Results: All five environments achieved cross-backend policy equivalence, with TOST equivalence confirming no sim-to-sim gap.The conclusion reports confirmation across direct translations, translations checked against optimized implementations, and new environment creation.
  • Results: End-to-end PPO speedups ranged from 1.5× to 42×, with throughput parity with Google’s MJX.The reported result indicates that translated implementations can shift training from environment-bound to model-bound.
  • Limitations: The methodology is most effective for reproducible environments with clear module boundaries and fixed-size state representations.Nondeterministic external dependencies and unbounded dynamic allocation may require additional engineering.
  • Limitations: Against mature hand-optimized engines, the methodology may achieve parity rather than speedup, as illustrated by HalfCheetah’s 1.04× result versus MJX.The authors recommend applying it primarily to unoptimized or new environments.
  • Implications: Re-translation after reference updates costs under $1, with the test suite serving as a regression guard.The authors frame fast verified simulation as a potential standard step in RL workflows.

A.1 Per-Environment Details

The appendix describes translations spanning emulation, competitive game simulation, rigid-body physics, and specification-driven environment creation. The implementations use different target architectures and demonstrate substantial throughput or training benefits in selected settings.

  • EmuRust: EmuRust translates the Game Boy emulator into Rust, using Rayon’s shared-memory thread pool to pack 128 environments into one process.At matched 32-core resources, this achieves a 1.5× comparison against PyBoy through higher per-core utilization and zero IPC overhead.
  • PokeJAX: PokeJAX converts the 100K+ line Pokemon Showdown server into a 55k-line JAX translation across approximately 30 modules.The redesign flattens server/client structure, uses fixed-size state arrays, and applies branch-parallel dispatch with jax.lax.switch.
  • HalfCheetah: HalfCheetah required articulated-body dynamics, contact modeling, and four solver revisions, totaling $3.26 with all 69 tests passing.The translation used pure JAX across five modules and 1k lines.
  • TCGJax: TCGJax translates web-extracted Pokemon TCG Pocket rules through a 29k-line Python reference into a 4k-line JAX implementation.L1 tests and rule verification caught an erroneous early use of rules from another trading card game.
  • TCGJax: 153K SPS at batch 4K enabled TCGJax to converge to reward 1.0 in approximately 12 minutes, compared with 23K SPS for the Python reference.The private reference also provides a contamination control because it is absent from public repositories.
  • Puffer Pong: Puffer Pong translates an already optimized C environment to Rust and JAX, enabling fused GPU rollouts through jax.lax.scan.The reported 42× PPO speedup reflects a CPU-to-GPU architectural change rather than like-for-like optimization.

A.3 Multi-Agent Validation

The multi-agent validation reran translations with Claude Sonnet 4.6 and Claude Opus 4.6 using identical prompts and test suites. Both agents produced functionally correct translations, supporting agent-agnostic generalization.

  • Multi-Agent Validation: Claude Sonnet 4.6 on Pong and Claude Opus 4.6 on HalfCheetah both converged to functionally correct translations.The comparison used identical inputs and test suites.
  • Multi-Agent Validation: The results confirm that the methodology is agent-agnostic across these representative environments.Table 6 summarizes the multi-agent comparison.

A.4 Verification Ablation Details

The ablation compares hierarchical verification with L3-only rollout testing on Pong and HalfCheetah. Hierarchical verification is especially important for complex rigid-body dynamics, where coarse end-to-end feedback failed to isolate bugs.

  • Failure Analysis: L3-only verification failed for HalfCheetah because coarse rollout feedback could not distinguish dynamics bugs from end-to-end failures.The observed failure involved Coriolis-force signs, contact Jacobians, vectorization rewrites, and stability patches.
  • Pong: On Pong, hierarchical verification converged in 13 iterations over 3.5 minutes, while L3-only verification required 15 iterations and 2.4× longer wall-clock time.Both approaches succeeded on the simpler game-logic environment.
  • Failure Analysis: Across Pong and HalfCheetah, L3-only verification failed when contact dynamics and multi-body kinematic chains were involved.The reported complexity threshold lies between simple game logic and rigid-body physics with at least 6 degrees of freedom.

A.6 Experimental Details

The experiments standardize backend comparisons with matched training settings and report throughput, compilation, memory, and baseline details. JAX timing excludes one-time compilation, whose cost is small relative to typical training runs.

  • Measurement protocol: JAX benchmarks exclude one-time JIT compilation from steady-state throughput measurements.A warm-up call occurs before measurement.
  • Measurement protocol: ∼3 s to ∼45 s: JIT compilation ranges across environments, amortizing to <1% over 10-minute training runs.For PokeJAX, 45 s of JIT compilation adds ∼2.5% over a typical 30-minute run.
  • Resource usage: ∼4 GB to ∼28 GB: reported GPU memory varies by environment and batch size.HalfCheetah uses ∼4 GB at batch 65K, Pong ∼2 GB, PokeJAX ∼28 GB at 65K, and TCGJax ∼8 GB at 16K.
  • Training setup: Training uses a 2.5 × 10−4 learning rate, 0.2 clip ratio, 4 epochs, GAE λ = 0.95, and γ = 0.99.Environment-specific batch sizes are matched between backends.
  • Baselines: Table 7 compares PufferLib training and matched ∼2M-GRU rows on 1× RTX 5090.The passage identifies the comparison structure but does not provide the table’s numerical entries.
  • Baselines: Table 8 reports peak SPS at batch 65k on A6000 Ada without code changes.The passage specifies the measurement setting but not the numerical throughput values.

A.9 TCG Pocket Agent Translation Metrics

TCG Pocket translation used logged Gemini sessions and a staged implementation process, followed by performance-oriented JAX transformations. The appendix describes static state, vectorization, JIT compilation, rollout fusion, and memory-conscious representations as key techniques.

  • Translation metrics: 29.3M input tokens across 256 messages: Phase 2 translated six logic-heavy modules with 79–95% cache hit rates.Phase 1 translated five core modules totaling 1.4k source lines in 20 iterations and 83K tokens.
  • Translation metrics: $4.98 for 4k lines: the TCG Pocket translation’s reported total cost is shown alongside cumulative L1 tests passing versus tokens consumed.Phase 1 cost $0.02 and Phase 2 cost $4.96.
  • JAX optimization: Fixed-size arrays replace dynamic structures, enabling JIT compilation of the entire game engine.In TCG Pocket, card-zone storage uses fixed (MAX_HAND_SIZE,) arrays with sentinel values for unused slots.
  • JAX optimization: jnp.where replaces batched Python branching, while vmap applies single-instance logic across environments and broadcasts shared constants with in_axes=None.Under vmap, lax.cond evaluates both branches, so jnp.where is preferred for batched contexts.
  • JAX optimization: 3.2×: lax.scan improved CartPole throughput over a Python loop calling jitted steps by fusing multi-step rollouts.The fusion eliminates per-step CPU→GPU dispatch overhead.
  • JAX optimization: Pre-allocated buffers, compact data types, and source-side observation normalization reduce allocation and memory overhead in JAX execution.Categorical state can use int8, while normalized observations are computed inside the JIT-compiled step function.

B.2 Rust Optimization Checklist

The appendix provides optimization checklists for Rust and JAX environments while requiring preservation of verified behavior and external interfaces. It also specifies profiling, throughput measurement, and hierarchical translation-validation procedures.

  • Rust optimization: 8−16×: Rayon parallelization typically scales near-linearly up to the number of physical CPU cores.The checklist applies par_iter_mut to step independent environments in parallel.
  • Rust optimization: ∼60%: EmuRust’s frame-skip fast path saved approximately this fraction of per-step time at frame skip 24.Intermediate frames skip rendering, while only the final frame produces the observation.
  • Rust optimization: One copy in memory regardless of batch size: Arc shares immutable ROM, card, or terrain data across environment instances.The approach clones references rather than duplicating the underlying data.
  • Optimization constraints: The optimization prompt preserves all Level 1–3 tests, the external API, and simulator or reward logic while maximizing SPS.It begins with profiling and targets the identified bottleneck first.
  • Optimization checklists: JAX optimization applies static shapes, vmap, jitted interfaces, reduced data types, and precomputed normalization constants; Rust optimization applies Rayon, pre-allocation, inlining, lookup tables, and frame skipping.Both workflows conclude with profiler-guided bottleneck elimination.
  • Translation and verification: The translation workflow uses generic prompts containing module specifications, target constraints, interface contracts, reference behavior, and Level 1 property-test instructions.Later phases integrate modules, validate matched-seed rollouts, and perform cross-backend policy evaluation.
  • Translation and verification: Cross-backend validation trains a policy in the performance environment, evaluates it in the reference environment, and computes a reward gap Δ.A diagnosed sim-to-sim gap triggers targeted tests, repair, and re-verification.
Loading 2603.12145v2…