Scaling Multitask RL: Is Representation Learning the Real Secret Sauce?
Representation Learning Enables Scalable Multitask Deep Reinforcement Learning
This paper introduces MR.Q, a model-free Reinforcement Learning (RL) algorithm augmented with auxiliary predictive objectives, to address the challenge of scaling multitask deep RL. By isolating representation learning from planning, MR.Q outperforms the state-of-the-art world-model-based method "Newt" across 200 diverse continuous control tasks in the MMBench suite, achieving superior sample and wall-clock efficiency.
TL;DR
A common assumption in deep Reinforcement Learning is that scaling to hundreds of tasks requires "World Models" and complex planning. This paper challenges that dogma. By using a model-free agent called MR.Q, which uses predictive objectives only to shape its internal representations (and not for planning), the authors outperformed the reigning world-model champion, Newt, across 200 tasks while being significantly faster and more computationally efficient.
The "Planning" Tax vs. The "Representation" Gap
Recent SOTA methods like DreamerV3 and TD-MPC2 have shown that model-based RL can scale. However, these methods come with a heavy "planning tax": they require latent rollouts and trajectory optimization, which are computationally expensive and sensitive to hyperparameters.
The authors of this paper noticed a critical bottleneck: Scaling model capacity in standard RL (like PPO or TD3) often fails to improve performance. As shown in the study, simply making a PPO model larger on a task like Humanoid results in flat performance. The problem isn't the number of parameters; it's that the reward signal alone is too sparse and non-stationary to train high-capacity encoders effectively, leading to "representation collapse" and "dormant neurons."
Methodology: MR.Q — Planning-Free Predictive Learning
The core insight is that the representations learned by world models are what actually matter, not the planning they perform.
MR.Q (Model-based Representations for Q-learning) takes a model-free approach (based on TD3) and adds auxiliary heads to predict:
- Latent Dynamics: What will the next latent state look like?
- Reward: What reward will be received?
- Termination: Will the episode end ()?
These predictions are used solely to backpropagate gradients into the shared encoder . The actor and critic then act on these high-quality features. No imagination is used.
Figure 1: Representation quality facilitates scaling. Adding Model-Based (MB) Representations allows PPO to actually benefit from increased model size, whereas standard PPO plateaus.
Experimental Battleground: MMBench
The authors tested MR.Q against Newt (the multitask version of TD-MPC2) on MMBench, which includes 200 tasks spanning MuJoCo, MetaWorld, ManiSkill, and even Atari games.
Key Discoveries:
- Sample Efficiency: At the 2M step mark, MR.Q achieved a normalized score 37% higher than Newt.
- Wall-Clock Speed: Because MR.Q doesn't waste time "thinking" (planning) during training and inference, it hits high performance levels hours earlier than model-based methods.
- Zero-Shot Transfer: MR.Q's representations are more general. When moved to 28 held-out tasks, it showed a 50% relative advantage in zero-shot performance compared to Newt.
Figure 2: Performance across suites. MR.Q (teal) consistently matches or beats Newt (red) in both sample efficiency and final score.
Deep Dive: Why does it work?
In the ablation study, the authors compared MR.Q to a standard TD3 (encoder-free). They used PCA (Principal Component Analysis) and SRank to measure the "health" of the representations.
- High Rank: MR.Q maintained a high-rank latent manifold, meaning it used its capacity to store diverse information.
- Healthy Neurons: MR.Q had a much lower fraction of "dormant neurons" (inactive units) compared to standard RL, indicating that the predictive loss acts as a regularizer that keeps the network "plastic" and ready to learn.
Figure 3: PCA projections show that MR.Q (left) learns well-separated, high-dimensional task clusters, while model-free RL without predictive objectives (right) suffers from feature collapse.
Critical Analysis & Conclusion
This paper provides a refreshing "Bitter Lesson" for RL: complex algorithmic components like online planning might be unnecessary if we can master representation learning.
Takeaway for Practitioners: If you are trying to scale RL, don't just add more layers or complex planning loops. Focus on auxiliary predictive losses that force your encoder to understand the dynamics of the environment.
Limitations: The study focuses on continuous control. It remains to be seen if this "representation-only" approach holds up in long-horizon tasks (like complex navigation or strategy games) where the search-based benefits of planning usually shine.
Future Work: The hybrid approach—combining MR.Q's efficient training with selective, high-level imagination—could be the next frontier for general-purpose robotic agents.