Embodied reinforcement learning speeds up with asynchronous resource management

EBRL: Asynchronous Embodied RL by Multi-Grained Resource Management

Machine LearningDistributed, Parallel, and Cluster Computing

Summary

Embodied reinforcement learning systems improve models by simulating environments, generating actions, and updating models, but these steps use CPUs and GPUs differently, causing inefficiencies. The authors introduce EBRL, a system that schedules simulation and training without waiting for synchronization, and dynamically manages CPU and GPU resources based on need. This approach speeds up the whole learning process by better using available hardware. Their tests show EBRL can be over three times faster in simulation throughput and train models 2.5 times faster than previous leading systems.

What this means in practice

  • For robotics developers: Train embodied agents more quickly by efficiently managing CPU and GPU during simulation and learning phases.
  • For autonomous vehicle teams: Accelerate the training of vehicle control policies by asynchronously overlapping environment simulation and policy updates to shorten training time.

Authors

Liang Mi, Weijun Wang, Bowen Gao, Tianze Yu, Zixu Hao, Han Xiao, Xin Ding, Mingzhe Huang, Xin He, Lu Shi, Hao Wu, Haipeng Dai, Guihai Chen, Yunxin Liu, Ting Cao

Abstract

Embodied reinforcement learning (RL) improves model capabilities with a pipeline of environment simulation, action generation, and model updates. These stages show heterogeneous CPU and GPU demands, making efficient resource utilization difficult. Recent systems overlap rollout (simulation and generation) with training for efficiency, but exclusive GPU allocation and synchronized barrier in rollout still leave substantial hardware resource waste. In this paper, we present EBRL, an asynchronous embodied RL training system with two core techniques. The asynchronous pipelined scheduler overlaps rollout and training, pipelines simulation and generation across environment groups, and carries out each environment independently, eliminating synchronization stalls. The fine-grained resource manager pools CPU cores and GPU streaming multiprocessors, and uses stage profiles and runtime feedback to adjust resource quotas and batch sizes to meet the shifting demands among stages. We implement EBRL on RLinf and evaluate it with four embodied policies and four simulation benchmarks across heterogeneous GPU testbeds. Experiments show that EBRL achieves 1.30-3.47 times the end-to-end rollout throughput and 2.5 times of training convergency compared to the SOTA embodied RL systems.