PaperlayerView exact revision on arXiv

Source-linked AI summary

Debiased Model-based Representations for Sample-efficient Continuous Control

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

arXiv:2605.11711v1cs.LGcs.AI

TL;DR

Existing model-based representation training can align latent vectors without ensuring that task-relevant information is retained. DR.Q adds mutual-information maximization and faded prioritized replay, and across 73 continuous-control tasks it matched or outperformed strong baselines.

  • Problem

    Minimizing distances between model-based representations does not ensure they retain task-relevant information.

    Source: §4.1 · p. 3

  • Method

    DR.Q maximizes mutual information between state-action and next-state representations while minimizing their deviations, and uses faded prioritized experience replay.

    Source: §1 · p. 2

  • Results

    Across 73 continuous-control tasks, DR.Q matched or outperformed strong baselines, including a 15.5% gain over SimBaV2 on DMC-Hard tasks.

    Source: §5.1 · p. 8, §6 · p. 9

  • Takeaways & Limitations

    The reported results support DR.Q's applicability across diverse continuous-control tasks using a single set of hyperparameters.

    Source: §6 · p. 9

    DR.Q underperformed on Hopper-v4 and failed on visual DMC humanoid-run within the 1M-step evaluation budget; unified hyperparameters may contribute to the Hopper result.

    Source: §6 · p. 9

Abstract

from arXiv · show

Model-based representations recently stand out as a promising framework that embeds latent dynamics information into the representations for downstream off-policy actor-critic learning. It implicitly combines the advantages of both model-free and model-based approaches while avoiding the training costs associated with model-based methods. Nevertheless, existing model-based representation methods can fail to capture sufficient information about relevant variables and can overfit to early experiences in the replay buffer. These incur biases in representation and actor-critic learning, leading to inferior performance. To address this, we propose Debiased model-based Representations for Q-learning, tagged DR.Q algorithm. DR.Q explicitly maximizes the mutual information between the representations of the current state-action pair and the next state besides minimizing their deviations, and samples transitions with faded prioritized experience replay. We evaluate DR.Q on numerous continuous control benchmarks with a single set of hyperparameters, and the results demonstrate that DR.Q can match or surpass recent strong baselines, sometimes outperforming them by a large margin. Our code is available at https://github.com/dmksjfl/DR.Q.

1. Introduction

The introduction identifies information loss and replay-induced primacy bias as problems in model-based representations, then presents DR.Q’s mutual-information objective and faded prioritized replay as remedies and reports strong benchmark performance.

1. Introduction

  • Across 73 tasks on three continuous-control benchmarks, DR.Q matches or outperforms strong domain-specific algorithms and general baselines, sometimes by a large margin.Figure 1 summarizes results from the benchmarks and names SimBaV2, MR.Q, and TDMPC2 among the strong baselines.Source: Figure 1 · p. 1, §1 · p. 2
  • RL agents often need many costly online interactions; prior work modifies model-free methods (Fujimoto et al., 2018; Kuznetsov et al., 2020; Lyu et al., 2022) or uses world models for planning (Hansen et al., 2022; 2024) and data augmentation (Janner et al., 2019; Voelcker et al., 2025), which can cost more to train.The passage also notes replay-buffer reuse (Chen et al., 2021; Lyu et al., 2024) and network-architecture changes (Nauman et al., 2024; Lee et al., 2025a;b) as model-free directions.Source: §1 · p. 1
  • Model-based representations learn latent dynamics in state-action representations and feed them to actor and critic networks, offering richer learning signals and faster adaptation to environmental dynamics.The introduction presents this as a way to bring model-based objectives into model-free algorithms.Source: §1 · p. 1
  • Existing objectives can minimize representation deviation without ensuring mutual information, while uniform sampling or TD-error PER can overfit to early bad experiences through primacy bias (Nikishin et al., 2022).The introduction says these factors bias representation and actor-critic learning and lead to inferior performance.Source: §1 · p. 2
  • DR.Q maximizes mutual information between current state-action and next-state representations alongside minimizing their deviation, and fades replay priorities to favor newer high-TD-error experiences.The replay scheme is intended to retain valuable samples while reducing primacy bias, producing more informative representations for actor-critic learning.Source: §1 · p. 2

