PixelBrax: Learning Continuous Control from Pixels End-to-End on the GPU
Fuente:
arXiv
Saved in:
| Main Authors: | , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866916593362010112 |
|---|---|
| author | McInroe, Trevor Garcin, Samuel |
| author_facet | McInroe, Trevor Garcin, Samuel |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2502_00021 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | PixelBrax: Learning Continuous Control from Pixels End-to-End on the GPU McInroe, Trevor Garcin, Samuel Machine Learning Performance 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. |
| title | PixelBrax: Learning Continuous Control from Pixels End-to-End on the GPU |
| topic | Machine Learning Performance |
| url | https://arxiv.org/abs/2502.00021 |