🤖 machine learning

Debiased Model-based Representations for Sample-efficient Continuous Control

This paper proposes DR.Q, a debiased model-based representation algorithm that enhances sample efficiency in continuous control by maximizing mutual information between state-action pairs and next states while employing faded prioritized experience replay to mitigate overfitting and representation bias.

Original authors: Jiafei Lyu, Zichuan Lin, Scott Fujimoto, Kai Yang, Yangkun Chen, Saiyong Yang, Zongqing Lu, Deheng Ye

Published 2026-05-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jiafei Lyu, Zichuan Lin, Scott Fujimoto, Kai Yang, Yangkun Chen, Saiyong Yang, Zongqing Lu, Deheng Ye

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 teaching a robot to walk, run, or play basketball. In the world of Artificial Intelligence, this is called Reinforcement Learning. The robot learns by trying things, failing, getting a "score" (reward), and trying again.

The problem is that robots are usually terrible at learning efficiently. They might need to practice for years (millions of steps) just to learn how to walk without falling. This is expensive and slow.

To fix this, researchers use a trick called Model-Based Representations. Think of this as giving the robot a "mental map" or a "dream world." Instead of just reacting to the raw pixels of a camera or raw numbers from sensors, the robot tries to understand the rules of how the world changes. If it moves its leg, what happens next? By learning these rules, it can learn faster.

However, the paper argues that the current way robots build these "mental maps" has two major flaws. The authors, Jiafei Lyu and colleagues, propose a new method called DR.Q (Debiased model-based Representations for Q-learning) to fix these flaws.

Here is how DR.Q works, explained through simple analogies:

The Two Problems with Current Methods

1. The "Fake Alignment" Problem (The Bad Map)
Current methods try to teach the robot to predict the future by making sure its "prediction" looks exactly like the "reality" in terms of numbers.

The Analogy: Imagine you are trying to learn a new language. A bad teacher tells you, "Just make sure your accent sounds exactly like mine." You might mimic the sounds perfectly, but you aren't actually learning the meaning of the words. You could be saying "apple" when you mean "car," but if the sound is close enough, the teacher is happy.
The Paper's Claim: Current AI does this. It minimizes the "distance" between what it predicts and what actually happens, but it doesn't guarantee the robot actually understands the connection. It might be memorizing noise or irrelevant details (like the color of the sky) instead of the important mechanics (like how the legs move).

2. The "Old News" Problem (The Bad Memory)
Robots learn from a "replay buffer," which is like a diary of everything they've ever done.

The Analogy: Imagine a student studying for a test. They have a diary of all their past mistakes.
- Old Method A (Uniform): They read every page of the diary equally, even the boring stuff from last year.
- Old Method B (Prioritized): They only read the pages where they got a terrible grade (big mistakes). This is good, but they keep reading the same terrible mistakes from the beginning of the diary over and over again. They get stuck in the past and don't learn from their recent progress.
The Paper's Claim: Robots get "stuck" on old, bad experiences. They overfit to early failures and ignore new, useful lessons.

The DR.Q Solution

DR.Q fixes these two problems with two clever tricks:

1. The "Deep Connection" Trick (Mutual Information)
Instead of just telling the robot, "Make your prediction look like reality," DR.Q says, "Make sure your prediction and reality are deeply connected."

The Analogy: Instead of just mimicking the teacher's accent, the robot is now forced to understand the relationship between the words. If you say "apple," the robot must understand that the next word is likely "pie" or "tree," not just that the sound is similar.
How it works: The algorithm adds a special math rule (called Mutual Information) that forces the robot's internal "mental map" to capture the most important information about how the world works, ignoring the useless noise. It ensures the map is a true reflection of the rules, not just a cheap copy.

2. The "Fresh & Important" Trick (Faded Prioritized Replay)
DR.Q changes how the robot reads its diary.

The Analogy: Imagine a diary where the pages are weighted by two things:
1. How much you learned from it (Did you make a big mistake? Great, study this!).
2. How new it is (Is this from yesterday or last year?).
How it works: DR.Q uses a "Faded" system. If a mistake is huge, it gets attention. But if that mistake is from a long time ago, its importance "fades" away. This ensures the robot focuses on recent, valuable lessons and stops obsessing over old, irrelevant failures. It balances learning from big mistakes with staying up-to-date.

The Results

The authors tested DR.Q on 73 different robot tasks, from simple walking to complex tasks like a humanoid robot doing a backflip or a dog running.

The Outcome: DR.Q learned faster and performed better than almost all other top methods.
The Analogy: While other robots were still stumbling around trying to figure out how to stand up, DR.Q was already running marathons. In some cases, it was significantly better, sometimes by a huge margin.
One Set of Rules: The best part? They used the exact same settings (hyperparameters) for every single task. They didn't have to tweak the robot's brain for every new game; it just worked.

Summary

The paper introduces DR.Q, a smarter way for robots to learn. It stops robots from memorizing useless details and stops them from obsessing over old mistakes. By forcing a deeper understanding of how the world works and focusing on fresh, important lessons, DR.Q allows robots to learn complex skills much faster and more reliably than before.

Technical Summary: Debiased Model-based Representations for Sample-efficient Continuous Control

Problem Statement

Reinforcement learning (RL) agents often suffer from sample inefficiency, requiring extensive online interactions to learn effective policies. While model-free methods have improved efficiency through various techniques, model-based approaches offer higher sample efficiency by learning environment dynamics but incur high training costs. Recent hybrid approaches attempt to learn model-based representations (embedding latent dynamics into state-action representations) to feed downstream actor-critic networks, combining the benefits of both paradigms.