2. Related Work

Related work studies dynamics-based representations across model-based and model-free RL, strategies for improving sample efficiency, and replay methods for selecting useful experience. DR.Q focuses on improving model-based representations and combines prioritized sampling with a forget mechanism.

Dynamics-based representation learning

  • Dynamics-based representation learning models system evolution in model-based RL and predicts future latent states in model-free methods, with effectiveness demonstrated across scenarios.The related work also identifies contrastive and self-supervised learning as representation-learning approaches used in RL.Source: §2 · p. 2

Sample-efficient RL algorithms

  • Prior work improves sample efficiency through exploration, replay-buffer reuse, parallel simulation, normalization, value-bias mitigation, model-based methods, and larger network capacities; DR.Q instead improves representations without changing network configurations.Source: §2 · p. 2

Experience replay methods

  • PER prioritizes transitions by TD error, while HER addresses sparse rewards; DR.Q's faded prioritized replay combines PER with forgetting to favor valuable samples.Other replay approaches adjust sampling toward a more uniform distribution, organize experiences into graphs, or lower sampling probabilities for older experiences.Source: §2 · pp. 2–3

3. Preliminary

The section introduces the RL formulation and value function, then explains how model-based representations encode state-action information and learn through latent dynamics consistency.

Reinforcement learning

  • RL is framed as an MDP (S, A, r, γ, P), where the agent learns a policy π: S→A to maximize cumulative discounted return.Qπ(s, a) measures expected return given state s and action a.Source: §3 · p. 3

Model-based representations

  • Model-based representations encode states and state-action pairs in a latent space, using state and state-action encoders plus a reward predictor.The encoders produce zs = f(s) and zsa = g(zs, a), while the reward predictor outputs r̂(zs, a).Source: §3 · p. 3
  • Their training minimizes expected squared deviations between the state-action and next-state representations and between predicted and true rewards.The objective is E[(zsa − zs′)^2 + (r̂ − r)^2].Source: §3 · p. 3

Notations

  • The notation defines H(X) as entropy, H(X|Y) as conditional entropy, and I(X; Y) as mutual information.Source: §3 · p. 3

4. Debiased Model-based Representations

DR.Q addresses representation and replay-sampling biases with mutual-information learning and faded prioritized replay, then uses the learned representations for downstream actor-critic training.

4. Debiased Model-based Representations

  • DR.Q separates learning state and state-action representations with model-based objectives from optimizing downstream value functions and policy.The framework defines mappings for states, state-action pairs, and rewards.Source: §4 · p. 3

4.1. Representation Learning with Mutual Information

  • Minimizing distance between state-action and next-state representations does not necessarily maximize their mutual information, so DR.Q adds an MI loss alongside MSE.The paper argues that distance minimization alone can leave task-relevant information underemphasized; increasing mutual information reduces uncertainty about the next-state representation given the state-action representation.Source: §4.1 · pp. 3–4, A. Missing Proofs · p. 16

4.2. Faded Prioritized Experience Replay

  • Faded PER combines TD-error priorities with experience-age decay, sampling transitions that are both recent and important while retaining high-error older samples that are not too old.The analysis states that equal-TD-error older transitions receive lower sampling probability and that expected sampling times of old experiences are bounded.Source: §4.2 · p. 5

4.3. Algorithm

  • DR.Q combines a mutual-information objective for informative representations with faded prioritized replay that favors valuable recent samples while reducing overfitting to old experiences.The algorithm is designed to debias model-based representations on both the representation-learning and sampling fronts.Source: §4.3 · p. 5

