Rose improves embodied reinforcement learning training using GPU optimization
Rose: Efficient Embodied RL through Graph Optimization and GPU Compute Resource Control
Distributed, Parallel, and Cluster Computing
Summary
Training robots or agents to learn by interacting with their environment takes a lot of computer power, especially on GPUs. The authors show that parts of this work often leave the GPU idle even when it could be working. They created a system called Rose that smartly reshuffles tasks and manages GPU resources to keep the GPU busier and speed up training. This approach helps train models about 1.2 to 1.9 times faster while using the GPU more efficiently.
What this means in practice
- •For robotics engineers: Accelerate training of robotic control policies by optimizing GPU use during embodied RL simulations to reduce overall model training time.
- •For graphics hardware system designers: Inform design of GPU runtime schedulers to better manage resource sharing across concurrent RL workloads with fine-grained task overlap.
Authors
Wen Zeng, Ke Hong, Xu Fu, Zixiao Huang, Zehao Wang, Zhuoran Zhu, Yuanqing Wang, Bo Dai, Zhen Guo, Hao Lin, Chunyang Zhu, Boxun Li, Chao Yu, Quanlu Zhang, Guohao Dai, Zhenhua Zhu, Yu Wang
Abstract
Embodied reinforcement learning (RL) coordinates policy inference, simulation, and model training with highly heterogeneous resource demands. Although individual components frequently underutilize GPU compute during execution, existing systems only overlap work at coarse component boundaries (e.g., pipelining inference and simulation across simulator groups), leaving substantial low-activity GPU windows within subcomponent phases and nested loops. Exploiting these fine-grained windows manually requires tedious dependency analysis and workload-specific coordination. Furthermore, even when independent work is scheduled concurrently, cross-process GPU time slicing and streaming multiprocessor (SM) contention can negate the benefits of overlap. We present Rose, a runtime-aware system that automates fine-grained overlap and GPU resource management for embodied RL while preserving synchronous training semantics. Through a lightweight programmable interface that exposes execution hierarchy and state accesses, Rose tracks runtime data readiness and GPU activity to perform a two-stage optimization. First, a window-guided graph optimizer identifies low-activity windows and applies graph transformations to overlap complementary work across components and loop levels. Second, leveraging NVIDIA Multi-Process Service (MPS) for cross-process GPU concurrency, a slack-aware resource manager tunes process-level SM quotas by exploiting branch synchronization slack to mitigate interference on the critical path. Both stages combine local profiling with end-to-end iteration validation to guarantee throughput gains with minimal tuning overhead. Across diverse policies, simulators, RL algorithms, and GPU platforms, Rose achieves 1.20--1.92$\times$ training throughput improvement and increases average SM activity by 23.8--39.3\% compared with the state-of-the-art systems.