Roxie: A New JAX-Based Reinforcement Learning Framework for MuJoCo Continuous Control

Developer Vittorio recently shared a new reinforcement learning (RL) framework called Roxie on Reddit. Built entirely in JAX, it's designed specifically for continuous control tasks in MuJoCo environments. For researchers who have spent years navigating the intersection of reinforcement learning and physics simulation, Roxie aims to address a long-standing pain point: efficiently scaling MuJoCo simulations across different hardware.

What Problem Does Roxie Solve?

In real-world RL engineering, MuJoCo is the go-to physics simulation engine — but scaling it across hardware is notoriously frustrating. Data transfers between GPU and CPU, managing environment parallelism, and optimizing the compilation of training loops: any one of these, handled poorly, can drag down overall throughput.

According to the author, Roxie's core idea is to drive MuJoCo's different hardware paths through a single unified interface. In other words, whether the underlying execution runs on a GPU or a pure CPU environment, developers interact with the same API — and the framework handles the mapping to the appropriate hardware path internally. This design is especially useful for teams that need to run experiments under varying compute conditions, eliminating the need to write separate adapter code for different hardware.

Seven Built-In Agents, Sharing a Unified Base Class

One of Roxie's most immediate selling points is its ready-to-use algorithm implementations. The framework includes seven mainstream algorithms for continuous control:

• DDPG (Deep Deterministic Policy Gradient)
• TD3 (Twin Delayed DDPG)
• D4PG (Distributed Distributional DDPG)
• TD4
• SAC (Soft Actor-Critic)
• MPO (Maximum a Posteriori Policy Optimization)
• PPO (Proximal Policy Optimization)

These algorithms span multiple technical lineages — from deterministic policy gradients to maximum-entropy methods to distributional value estimation. A notable engineering detail: all agents share the same base class. This means algorithmic differences are isolated to the smallest possible pluggable modules, reducing maintenance overhead and making cross-algorithm comparisons fairer — since they all run on exactly the same underlying infrastructure.

These seven algorithms represent the major branches of current RL research for continuous control. DDPG was the first to apply deep networks to deterministic policy gradients; TD3 improved on it by using twin Q-networks and delayed policy updates to mitigate overestimation. D4PG introduced distributional value functions, modeling Q-values as distributions rather than scalars to improve sample efficiency. SAC belongs to the maximum-entropy RL family, encouraging exploration by adding a policy entropy term to the reward — and has consistently led continuous control benchmarks. MPO, developed at DeepMind, uses an EM-based framework with KL divergence constraints on policy update step sizes, making it relatively robust to hyperparameters. PPO is arguably the most engineering-stable policy gradient method, using clipped importance sampling ratios to prevent overly large policy updates. Sharing a single base class means the data pipeline, environment interaction, and logging infrastructure are completely consistent — algorithmic differences show up only in loss functions and minimal network structure variations, greatly reducing experimental noise caused by conflating "algorithmic differences" with "engineering differences."

Fully Compiled Loop: Keeping Hardware Running at Full Capacity

A key performance design in Roxie is compiling both the acting (sampling) and learning phases into a single XLA program.

This is a textbook approach to leveraging the JAX ecosystem's strengths. By compiling the entire execution chain into XLA, Roxie ensures that no operation needs to return to the host during runtime. In traditional RL implementations, Python's main loop frequently intervenes, and repeated communication between the CPU and accelerator often becomes a throughput bottleneck. A fully compiled loop allows thousands of parallel environments to be continuously "fed" data, so the accelerator never idles waiting for input.

For research scenarios requiring large-scale parallel simulation, this design significantly boosts sample throughput — and it's precisely where JAX has a competitive edge over traditional imperative frameworks.

XLA (Accelerated Linear Algebra) is a domain-specific compiler developed by Google, originally for TensorFlow and now the core execution backend for JAX. XLA works by ahead-of-time (AOT) compiling computation graphs into hardware-optimized machine code, rather than interpreting operations one-by-one like traditional Python/NumPy. JAX uses the jit (just-in-time compilation) decorator to pass Python functions to XLA; compiled functions run on GPU/TPU with virtually no Python interpreter overhead.

In the RL context, another major XLA advantage is operator fusion: the compiler automatically merges multiple small operations — matrix multiplications, activation functions, normalization — into a single GPU kernel call, dramatically reducing memory read/write cycles. By compiling both sampling and learning into the same XLA program, Roxie treats the entire RL inner loop as a single indivisible compilation unit, completely eliminating the latency from the Python main loop repeatedly handshaking between GPU and CPU.

500M-Step Benchmarks and Pretrained Weights

To demonstrate the framework's viability, Roxie ships with a solid set of benchmarks. The author trained all 7 agents with a 500 million-step budget on 3 humanoid tasks from dm_control, covering both GPU and no-GPU (pure CPU) compute configurations.

This benchmark design carries two implications: the 500M-step scale is sufficient to reveal algorithmic stability and final performance over long training runs; and successfully running on both GPU and CPU directly validates the framework's core claim of a unified interface driving multiple hardware paths.

For users who don't want to train from scratch, the author has uploaded pretrained checkpoints to Hugging Face (vittorione/roxie-release-v1) for direct download and use. The framework code is hosted on GitHub (github.com/vittorione94/roxie), and the author has openly invited the community to try it out and provide feedback.

Who Should Pay Attention to Roxie?

Roxie has a clear positioning — it's not a general-purpose, all-encompassing RL library, but a vertically focused tool for MuJoCo continuous control. For researchers and engineers working on robot control, motion learning, and physics-simulation-based RL, Roxie's unified multi-algorithm implementation, fully compiled training loop, and cross-hardware support all deliver practical, immediately applicable value.

As a freshly released open-source project, Roxie's long-term ecosystem health, documentation quality, and community activity remain to be seen. But from an engineering design perspective, it capitalizes on the performance dividends of the JAX + MuJoCo combination — and for practitioners in the field, it's well worth a try.