A GPU-Parallel Framework for Heterogeneous Multi-Task Reinforcement Learning

AI summary

Overview

Research area: Robotics, specifically GPU-accelerated multi-task reinforcement learning (MTRL) for robot manipulation, plus robot-learning benchmark design.

Technical level: Intermediate. The paper assumes familiarity with reinforcement learning (policy optimization, PPO, behavior cloning, reward shaping), but the core ideas — training one policy across many tasks and reusing demonstrations — are described in accessible terms.

Scope: The paper introduces Hebero, a GPU-parallel Isaac Lab benchmark that jointly trains and evaluates one policy across 40 heterogeneous manipulation tasks, together with DGPO (Demonstration-Guided Policy Optimization) and its reference recipe IW-ABC for demonstration-guided multi-task learning.

What This Paper Is About

Existing robot-learning benchmarks rarely combine large-scale GPU-parallel simulation with heterogeneous manipulation tasks and a standardized multi-task evaluation protocol. The authors ask whether GPU-parallel infrastructure can support capability breadth — a single policy acquiring many structured manipulation skills in one training process — rather than producing one specialist policy per task. Their answer is Hebero (Heterogeneous Benchmark for Robot Learning), which compiles all 40 LIBERO tasks into one heterogeneous GPU-parallel Isaac Lab process with shared state and visual interfaces, demonstrations, and a common evaluation contract, and DGPO, a shared demonstration-guided training framework whose IW-ABC recipe coordinates imitation strength and task priority using a cheap per-task learning-progress signal.

Key Contributions

1. Hebero, an efficient and scalable multi-task RL benchmark. All 40 LIBERO tasks are compiled into one heterogeneous GPU-parallel RL process with common state and visual interfaces, demonstrations, and evaluation protocols. GPU parallelism increases throughput, and scaling environments per task improves success under a fixed wall-clock budget.

2. A demonstration-guided training framework (DGPO) and its reference recipe. DGPO reuses demonstrations for dense tracking rewards and asymmetric value learning, providing a shared training stack with IW-ABC as its reference recipe. Applied to four Piper tasks, this recipe trains a joint state-based policy that transfers to the physical robot with 82.5% success and no weight updates.

3. Adaptive imitation guidance and task balancing (IW-ABC). Within DGPO, IW-ABC uses the same lightweight per-task learning progress signal to guide both imitation and online RL across heterogeneous tasks. Compared with FAMO-ABC, the strongest state-input baseline, it improves Mean SR from 82.3% to 90.1% and Long SR from 53.3% to 70.0%.

4. A controlled comparison stack. The shared DGPO components (demonstration command stream, tracking reward, privileged critic inputs, task-conditioned PPO backbone, network capacity) are held fixed while demonstration interfaces vary, enabling controlled comparisons of PPO, BC→PPO, DAPG, ABC, and an RFCL-inspired reset curriculum under matched conditions.

Main Findings

• Simulation throughput: Across 40 tasks and 1,600 environments, Hebero delivers approximately 7.5k simulation steps per second (7,500 SPS) on one L20 GPU with 10.6 GiB VRAM, which the paper reports as 3.4 times the throughput of its 96-core MuJoCo baseline (2,200 SPS, 1.5 TiB host PSS).

• End-to-end multi-GPU training: On eight L20 GPUs with 40 tasks and 25,600 environments, Hebero reaches 78,500 SPS end-to-end RL training throughput at 144.9 GiB VRAM. The paper lists MTBench MT50-rand (Isaac Gym, 50 tasks, 24,576 environments) at approximately 69,000 SPS with hardware and peak memory not reported.

• More parallel environments help learning: Increasing replicas per task yields higher Mean SR at the common training-budget endpoint within a fixed wall-clock budget, and reaches matched mean training reward sooner.

• BC→PPO underperforms ABC: Behavior-cloning initialization plus PPO scores 47.5 ± 2.6 Mean SR, whereas ABC (adaptive behavior cloning without importance weighting) reaches 75.5 ± 3.6, indicating the value of online adaptive demonstration guidance over offline BC initialization.

• IW-ABC is the strongest state-input learner: 90.1 ± 3.8 Mean SR, 70.0 ± 10.0 Long SR, 35/40 task coverage, 82.1 SR-AUC, using a 0.32M-parameter actor.

• Improvement over the strongest baseline: IW-ABC outperforms FAMO-ABC (82.3 ± 2.9 Mean SR, 53.3 ± 5.8 Long SR, 30/40 coverage, 75.6 SR-AUC) by 7.8 percentage points in Mean SR and improves Long SR from 53.3% to 70.0%.