However, the authors identify two critical sources of bias in existing model-based representation methods that lead to suboptimal performance:

Representation Bias: Existing methods typically minimize the numerical deviation (e.g., Mean Squared Error) between the current state-action representation ( $z_{sa}$ ) and the next state representation ( $z_{s'}$ ). The authors argue that minimizing Euclidean distance does not inherently maximize the mutual information between these representations. Consequently, learned representations may fail to capture sufficient relevant information or may overfit to redundant elements, particularly in high-dimensional spaces.
Sampling Bias: Standard experience replay strategies, such as uniform sampling or Prioritized Experience Replay (PER), are susceptible to primacy bias. PER favors transitions with large temporal difference (TD) errors, which often correspond to early, potentially suboptimal experiences. This causes the agent to overfit to early data that may deviate significantly from the current policy distribution, leading to training instability and inferior performance.

Methodology: The DR.Q Algorithm

To address these issues, the authors propose DR.Q (Debiased model-based Representations for Q-learning), an off-policy RL algorithm designed to learn informative model-based representations. The framework separates the learning process into two phases: (i) learning state and state-action representations using model-based objectives, and (ii) optimizing downstream value functions and policies.

1. Mutual Information Maximization in Representation Learning

DR.Q modifies the encoder training objective to explicitly maximize the mutual information between the state-action representation ( $z_{sa}$ ) and the next state representation ( $z_{s'}$ ), in addition to minimizing their deviation.

Latent Consistency Loss: Minimizes the MSE between the predicted next state representation and the target next state representation.
Mutual Information Loss: To avoid the intractability of direct mutual information calculation, DR.Q employs the InfoNCE loss (a lower bound of mutual information). This loss encourages the representations to be not only numerically close but also to encode shared, informative knowledge about the latent dynamics.
Theoretical Justification: The authors prove (Theorem 4.1) that minimizing MSE does not guarantee increased mutual information. Furthermore, they show (Lemma 4.2) that maximizing mutual information strictly reduces the conditional entropy of the next state given the current state-action, leading to a more deterministic and discriminative latent dynamics model, which tightens the upper bound on value estimation error.

2. Faded Prioritized Experience Replay (Faded PER)

DR.Q introduces a novel sampling strategy to mitigate primacy bias by combining Prioritized Experience Replay (PER) with a forget mechanism.

Mechanism: The sampling probability $P(i)$ for a transition $e_i$ is proportional to the product of its TD error magnitude and a decay factor based on its age in the replay buffer:
$P(i) \propto |\delta(i)|^\alpha (1 - \epsilon)^i$
where $\delta(i)$ is the TD error, $\alpha$ is a smoothing coefficient, and $\epsilon$ is the decay rate.
Effect: This strategy assigns high priority to transitions that are both new (recently added) and informative (large TD errors). It ensures that older experiences, even if they had large errors initially, are gradually deprioritized as the policy evolves, preventing the agent from overfitting to outdated data.
Theoretical Analysis: The authors provide theoretical bounds (Theorem 4.3) showing that under Faded PER, the sampling probability of older experiences with identical TD errors is strictly lower than that of newer ones, and the expected sampling times for old experiences are bounded.

Key Contributions

DR.Q Algorithm: A general off-policy RL algorithm that learns model-based representations using a single set of hyperparameters across diverse continuous control benchmarks.
Debiased Representation Learning: The introduction of an InfoNCE-based mutual information loss to complement standard latent consistency losses, ensuring representations are sufficiently informative and expressive.
Faded Prioritized Experience Replay: A sampling strategy that effectively balances the benefits of prioritizing high-error transitions with the need to forget outdated experiences, thereby mitigating primacy bias.
Comprehensive Evaluation: Extensive empirical validation across 73 tasks from three standard benchmarks (MuJoCo, DMC Suite, and HumanoidBench), demonstrating robustness and generalizability.

Experimental Results

The authors evaluated DR.Q against strong baselines, including model-free methods (e.g., SimBaV2, FoG, TD7) and model-based methods (e.g., MR.Q, TDMPC2, DreamerV3).

Performance: DR.Q consistently matches or outperforms state-of-the-art baselines. Notable improvements include:
- A 15.5% gain over SimBaV2 on DMC-Hard tasks.
- A 58.9% improvement over FoG on HumanoidBench (with dexterous hands) tasks.
- A 26.8% lead over MR.Q on DMC-Visual tasks.
Sample Efficiency: DR.Q achieves superior sample efficiency, often surpassing baselines by a large margin. For instance, it is the first method reported to achieve an average return exceeding 700 on the challenging dog-run task within 1 million environment steps.
Ablation Studies: Experiments confirm that removing the InfoNCE loss or the Faded PER mechanism leads to significant performance degradation, particularly on high-dimensional tasks with redundant state information (e.g., HumanoidBench with hands).
Representation Quality: t-SNE visualizations indicate that DR.Q learns more continuous and concentrated clusters compared to MR.Q, suggesting better structural learning of the latent space.

Significance and Claims

The paper claims that DR.Q represents a concrete step toward building high-performing and general model-free RL algorithms. By explicitly addressing the biases in representation learning (via mutual information maximization) and experience replay (via Faded PER), DR.Q achieves strong performance without requiring task-specific hyperparameter tuning or complex architectural changes.

The authors note limitations, such as inferior performance on specific tasks like Hopper-v4 (attributed to the use of unified hyperparameters) and challenges on extremely difficult visual tasks like humanoid-run within the limited training budget. They also acknowledge that the algorithm is not designed for hard exploration or non-Markovian tasks. Despite these limitations, the work demonstrates that debiasing model-based representations is a viable path to enhancing sample efficiency in continuous control.