4.3.1. ENCODER TRAINING

  • Encoder training combines two-hot reward cross-entropy, latent-dynamics MSE, and an InfoNCE loss that lower-bounds mutual information between predicted and target next-state representations.The latent-dynamics target comes from a target state encoder; DR.Q uses short-horizon rollouts to train the encoder components.Source: §4.3.1 · pp. 5–6

4.3.2. ACTOR-CRITIC TRAINING

  • For downstream learning, DR.Q trains a deterministic policy with twin critics, using clipped double Q targets, multi-step returns, and Huber loss for the critics.The actor and critic target networks are periodically copied from their corresponding current networks.Source: §4 · p. 3, §4.3.2 · p. 6

5. Experiment

Across 73 continuous-control tasks, DR.Q broadly matched or exceeded baselines, while ablations showed that its InfoNCE loss and faded prioritized replay each contribute to performance.

5. Experiment

  • Experiments covered 73 MuJoCo, DMC, and HumanoidBench tasks, using one DR.Q hyperparameter set across environments and training for 1M steps on MuJoCo and 500K on other tasks.Results were normalized by TD3 scores for MuJoCo, by 1000 for DMC, and by task success scores for HumanoidBench.Source: §5 · p. 7

5.1. Main Results

  • 15.5% over SimBaV2 on DMC-Hard, 58.9% over FoG on HumanoidBench (w/ hand), and 26.8% over MR.Q on DMC-Visual were reported as DR.Q gains; overall, it matched or outperformed baselines across domains.DR.Q consistently outperformed MR.Q across all benchmarks.Source: §5.1 · p. 8
  • DR.Q was the first method reported to exceed 700 average return on dog-run under 1M environment steps, and achieved better sample efficiency than strong baselines on numerous challenging tasks.The comparison included MR.Q and SimBaV2; BRO was omitted because logged results were unavailable for partial tasks.Source: §5.1 · p. 8

5.2. Ablation Study

  • Removing InfoNCE generally reduced performance, especially on high-dimensional HumanoidBench tasks with redundant state information, although DR.Q remained at least as competitive as MR.Q without the loss.The ablation set λm = 0; Figure 4 (top) reports results on four representative tasks.Source: §5.2 · p. 8, Figure 4 · p. 8
  • Faded PER combines TD-error prioritization with sampling-weight decay; excluding either PER or forgetting risks lower performance and sample efficiency, while combining both maximized gains across evaluated tasks.The variants separately removed the forget mechanism or PER; Figure 4 (bottom) shows the ablation results.Source: §5.2 · pp. 8–9, Figure 4 · p. 8

6. Conclusion

DR.Q combines mutual-information maximization and faded prioritized replay to train model-based representations, achieving competitive or better performance across diverse continuous-control benchmarks. The conclusion also notes task-specific limitations and directions for future work.

6. Conclusion

  • Across 73 continuous-control tasks from three standard benchmarks, DR.Q achieves competitive or better performance than prior strong model-free or model-based baselines using one hyperparameter set.The paper describes this as a step toward high-performing, general model-free algorithms.Source: §6 · p. 9
  • DR.Q combines InfoNCE to maximize mutual information between state-action and next-state representations with faded prioritized replay to favor recent, important transitions.The replay mechanism combines prioritized experience replay with a forget mechanism, alongside minimizing representation deviations.Source: §6 · p. 9
  • Future work could evaluate DR.Q on discrete-action tasks such as Atari, explore better model-based representation paradigms, or build a stronger model-free algorithm based on DR.Q.The paper did not evaluate DR.Q on Atari, focusing instead on continuous control because those experiments are expensive.Source: §6 · p. 9
  • DR.Q underperforms on Hopper-v4 and fails on visual humanoid-run within the reported 1M-step budget.The paper notes that unified hyperparameters may contribute to Hopper-v4's inferior performance; all methods failed to achieve meaningful humanoid-run scores, while DrQv2 required 15M steps. InfoNCE and Faded PER may add computation overhead, although DR.Q keeps UTD at 1; the method is not designed for hard-exploration or non-Markovian tasks.Source: §6 · p. 9

