TOPPO: Rethinking PPO for Multi-Task Reinforcement Learning with Critic Balancing

Abstract

Soft Actor-Critic (SAC) and its variants dominate Multi-Task Reinforcement Learning (MTRL) due to their off-policy sample efficiency, while on-policy methods such as Proximal Policy Optimization (PPO) remain underexplored. We diagnose that PPO in MTRL suffers from a previously overlooked issue: critic-side gradient ill-conditioning, which may cause tail tasks to stall while easy tasks dominate the value function’s updates. To address this, we propose TOPPO (Tail-Optimized PPO), a reformulation of PPO via Critic Balancing—a set of modules that improve gradient conditioning and balance learning dynamics across tasks. Unlike prior approaches that rely on modular architectures or large models, TOPPO targets the optimization bottleneck within PPO itself. Empirically, on the Meta-World+ MT50 benchmark, TOPPO achieves stronger mean and tail-task performance than prior SAC-family and ARS-family methods while using substantially fewer parameters and reaching their final performance in fewer environment steps. Notably, TOPPO matches or surpasses strong SAC baselines early in training and maintains superior performance at full budget. Ablations confirm the effectiveness of each module in TOPPO and provide insights into their interactions. Our results demonstrate that, with proper optimization, on-policy methods can rival or exceed off-policy approaches in MTRL, challenging the prevailing reliance on SAC and highlighting critic-side gradient conditioning as the central bottleneck.

Mean success hides the tail

In Meta-World, one policy and one value function are shared by many robot-arm manipulation tasks, and the only task information is a one-hot ID. MT10 has 10 tasks and MT50 has 50. The usual score is the mean success over tasks, so a handful of tasks can fail without moving it much.

So we also report worst-k success: the mean over each seed’s k lowest tasks. By that measure the field’s default backbone looks very different. Every published non-ARS SAC baseline scores 0.0% on its ten worst MT50 tasks, and the best published worst-10 is 29.3%, from ARS-LoRA. A carefully rebuilt multi-task PPO gets 6.6%. Where does PPO still fail on the hard tasks?

The critic is ill-conditioned

For every minibatch we compute each task’s gradient on the shared weights, separately for the actor and the critic. On MT10, the per-task critic gradient norms spread by a median of 497× across tasks; the actor’s spread by 4.1×. The actor has a built-in fix: advantages are normalized within each task. The critic has none. Its squared loss carries each task’s return scale, squared.

We read the critic problem along three lines. Tasks with large gradients set the scale of the mean (scale spread). Critic gradients collapse onto a few shared feature directions (co-linear collapse). And even with both reduced, a plain mean can hand most of the update to a few tasks (unfair aggregation). Our hypothesis is that loud tasks dominate the shared critic, quiet tasks get worse value estimates, and their policy updates, and the data they collect next, suffer in turn. The diagnostics are correlational: they support this story but do not prove the full causal chain.

Try it: the mean follows the loudest task

Take two orthogonal task gradients with norms 10 and 1. Their mean points almost exactly along the loud one (cosine 0.995). FairGrad (Ban & Ji, ICML 2024) at α = 1 instead chooses weights so that every task contributes equally to the update; here that means weighting each gradient by one over its norm. In the sketch below, four tasks share one critic. Drag the arrowheads, or tab to one and use the arrow keys.

TOPPO: one module per failure mode

TOPPO keeps the PPO loop and its task-shared actor and critic MLPs. It adds one module where each part of the diagnosis points, and a light fix on the actor. Within each minibatch it combines per-task gradients (FairGrad on the critic, PCGrad on the actor), clips globally and takes one Adam step.

Why α = 1

FairGrad picks task weights by an α-fair criterion. At α = 1 its solution has two properties that fail for every other α.

Theorem 1. For linearly independent task gradients with Gram matrix G ≻ 0, let w > 0 solve Gw = w⁻¹. Then d = Σ_i w_i g_i has norm exactly √K, and d does not change when any task’s gradient is multiplied by a positive constant.