• Other IW pairings: IW-DAPG scores 68.5 ± 3.0 Mean SR and IW-RFCL 66.4 ± 2.8, both below IW-ABC. IW-PPO alone scores 44.9 ± 2.5 Mean SR but 20.0 ± 0.0 Long SR, trading lower Mean SR for higher Long SR than PPO (50.8 ± 3.1 Mean SR, 10.0 ± 0.0 Long SR); pairing IW with ABC improves both.

• Visual policy is stronger: Vis IW-ABC reaches 93.5 ± 2.6 Mean SR, 81.9 ± 11.5 Long SR, 38/40 task coverage, and 83.3 SR-AUC, with 5.92M total deployed parameters including its frozen encoder.

• Long-horizon coverage: ABC, IW-ABC, and Vis IW-ABC cover 4, 7, and 9 of the ten Long tasks, respectively.

• Real-robot transfer: One jointly trained state-input policy executes four representative RoboTwin Piper tasks on a physical Piper robot at 82.5% overall success across 20 physical trials per task, with policy weights fixed throughout deployment.

• IW concentrates on lagging tasks: Most Object and Spatial tasks succeed early, while the harder Goal and Long tasks learn more slowly; IW assigns these larger weights until their success approaches the multi-task mean.

• ABC anneals by design: Once the vectorized rollout-step counter reaches the annealing horizon, all BC coefficients reach the residual floor, relaxing constraints on deviation from demonstration actions while IW continues prioritizing lagging tasks.

Methodology in Plain English

The authors start from a benchmark problem rather than an algorithm problem. They take 40 LIBERO manipulation tasks — ten each from the Goal, Object, Spatial, and Long suites — and compile each one offline into a cached JSON descriptor holding its instruction, scene, assets, spawn regions, goal predicates, and thresholds. MuJoCo MJCF assets are converted offline to cached USD files that preserve geometry, appearance, dynamics, joints, and articulation limits.

At runtime, every environment instantiates the assets for its assigned task, but all environments share one Isaac Lab simulator, one renderer, one rollout buffer, and one policy update. For K tasks and g environments per task, the total is N = Kg environments, with environment i assigned to task i mod K. Task success is the conjunction of relational, articulation, and activation predicates evaluated from live simulator state, which provides a common sparse-reward event and the evaluation outcome.

To make this common contract usable by one agent, Hebero standardizes interfaces: a normalized 7-dimensional action (six for relative end-effector translation and rotation through an operational-space controller, one for the binary gripper), shared proprioception and previous-action inputs, and a task encoding. State actors see a 300-feature input including an object–target pose buffer of 26 entity slots with 9 features each; visual actors see a 450-feature input using pooled frozen-ViT features from third-person and wrist images (2 × 192 = 384 features).

Rewards are assembled from a task reward (goal reward minus smoothness/control regularizers minus joint-limit penalties) plus, for the DGPO configurations, a dense demonstration-tracking reward that sums exponential kernels over robot-motion errors and world-state errors such as rigid-object pose and articulation joint positions. A demonstration command stream maintains each environment's task identity, active demonstration, and timestep cursor, so the tracking reward, critic, and learner interfaces share the same references without duplicating trajectory storage.

Evaluation deliberately resamples object placement. For each task and movable object, coordinate-wise bounds are computed from the 50 recorded demonstration starts, and each evaluation episode independently redraws the object's x and y coordinates uniformly within that box, with collision-based rejection sampling. This is intended to test placement robustness rather than trajectory memorization, in contrast to LIBERO's fixed evaluation starts.

The learning comparison is kept fair by fixing the shared stack — reward, observations, critic inputs, PPO backbone, and network capacity — and varying only how each learner consumes demonstrations: none (PPO), offline actor initialization (BC→PPO), an offline likelihood regularizer (DAPG), a live time-aligned action target (ABC), or a demonstration-state reset curriculum (RFCL). The authors' own recipe, IW-ABC, uses one per-task success-rate EMA as a progress signal. That same scalar does two jobs: it scales down the behavior-cloning penalty as a task succeeds and as training time advances, and it computes an importance weight that gives tasks below the multi-task mean success a larger share of the PPO gradient. All state-input learners share a compact [512, 256, 128] task-conditioned actor–critic, train on 40 tasks for 30,000 online PPO iterations with 50 demonstrations per task, use three training seeds, and run on eight NVIDIA L20 GPUs with 2,000 parallel environments per GPU for approximately two days each.

