arXiv · 2502.00021
PixelBrax: Learning Continuous Control from Pixels End-to-End on the GPU
Abstract
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.
Explore related subjects
Keep this discovery
Trevor McInroe, Samuel Garcin. 2025-01-16. PixelBrax: Learning Continuous Control from Pixels End-to-End on the GPU. https://arxiv.org/abs/2502.00021
Cite the original work for its findings. Save a collection to share your selection of sources.