A. Missing Proofs · B. Pseudo-code of DR.Q · C. Experiment Setup

The appendix shows that minimizing representation MSE need not increase mutual information, gives bounds for faded prioritized replay, outlines DR.Q’s training loop, and introduces the benchmark and experiment setup.

A. Missing Proofs

  • Minimizing representation MSE does not necessarily increase mutual information: two toy examples show it can decrease or increase, with no definite correlation between the terms.The Gaussian examples show MSE and mutual information moving oppositely as noise grows, and both becoming small near k = 1 or large as k grows.Source: A. Missing Proofs · p. 16
  • Because the state entropy is fixed during optimization, increasing I(Zs′; Zsa) strictly decreases the conditional entropy H(Zs′|Zsa).The proof states this relationship follows from the definition of mutual information.Source: A. Missing Proofs · p. 16
  • With equal TD-error magnitudes, faded prioritized replay gives older transitions higher sampling probability, while expected sample times are positive and below the batch size.The theorem also bounds sampling probabilities and expected sample times using the replay-buffer decay rate.Source: A. Missing Proofs · pp. 16–17

B. Pseudo-code of DR.Q

  • DR.Q initializes policy, critic, predictor, encoder, and target networks, collects transitions, then trains encoders, critics, and actor while updating priorities.Encoder training uses reward, dynamics, and mutual-information losses; replay sampling uses faded PER.Source: B. Pseudo-code of DR.Q · p. 18

C. Experiment Setup

  • The experiment setup section presents benchmark details and the setup for running DR.Q and baseline methods.Source: C. Experiment Setup · p. 18

C.1. Environment Details … D. Full Main Results

The sections specify evaluation environments, comparison baselines, and fixed DR.Q settings, then describe the scope and uncertainty reporting of the full-results comparisons. The supplied passages provide evaluation protocols but no benchmark outcome values.

C.1. Environment Details

  • Gym MuJoCo evaluation uses five commonly adopted v4 tasks for 1M environment steps without action repeat, with returns normalized by random-policy and TD3 scores.The normalization is TD3-Normalized(x) := (x − random score) / (TD3 score − random score).Source: §C.1 · p. 18
  • Proprioceptive DMC evaluation covers 28 tasks, split into 21 DMC-Easy and seven DMC-Hard tasks; visual tasks use three prior observations resized to 84 × 84 RGB pixels.Both DMC settings use action repeat 2 and 500K steps, equivalent to 1M environment steps, with average returns reported directly.Source: §C.1 · pp. 18–19
  • HumanoidBench evaluation includes 14 locomotion tasks without dexterous hands and 14 with hands, extending coverage to tasks with larger observation and action spaces.Returns are normalized using random-policy and task-success scores.Source: §C.1 · pp. 19–20

C.2. Baselines

  • Gym MuJoCo baselines include PPO, TD3+OFE, TQC, REDQ, DroQ, and DreamerV3, with several reported scores sourced from prior papers and averaged across seeds.The passages specify different source papers and seed counts for individual baselines and benchmarks.Source: §C.2 · p. 21
  • The baseline discussion cautions that some methods described as using one hyperparameter set instead use different task configurations or hyperparameters, citing FoG and SimBaV2 as examples.DR.Q is compared against representative strong baselines.Source: §C.2 · p. 22

C.3. Hyperparameters

  • DR.Q keeps its hyperparameters and algorithmic configurations unchanged across all benchmark tasks.The hyperparameters are summarized in Table 5.Source: §C.3 · p. 23, Table 5 · p. 23

D. Full Main Results

  • Full results compare DR.Q with MR.Q, SimBaV2, TDMPC2, and FoG, selected for strong performance across challenging benchmarks; some learning curves are omitted because raw logs or result files were unavailable.The section also says it provides learning curves for DR.Q and some typical algorithms in each domain.Source: D. Full Main Results · p. 24
  • Figures show shaded 95% bootstrapped confidence intervals and tables show them in gray brackets; aggregate mean, median, and IQM intervals use rliable over n × T samples.Here T is the number of evaluated tasks in the benchmark.Source: D. Full Main Results · p. 24

