🤖 AI

Performance Asymmetry in Model-Based Reinforcement Learning

This paper reveals a critical "Performance Asymmetry" in current Model-Based Reinforcement Learning where agents excel on some Atari tasks while failing on others, and proposes a novel Joint Embedding DIffusion (JEDI) world model that resolves this imbalance to achieve state-of-the-art performance across both human-preferred and agent-preferred task subsets.

Original authors: Jing Yu Lim, Rushi Shah, Zarif Ikram, Samson Yu, Haozhe Ma, Tze-Yun Leong, Dianbo Liu

Published 2026-02-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jing Yu Lim, Rushi Shah, Zarif Ikram, Samson Yu, Haozhe Ma, Tze-Yun Leong, Dianbo Liu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a talent scout looking for the ultimate video game champion. You have a list of 26 different Atari games, ranging from simple ones like Pong to complex ones like Bank Heist.

Recently, a new generation of AI agents (let's call them "Robo-Players") has been trained using a method called Model-Based Reinforcement Learning (MBRL). When you look at their average score across all 26 games, they seem incredible—they are beating human players! The headlines say, "Robots are now superhuman!"

But this paper, titled "Performance Asymmetry," argues that this headline is misleading. It's like saying a student is a "genius" because they got 100% on a few easy math quizzes but failed every single history exam.

Here is the breakdown of what the authors found, using simple analogies:

1. The "Average" Trap (The Hidden Problem)

The researchers discovered a phenomenon they call Performance Asymmetry.

The Situation: The Robo-Players are absolutely crushing the "easy" games (which the authors call Agent-Optimal tasks). In some of these, they are 20 times better than humans.
The Catch: However, in the "hard" games (called Human-Optimal tasks), these same robots are terrible. They are often performing worse than a random person just mashing buttons.
The Analogy: Imagine a chef who is a world-class master at making perfect, simple toast. They can make toast so good it's 20 times better than anyone else. But if you ask them to cook a complex, multi-course gourmet dinner, they burn the food and serve you raw eggs. If you just look at their "average" cooking score, they might still look like a 5-star chef because the toast scores were so high. But in reality, they are a one-trick pony.

2. Why the Old Metrics Failed

The standard way to measure these AI agents is by taking the Arithmetic Mean (the simple average).

The Flaw: The arithmetic mean is easily fooled by extreme numbers. If an agent gets a "100" on a few easy games, it can drag the average up so high that it hides the fact that they got "0" on the hard games.
The Solution: The authors propose a new metric called Sym-HNS. Think of this like a GPA that punishes failing grades. In this new system, if you fail even one subject, your overall grade drops significantly, no matter how good you are at the others. This forces the AI to be balanced rather than just extreme in one direction.

3. Why Are the Robots Failing the Hard Games?

The authors dug into why the robots are failing the "Human-Optimal" games. They found two main culprits:

The "Pixel" Problem: The best-performing robots so far look at the game screen like a human does (pixel by pixel). This is like trying to learn to drive a car by staring at a high-definition photo of the road. It's too much information! The "hard" games have complex rules and many possible moves (like shooting a bomb that explodes later). Processing all those pixels makes the robot's brain get overwhelmed (the "Curse of Dimensionality").
The "Shooter" Problem: Many of the hard games involve shooting or timing actions (like Bank Heist where you throw a bomb and have to run away before it explodes). These games have high "variance"—one wrong move leads to instant disaster. Robots looking at raw pixels struggle to predict these complex cause-and-effect chains.

4. The New Solution: JEDI

To fix this, the authors built a new AI agent called JEDI (Joint Embedding DIffusion).

How it works: Instead of staring at every single pixel (like looking at a high-res photo), JEDI learns to see the game in "compressed concepts" (like looking at a map or a summary). It understands the essence of the game state without getting bogged down in visual noise.
The Magic: They combined this "conceptual view" with a Diffusion Model (a type of AI usually used to generate art). This allows the robot to imagine future scenarios and plan ahead, even in complex, chaotic games.
The Result: JEDI is the first agent to fix the "one-trick pony" problem.
- It is excellent at the complex "Human-Optimal" games (finally beating humans at Bank Heist).
- It remains competitive at the "easy" games.
- It does all this while using less computer memory and running faster than the previous champions.

Summary

The paper teaches us a valuable lesson: Don't just look at the average.

If an AI is great at some things and terrible at others, it isn't truly "intelligent" yet; it's just overfitting to specific tasks. The authors showed that by changing how we measure success (using Sym-HNS) and changing how the AI sees the world (using latent diffusion instead of raw pixels), we can build robots that are truly well-rounded, balanced, and ready for the real world.

In short: They stopped the AI from being a "Toast Master" and turned it into a "Gourmet Chef" who can handle both simple and complex dishes.

1. Problem Statement: Performance Asymmetry

The paper identifies a critical flaw in the evaluation of Model-Based Reinforcement Learning (MBRL) agents on the Atari100k benchmark. While recent State-of-the-Art (SOTA) agents (e.g., DIAMOND) achieve a mean Human-Normalized Score (HNS) greater than 1.0 (suggesting super-human performance), this aggregate metric masks a severe Performance Asymmetry.

The Phenomenon: MBRL agents dramatically outperform humans on specific tasks (termed Agent-Optimal tasks) but drastically underperform humans on others (termed Human-Optimal tasks).
The Disparity: The SOTA pixel-based agent (DIAMOND) showed a 21× performance gap between these two subsets. It scored the worst among all baselines on Human-Optimal tasks despite having the highest overall mean HNS.
Root Causes of Masking:
1. Arithmetic Mean: Standard evaluation uses the arithmetic mean, which is invariant to extreme values. High scores on a few tasks can mathematically compensate for near-zero performance on many others.
2. HNS Range: The HNS metric has a large value range (from negative to >8). Small improvements in high-scoring tasks dominate the aggregate, rendering significant progress (or regression) in low-scoring tasks invisible.
Task Characteristics: The authors found that Human-Optimal tasks generally have:
- Larger action spaces (increasing the search space exponentially).
- A higher prevalence of "Shooter" tasks (games requiring ranged weapons/FIRE actions), which introduce high variance in outcomes and complex timing dependencies.
Pixel vs. Latent: The asymmetry is most pronounced in pixel-based agents. These agents struggle with Human-Optimal tasks due to the curse of dimensionality (processing raw $64\times64\times3$ images) and the complexity of Shooter tasks. Conversely, they excel in Agent-Optimal tasks (like Breakout) due to the high-fidelity generation capabilities of pixel diffusion models.

2. Methodology

A. New Evaluation Framework

To expose and quantify this asymmetry, the authors propose:

Task Partitioning: They split the 26 Atari100k tasks into two balanced subsets (13 tasks each) based on the average performance of four existing baselines:
- Agent-Optimal (AO): Tasks where the average agent HNS > 75%.
- Human-Optimal (HO): Tasks where the average agent HNS < 75%.
Sym-HNS Metric: They introduce a new aggregate metric, Sym-HNS, calculated as the harmonic mean of the mean HNS on AO and HO subsets:
$\text{Sym-HNS} = \frac{2 \cdot \mu_{\tau}(AO) \cdot \mu_{\tau}(HO)}{\mu_{\tau}(AO) + \mu_{\tau}(HO)}$
Unlike the arithmetic mean, the harmonic mean penalizes extreme imbalances, ensuring that an agent cannot achieve a high score by neglecting one subset.

B. Proposed Solution: JEDI (Joint Embedding DIffusion)

To address the asymmetry, the authors propose JEDI, a novel latent end-to-end Joint Embedding Diffusion world model.

Architecture:
- Encoder ( $E_\phi$ ): Compresses high-dimensional pixel observations ( $x_t$ ) into a low-dimensional latent space ( $z_t$ ).
- Latent Diffusion Dynamics Model ( $D_\theta$ ): Predicts future latent states ( $z_{t+1}$ ) conditioned on current states and actions. It uses a single-step conditioned denoising operator.
- Reward/Termination Model ( $R_\psi$ ): Predicts rewards and episode termination from latent states.
Training Objective (The Core Innovation):
- JEDI adopts a Joint Embedding Predictive Architecture (JEPA) style framework but integrates it directly with Diffusion.
- Loss Function: The encoder is trained end-to-end using the gradients from the diffusion denoising loss (specifically the EDM loss).
- No Reconstruction: Crucially, JEDI does not use a decoder for reconstruction loss, nor does it use hidden state extraction or distillation. The encoder learns representations solely to minimize the diffusion prediction error.
- Stability Mechanisms:
  - Stop-Grad: Applied to the target latent to prevent representation collapse.
  - Tanh Clamp: A differentiable clamp function ( $\tanh(z/s)$ ) is applied to latent outputs to stabilize the diffusion process in the latent space.
  - Random Switching: During training, the model randomly switches between using the predicted next state or the ground truth latent as input for the subsequent time step's loss to balance horizon information.
Policy Training: The policy is trained via an Actor-Critic framework using REINFORCE to avoid backpropagating through the multi-step diffusion process, ensuring computational efficiency.

3. Key Contributions

Discovery of Performance Asymmetry: First to systematically identify and quantify the massive performance gap between Agent-Optimal and Human-Optimal tasks in MBRL, highlighting the issue of task overfitting.
Sym-HNS Metric: Proposed a balanced evaluation metric using the harmonic mean to prevent the masking of poor performance in specific task subsets.
JEDI World Model: Introduced the first end-to-end latent diffusion world model that trains the encoder directly via diffusion gradients within a JEPA-style framework.
Reversal of Asymmetry: Demonstrated that moving from pixel-based to latent-based diffusion models, combined with JEPA-style training, can simultaneously improve performance on complex Human-Optimal tasks and maintain competitiveness on Agent-Optimal tasks.

4. Results

Sym-HNS Performance: JEDI achieves SOTA Sym-HNS, reversing the worsening asymmetry trend seen in previous SOTA agents (like DIAMOND).
Human-Optimal Tasks: JEDI significantly outperforms all baselines on HO tasks, achieving super-human performance on BankHeist (previously a failure point for pixel agents) and improving scores on other Shooter tasks.
Agent-Optimal Tasks: JEDI remains competitive on AO tasks. Notably, it achieves SOTA on Breakout (HNS 5.35 vs. DIAMOND's 4.54), proving that latent diffusion can handle high-visual-detail tasks without pixel inputs.
Efficiency:
- Memory: Uses 43% less GPU memory than DIAMOND.
- Speed: Achieves 3× faster sampling and 2× faster training due to latent compression ( $12\times$ reduction in state dimension).
Robustness: JEDI outperforms latent non-diffusion agents in stochastic environments (random frameskipping).
Ablation Studies: Confirmed that the end-to-end diffusion gradients are essential for HO performance, and that the random switching and latent clamping mechanisms are critical for stability.

5. Significance

This paper fundamentally challenges the current paradigm of MBRL evaluation and architecture design:

Evaluation Shift: It argues that "average super-human" scores are misleading and that the community must adopt balanced metrics (like Sym-HNS) to ensure agents are robust across diverse task complexities, not just overfitting to specific game mechanics.
Architectural Insight: It demonstrates that latent diffusion models can overcome the curse of dimensionality inherent in pixel-based RL, particularly for complex tasks with high action spaces and stochastic dynamics.
JEPA + Diffusion Synergy: It proves that off-the-shelf diffusion losses are directly compatible with JEPA frameworks, enabling efficient, end-to-end representation learning without the computational overhead of reconstruction decoders.
Future Direction: The work opens new avenues for understanding the fundamental differences between tasks humans excel at versus those agents excel at, suggesting that future RL research must explicitly address task complexity and action space variance.