# PixelBrax: Learning Continuous Control from Pixels End-to-End on the GPU

arXiv: https://arxiv.org/abs/2502.00021 · HTML: https://mlsanity.com/paper/2502.00021/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 16 Jan 2025.

## Built on by (most cited first, 1 of 1)

- [Studying the Interplay Between the Actor and Critic Representations in Reinforcement Learning](https://mlsanity.com/paper/2503.06343/) (2025)

## Abstract (arXiv)

We present PixelBrax, a set of continuous control tasks with pixel observations. We combine the Brax physics engine with a pure JAX renderer, allowing reinforcement learning (RL) experiments to run end-to-end on the GPU. PixelBrax can render observations over thousands of parallel environments and can run two orders of magnitude faster than existing benchmarks that rely on CPU-based rendering. Additionally, PixelBrax supports fully reproducible experiments through its explicit handling of any stochasticity within the environments and supports color and video distractors for benchmarking generalization. We open-source PixelBrax alongside JAX implementations of several RL algorithms at github.com/trevormcinroe/pixelbrax.