D.1. Gym MuJoCo Results

On Gym MuJoCo, DR.Q beats SimBaV2 in three of five environments, but its average normalized score is slightly lower, mainly because of Hopper-v4. The section reports final performance at 1M environment steps and learning curves across 10 seeds.

D.1. Gym MuJoCo Results

  • Poor performance on Hopper-v4 mainly explains DR.Q’s lower average score; MR.Q and FoG also perform poorly there.Source: §D.1 · p. 25
  • DR.Q outperforms SimBaV2 in 3 of 5 environments, but its average normalized score is slightly lower.The full comparison reports final average performance at 1M environment steps; aggregate mean, median, and IQM use TD3-normalized scores.Source: Table 6 · p. 25, §D.1 · p. 25
  • Gym MuJoCo learning curves show average episode returns across 10 random seeds, with shaded 95% bootstrap confidence intervals.Source: Figure 5 · p. 25

D.2. DMC Suite Easy Results … D.5. HumanoidBench (w/ Hand) Results

These sections present final performance comparisons and full learning curves for DMC-Easy, DMC-Hard, and HumanoidBench with and without dexterous hands. The captions specify evaluation steps, aggregate-score conventions, and 95% bootstrap confidence intervals, but do not provide numerical results or identify winners.

D.2. DMC Suite Easy Results

  • DMC-Easy learning curves show average returns by environment with shaded 95% bootstrap confidence intervals.The solid lines denote average return.Source: Figure 6 · p. 26
  • DMC-Easy results compare final performance at 500K steps (1M environment steps with action repeat 2), with aggregate mean, median, and IQM in units of 1k.Bracketed values are 95% bootstrap confidence intervals.Source: Table 7 · p. 27

D.3. DMC Suite Hard Results

  • DMC-Hard results compare final average returns at 500K steps (1M environment steps with action repeat 2), reporting aggregate mean, median, and IQM in units of 1k.Bracketed values are 95% bootstrap confidence intervals.Source: Table 8 · p. 28
  • DMC-Hard learning curves show average returns by environment with shaded 95% bootstrap confidence intervals.The solid lines denote average return.Source: Figure 7 · p. 28

D.4. HumanoidBench (w/o Hand) Results

  • HumanoidBench without hands reports final average returns at 1M environment steps, equivalent to 500K steps with action repeat 2.Aggregate mean, median, and IQM use success-normalized scores; bracketed values are 95% bootstrap confidence intervals.Source: Table 9 · p. 29
  • HumanoidBench without hands includes full learning curves showing task-level average returns with light-colored 95% bootstrap confidence intervals.Source: Figure 8 · p. 30

D.5. HumanoidBench (w/ Hand) Results

  • HumanoidBench with dexterous hands reports final average returns at 1M environment steps, equivalent to 500K steps with action repeat 2.Aggregate mean, median, and IQM use success-normalized scores; bracketed values are 95% bootstrap confidence intervals.Source: Table 10 · p. 31
  • HumanoidBench with hands includes full learning curves showing task-level average returns with light-colored 95% bootstrap confidence intervals.Source: Figure 9 · p. 32

D.6. Visual DMC Suite Results

On 12 selected visual DMC tasks, DR.Q generally outperforms MR.Q and TDMPC2 in sample efficiency and final performance, especially on visual-dog-stand and visual-hopper-stand.

