PixelBrax让强化学习在GPU上端到端处理像素控制，速度提升百倍。
PixelBrax: Learning Continuous Control from Pixels End-to-End on the GPU
- 用JAX渲染+Brax物理引擎，实现像素输入的端到端训练
- 千并发环境运行速度比传统基准快100倍以上
- 支持可复现实验与泛化测试，适合研究高效强化学习
我们提出PixelBrax，一套基于像素观测的连续控制任务。结合Brax物理引擎与纯JAX渲染器，使强化学习实验可在GPU上端到端运行。PixelBrax可同时渲染数千个并行环境，运行速度比依赖CPU渲染的现有基准快两个数量级。此外，通过显式处理环境中的随机性，支持完全可复现的实验，并提供颜色与视频干扰项以评估模型泛化能力。项目已开源，包含JAX实现的多个强化学习算法，代码位于github.com/trevormcinroe/pixelbrax。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。