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

TL;DR

PixelBrax introduces GPU-accelerated, pixel-based continuous control environments enabling fast, scalable reinforcement learning experiments with reproducibility and generalization benchmarking capabilities.

Contribution

It combines Brax physics engine with a JAX renderer for end-to-end GPU-based RL experiments on pixel observations, supporting large-scale parallelism and reproducibility.

Findings

Runs thousands of environments in parallel with high speed.

Achieves two orders of magnitude faster performance than CPU-based benchmarks.

Supports reproducibility and generalization testing with distractors.

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.