PixelBrax: Learning Continuous Control from Pixels End-to-End on the GPU
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.
Code (1)
Tasks
Benchmarkingcontinuous-controlContinuous ControlCPUGPUReinforcement Learning (RL)Methods 이 논문이 사용한 방법론
Similar Papers 제목 키워드 기반
PixelSmile: Toward Fine-Grained Facial Expression Editing
Fine-grained facial expression editing has long been limited by intrinsic semantic overlap. To address this, we construct the Flex Facial Expression (FFE) dataset with continuous affective annotations and establish FFE-B…
Contrastive LearningrQdia: Regularizing Q-Value Distributions With Image Augmentation
rQdia regularizes Q-value distributions with augmented images in pixel-based deep reinforcement learning. With a simple auxiliary loss, that equalizes these distributions via MSE, rQdia boosts DrQ and SAC on 9/12 and 10/…
continuous-controlContinuous ControlDeep Reinforcement LearningImage Augmentation+1Data-Efficient Learning of Feedback Policies from Image Pixels using Deep Dynamical Models
Data-efficient reinforcement learning (RL) in continuous state-action spaces using very high-dimensional observations remains a key challenge in developing fully autonomous systems. We consider a particularly important i…
Model-based Reinforcement LearningModel Predictive Controlreinforcement-learningReinforcement Learning+1From Pixels to Torques: Policy Learning with Deep Dynamical Models
Data-efficient learning in continuous state-action spaces using very high-dimensional observations remains a key challenge in developing fully autonomous systems. In this paper, we consider one instance of this challenge…
Model-based Reinforcement LearningModel Predictive Controlreinforcement-learningReinforcement Learning+1Solving Continuous Control via Q-learning
While there has been substantial success for solving continuous control with actor-critic methods, simpler critic-only methods such as Q-learning find limited application in the associated high-dimensional action spaces.…
continuous-controlContinuous ControlMulti-agent Reinforcement LearningQ-Learning