Why This Matters

The paper reframes the scaling question for robot RL: the bottleneck in massively parallel simulation is no longer sample collection but value stability, task imbalance, and how prior data is used. Hebero supplies a benchmark where those algorithmic questions can be compared under one protocol, and IW-ABC shows that a very cheap signal — a per-task success-rate EMA already computed from training rollouts — can coordinate both imitation pressure and task prioritization without extra evaluation. The 82.5% real-robot result on a physical Piper robot with no weight updates also argues that simulation-trained multi-task policies are deployable, at least for the four tasks tested.

Real-world applications (implied by the work, not reported as deployed systems):

• Warehouse and logistics manipulation, where one policy would handle varied pick, place, and rearrangement tasks instead of requiring a separate policy per SKU flow.
• Household service robots, since the LIBERO task families include kitchen-style goal, object, spatial, and long-horizon manipulation.
• Manufacturing or lab automation cells that need a single controller to switch between structured tasks with different objects and goals.
• Robot foundation-model pipelines, where a compact task-conditioned actor (the state policy uses a 0.32M-parameter actor) can be paired with an external perception pipeline rather than learning end-to-end vision.

Industry relevance: The throughput and memory results matter to anyone paying for GPU time. The paper reports 7,500 SPS on a single L20 at 10.6 GiB VRAM and 78,500 SPS end-to-end on eight L20 GPUs at 144.9 GiB VRAM, which makes large-scale joint training practical on commodity-class accelerators. The finding that more parallel replicas per task improve success within a fixed wall-clock budget gives an actionable scaling lever, and the state-input policy's sim-to-real transfer with an external model-based perception pipeline suggests a deployment route that avoids the harder problem of visual sim-to-real transfer.

Future Directions

• Extending beyond LIBERO and RoboTwin. The authors state that budget constraints limited training and evaluation to these two task sources, and they call for extending the benchmark construction methodology and DGPO training framework to more challenging tasks and larger task sets.

• Fixing target misalignment under trajectory drift. IW-ABC's time-indexed action targets can misalign when trajectories drift, even with the tracking reward in place. A more robust demonstration-alignment mechanism is an open problem.

• Testing whether the scaling trend holds further. The paper shows that more environment replicas per task improve success under a fixed wall-clock budget; how far that trend continues, and where it saturates, is not established.

• Validating the success-based progress signal more broadly. IW-ABC relies on a per-task success-rate EMA that is independent of reward scale; whether it composes as cleanly with other task-balancing methods and other demonstration interfaces beyond those tested remains open.

Target Audience

Researchers and engineers working on multi-task and demonstration-guided reinforcement learning, robot manipulation, and GPU-accelerated simulation benchmarks will benefit most. It is also relevant to practitioners deciding how to budget GPU resources for multi-task training and to teams evaluating sim-to-real transfer of joint policies. Readers without a reinforcement learning background can follow the benchmark and scaling results, but the algorithmic sections assume familiarity with PPO, behavior cloning, and advantage/value learning.

Authors’ abstract

GPU-parallel simulation provides abundant robot interaction, but existing benchmarks rarely combine this scale with heterogeneous manipulation tasks and standardized multi-task RL evaluation. We introduce Hebero (Heterogeneous Benchmark for Robot Learning), a GPU-parallel Isaac Lab benchmark that enables efficient joint training and evaluation of a single policy across all 40 heterogeneous tasks. Scaling experiments show that increasing parallel replicas per task improves success under a fixed wall-clock budget. To support learning with sparse rewards and limited demonstrations, we propose Demonstration-Guided Policy Optimization (DGPO), which reuses demonstrations for dense tracking rewards and asymmetric value learning. Its shared stack supports controlled comparisons of learner-specific demonstration interfaces within PPO. Within DGPO framework, we introduce IW-ABC, which uses a lightweight per-task learning progress signal to coordinate adaptive behavior cloning (ABC), relaxing demonstration guidance with task progress, and importance weighting (IW), emphasizing lagging tasks in PPO updates. With 50 demonstrations per task, IW-ABC achieves 90.1% state-input mean success, outperforming the strongest baseline FAMO-ABC by 7.8 percentage points. Its visual counterpart reaches 93.5% mean success. Real-world experiments further demonstrate that a single multi-task policy trained in simulation can successfully perform four tasks on a physical Piper robot. The project page is available at https://hebero-rl.github.io/.