The first property keeps the critic’s step size predictable from minibatch to minibatch. The second has a sharp consequence: PopArt acts on the critic gradients as exactly such a per-task rescaling, so within a minibatch FairGrad gives the same update with or without it. That predicts that PopArt’s benefit should shrink once FairGrad is on, and on V2 rewards it does: PopArt adds 6.64 worst-10 points without FairGrad and −0.22 with it.

The solve has to be fast, since it runs on every minibatch. A projected Newton solver with a closed-form diagonal warm start takes a median 0.40 ms per solve, against 32.5 ms for SciPy, an 83× median speedup on 1,984 Gram matrices recorded during MT50 training.

Results on Meta-World

All runs below use one protocol: Meta-World V2 rewards, 20M (MT10) or 100M (MT50) environment steps, final checkpoint, 10 seeds, deterministic evaluation on all 50 goals of every task.

On MT50, TOPPO leads the ARS-LN-1024 rerun by 5.1 points on the mean and 13.2 on worst-10, with 22.7× fewer parameters, and first reaches that rerun’s final mean at about 28M of its 100M steps. Against the published numbers, which use a different protocol and worst-k convention, its mean is above the strongest published mean (ARS-LN-1024, 88.9%), and its worst-10 exceeds the best published one (ARS-LoRA, 29.3%) by 27.2 points.

MT10 is near ceiling, and it is not a win: TOPPO’s 97.2% is below ARS-LoRA (99.5% published, 98.4% in our rerun).

What matters, and when

• LayerNorm on the critic is the largest single step: the mean goes from 78.5% to 87.5% and worst-10 from 6.6% to 42.3%. Alone, it already beats the ARS-LN-1024 rerun’s mean.
• PopArt matters far more when tasks’ reward scales differ more. On V1 rewards its IQM gain over LN-c + FG-c is 17.01 points; on V2 it is 0.32.
• PCGrad on the actor lifts worst-10 from 52.4% to 56.5% and narrows its seed spread. In a separate sweep of eight actor combiners (run with EMA PopArt), none of the other seven differs from PCGrad at p < 0.05: it is one workable choice, not the best one.
• The rungs interact, so read the ladder as one path, not a unique split of the credit.

What this work does not show

• The critic diagnosis is correlational, and how much each module helps changes with the reward version and value clipping. No claim that every module is necessary follows.
• TOPPO costs wall-clock time: 19.8 versus 8.7 RTX 3090 hours per 100M steps for vanilla PPO. Most of it is extracting and combining per-task gradients; forming the Gram matrix and solving FairGrad take about 1%.
• Ported to multi-task SAC, the critic stack gives no reliable gain, so the claim is about PPO.
• In an AMP humanoid walk-steering test, where one critic serves a tracking and a style objective, the full stack cut mean motion-style FID by 76%, and all ten seeds improved, but the character walks more slowly and the task return drops.
• All experiments are state-based, with task IDs and separate actor and critic networks. The MT50 tail is far from solved: worst-5 success is 24.2%.

Code and training logs

The code covers TOPPO and every Meta-World experiment, the SAC and ARS baselines we reran, and the AMP humanoid test. The training logs hold 1,193 runs in 120 groups: the runs behind the paper’s tables and figures, plus related design and follow-up runs. A few timing and diagnostic measurements are not on W&B; the repository documents them separately. No trained checkpoints are released; the repository gives the commands that regenerate them.

Citation

@inproceedings{li2026topporethinkingppomultitask,
  title         = {{TOPPO}: Rethinking {PPO} for Multi-Task Reinforcement Learning with Critic Balancing},
  author        = {Yuanpeng Li and Gefei Lin and Annie Qu and Rui Miao},
  booktitle     = {Advances in Neural Information Processing Systems},
  year          = {2026},
  eprint        = {2605.11473},
  archivePrefix = {arXiv},
  primaryClass  = {cs.AI},
  url           = {https://arxiv.org/abs/2605.11473}
}