D.6. Visual DMC Suite Results

  • The experiments select 12 visual DMC tasks and omit simpler tasks where MR.Q or baseline performance has already saturated.Examples of omitted tasks include visual-cartpole-balance and visual-walker-stand.Source: §D.6 · p. 33
  • DR.Q generally outperforms MR.Q and TDMPC2 in sample efficiency and final performance, especially on visual-dog-stand and visual-hopper-stand.The reported results cover 10 seeds; Table 11 compares average returns at 1M environment steps, equivalent to 500K steps with action repeat 2.Source: Table 11 · p. 33, §D.6 · p. 33, Figure 10 · p. 33
  • Table 11 reports aggregate mean, median, and interquartile mean using success-normalized scores, with bracketed values denoting 95% bootstrap confidence intervals.The table reports average returns at 1M environment steps, equivalent to 500K steps due to action repeat 2.Source: Table 11 · p. 33

E. Additional Experiments

The additional experiments, omitted from the main paper for space, evaluate variants and baselines using the Section C setup over 10 seeds and 1M environment steps.

Experimental setup

  • All additional variants and baselines use the Section C setup and run for 10 seeds and 1M environment steps.These experiments were omitted from the main paper due to space constraints.Source: E. Additional Experiments · p. 34

InfoNCE loss ablation

  • Figure 11 presents an extended ablation study of the InfoNCE loss, with results averaged across 10 seeds.Source: Figure 11 · p. 34

Sampling-strategy ablation

  • Figure 12 presents an extended sampling-strategy ablation, reporting average returns across 10 seeds with shaded 95% confidence intervals.Source: Figure 12 · p. 34

E.1. Extended Ablation Study … E.4. Visualization Results of Sampling Strategies

The extended analyses report that DR.Q’s InfoNCE and sampling mechanisms matter for performance, and compare DR.Q with MR.Q under modified hyperparameters, noisy inputs, and representation and sampling visualizations.

E.1. Extended Ablation Study

  • In almost all evaluated tasks, DR.Q outperforms DR.Q (w/o InfoNCE) in sample efficiency and final performance; excluding InfoNCE especially hurts h1-sit-hard-v0 and h1hand-pole-v0.Source: §E.1 · pp. 34–35
  • Removing LAP or the forget mechanism can hurt performance, especially on dexterous-hand HumanoidBench tasks; removing LAP severely hurts humanoid-run and humanoid-walk.Source: §E.1 · p. 35
  • Removing the latent dynamics loss generally reduces sample efficiency and final performance, though its removal has only minor influence on some tasks such as acrobot-swingup.The loss’s effect varies by task; DR.Q without it can significantly underperform on h1hand-stair-v0 and h1hand-pole-v0.Source: §E.1 · p. 35

E.2. Comparison of DR.Q against MR.Q with Modified Hyperparameters

  • Scaling MR.Q’s encoder generally improves its performance, but it still trails DR.Q in numerous environments, supporting the role of InfoNCE and faded PER.The comparison sets MR.Q’s hyperparameters to DR.Q’s modified settings, including encoder learning rate and hidden dimension.Source: §E.2 · p. 36, Figure 14 · p. 36

E.3. On the Mutual Information Loss

  • With 50-dimensional Gaussian noise appended, MR.Q generally loses performance and sample efficiency while DR.Q is less affected and does not degrade on HumanoidBench tasks.The authors note that the injected noise is smaller than the 100-dimensional redundancy introduced by dexterous hands.Source: §E.3 · pp. 36–37, Figure 15 · p. 37
  • After 200K training steps, MR.Q often shows discontinuous representation clusters, while DR.Q shows continuous, concentrated clusters and lacks MR.Q’s void areas on h1hand-stand-v0.The visualizations use t-SNE on 5,000 replay-buffer samples from four tasks.Source: §E.3 · p. 37, Figure 16 · p. 38

E.4. Visualization Results of Sampling Strategies

  • LAP alone can favor old transitions with large TD errors, whereas DR.Q’s forget mechanism directs sampling toward recent valuable transitions so old samples do not outrank newer ones.The comparison records TD errors for the most recent 100K samples during 200K-step runs on eight continuous-control tasks.Source: §E.4 · p. 37, Figure 17 · p. 38

Related papers

Loading 2605.11711v1…