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 representations can miss task-relevant information and overfit replay-buffer experiences. DR.Q combines mutual-information representation learning with faded prioritized replay, matching or surpassing strong baselines across 73 continuous-control tasks, including a 15.5% gain over SimBaV2 on DMC-Hard.

  • Problem

    Existing representation learning can falsely align representations while underemphasizing task-relevant information, limiting sample-efficient reinforcement learning.

  • Method

    DR.Q maximizes mutual information between state-action and next-state representations while using faded prioritized replay to emphasize recent, valuable experiences.

  • Results

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

  • Takeaways & Limitations

    DR.Q shows broad applicability across diverse continuous-control domains using a single set of hyperparameters.

  • Takeaways & Limitations

    DR.Q performs worse on some tasks, fails on challenging visual humanoid-run, and may incur extra computation overhead.

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

DR.Q addresses information loss and bias in model-based representations by maximizing mutual information alongside representation alignment, achieving strong performance across 73 continuous-control environments.

  • Background: Model-based representations enrich model-free actor-critic learning by encoding latent environmental dynamics for richer learning signals and faster adaptation.They train state and action representations by modeling latent dynamics, then feed them to downstream actor and critic networks.
  • Problem: Existing objectives that minimize deviation between current state-action and next-state representations may not increase their mutual information, limiting captured information.This shortcoming is identified as a factor negatively affecting model-based representations.
  • Contribution: DR.Q maximizes mutual information between current state-action and next-state representations while also minimizing their deviations.This encourages representations to be both numerically close and informative about each other.
  • Evaluation: 73 environments across MuJoCo, DMC suite, and HumanoidBench were used to evaluate DR.Q on standard continuous-control online RL benchmarks.The tasks span diverse characteristics and varying complexities.
  • Results: DR.Q can match or outperform strong domain-specific algorithms and general baselines, sometimes by a large margin.The study also releases code, model weights, and logs to facilitate future research.

2. Related Work

Prior work develops representations that capture data patterns and system dynamics for reinforcement learning, while other efforts target sample efficiency and more effective replay-buffer sampling. Dynamics-based representation learning has demonstrated advantages across varied scenarios, and prioritized experience replay improves transition utilization using TD errors.

  • Dynamics-based representation learning: Dynamics-based representation learning learns intermediate features and latent models of how systems evolve from current states under legal actions.It appears in both general representation-learning research and model-based reinforcement learning.
  • Dynamics-based representation learning: Contrastive and self-supervised methods are established approaches for producing high-quality representations in reinforcement learning.The cited examples include contrastive methods and self-supervised methods.
  • Dynamics-based representation learning: Dynamics-based representation learning has demonstrated effectiveness and advantages across various reinforcement-learning scenarios.The passage attributes this evidence to numerous model-based reinforcement-learning works.
  • Sample-efficient RL algorithms: Sample efficiency is a key metric for online reinforcement learning, with higher efficiency meaning agents learn faster and better under a fixed online-interaction budget.Existing efforts include improving agents’ exploration ability.
  • Experience replay methods: Prioritized experience replay assigns transition priorities according to TD errors instead of treating all replay-buffer transitions equally.PER was introduced to better utilize gathered samples and has inspired subsequent work.

3. Preliminary

This section introduces reinforcement learning as an MDP-based policy optimization problem and defines the value function used to measure expected returns. It also describes model-based representations, which enforce latent-space dynamics consistency through learned state, state-action, and reward components.

  • Reinforcement learning (RL): RL problems are formulated as Markov Decision Processes specified by state space S, action space A, reward r, discount factor γ, and dynamics function P.An RL agent learns a policy π : S → A that maximizes cumulative discounted return.
  • Reinforcement learning (RL): The value function Qπ(s, a) measures the expected cumulative discounted return given state s and action a.The definition conditions the return on s0 = s and a0 = a.
  • Model-based representations: Model-based representations learn implicit state-action or state representations by enforcing dynamics consistency in latent space.They leverage objectives from model-based reinforcement learning.
  • Model-based representations: The representation framework trains a state encoder f(·), state-action encoder g(·), and reward function ˆr(·), with zs = f(s) as the state representation.The state encoder receives the state as input and outputs its representation.
  • Model-based representations: For prediction step i, ˆzs′,i is a linear mapping of the state-action representation zsa,i, while ˜zs′,i is produced by the target state encoder.The notation applies for i ∈[1, H].

4. Debiased Model-based Representations

DR.Q debiases model-based representations by maximizing mutual information between state-action and next-state representations alongside deviation minimization, while using faded prioritized experience replay to limit sampling bias. It separates representation learning from downstream value-function optimization and implements the mutual-information objective with InfoNCE.

  • Framework: DR.Q separates training into model-based learning of state and state-action representations, followed by downstream value-function optimization.The method trains mappings for state representations, state-action representations, and reward prediction.
  • Mutual-information debiasing: Minimizing Euclidean deviation between state-action and next-state representations does not necessarily increase their mutual information.Theorem 4.1 identifies numerical closeness as insufficient to guarantee informative latent dynamics.
  • Mutual-information debiasing: DR.Q adds a mutual-information loss to the MSE objective, reducing H(Zs′|Zsa) as I(Zs′; Zsa) increases.Theorem 4.2 states that increasing mutual information strictly reduces the conditional entropy of the next-state representation given the state-action representation.
  • Faded prioritized replay: Faded PER combines TD-error prioritization with experience forgetting, favoring transitions that are both recent and have large TD errors.This addresses primacy bias from overfitting to old experiences while avoiding the assumption that every recent transition is valuable.
  • Mutual-information debiasing: DR.Q uses InfoNCE as a tractable lower-bound objective for its mutual-information loss in high-dimensional representation spaces.Direct mutual-information calculation between next-state representations and predictions is intractable.

5. Experiment

DR.Q is evaluated on 73 continuous-control tasks spanning MuJoCo, DMC, and HumanoidBench using one unchanged hyperparameter set, and it matches or outperforms strong baselines across domains. Ablations show that both InfoNCE and the combined faded-PER components are important for performance and sample efficiency.

  • Experimental setup: DR.Q is evaluated on 73 tasks across MuJoCo, DMC, and HumanoidBench, covering simple low-dimensional through high-dimensional locomotion settings.DMC includes both proprioceptive and visual-control tasks.
  • Experimental setup: DR.Q uses a single hyperparameter set without algorithmic changes across all environments and benchmarks.Agents are trained for 1M steps on MuJoCo and 500K steps on other tasks.
  • Performance summary: DR.Q matches or outperforms selected domain-specific and general baselines across domains, including a 15.5% gain over SimBaV2 on DMC-Hard tasks and a 58.9% improvement against FoG on HumanoidBench (w/ hand) tasks.The comparison summarizes normalized average returns against DR.Q.
  • The InfoNCE loss: Removing the InfoNCE loss generally worsens performance, particularly on high-dimensional HumanoidBench tasks with redundant state information.The ablation sets λm = 0.
  • Faded PER: Excluding either prioritized replay or the forget mechanism can degrade performance and sample efficiency, whereas combining both yields the maximum gains across evaluated tasks.The variants are DR.Q (only forget) and DR.Q (only LAP).

6. Conclusion

DR.Q is presented as a general off-policy algorithm for model-based representations across diverse continuous-control tasks with one hyperparameter set. Its debiasing combines mutual-information maximization with transition reweighting, but performance remains limited on some tasks, especially challenging visual benchmarks.

  • Conclusion: DR.Q trains model-based representations for diverse continuous-control tasks using a single set of hyperparameters.The paper describes DR.Q as a general off-policy reinforcement-learning algorithm.
  • Conclusion: DR.Q introduces InfoNCE to maximize mutual information between state-action and next-state representations while minimizing their deviations.This is the first stated component of DR.Q’s debiasing approach.
  • Conclusion: DR.Q increases the frequency of recent and important transitions to debias model-based representation learning.The passage identifies this transition-sampling strategy as the second component of DR.Q’s debiasing approach.
  • Limitations: DR.Q performs worse on Hopper-v4, potentially because unified hyperparameters are used across all benchmarks.The passage identifies this as a limitation despite DR.Q’s strong performance on numerous challenging tasks.
  • Limitations: DR.Q fails on visual humanoid-run, although all methods reportedly fail to achieve meaningful scores on this task.DrQv2 requires 15M environment steps to achieve meaningful performance on visual-humanoid-run, whereas the reported DR.Q and baseline runs were shorter.

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

The appendices show that DR.Q’s MSE and mutual-information objectives are not definitively correlated, while its faded prioritized replay has bounded sampling behavior. The algorithm combines these components in encoder, critic, and actor updates, and the experiments benchmark DR.Q against baseline methods.

  • A. Missing Proofs: When σ2 becomes large, the MSE term increases while mutual information decreases; for k →∞, both terms become large.Conversely, when k approaches 1, both the MSE term and mutual information become small.
  • A. Missing Proofs: The proofs conclude that minimizing E[∥Zsa − Zs′∥2] does not necessarily increase I(Zsa; Zs′).Toy constructions show that minimizing the MSE term can either minimize or increase mutual information.
  • A. Missing Proofs: Increasing I(Zs′; Zsa) strictly reduces H(Zs′|Zsa), because H(Zs′) is an environment property unchanged by optimization.The lemma links mutual-information maximization to lower conditional entropy.
  • A. Missing Proofs: Faded prioritized replay samples equally erroneous transitions with higher probability when they are earlier, and bounds each transition’s expected sample times below N.The theorem states 0 < E[ni] ≤ N/(1+C(1−ϵ)^(k−i)) < N.
  • B. Pseudo-code of DR.Q: DR.Q initializes policy, twin critic, predictor, state encoder, state-action encoder, their target networks, and an empty replay buffer before interacting with the environment.Actions use the target policy with clipped Gaussian exploration noise and are clipped to [−1, 1].
  • B. Pseudo-code of DR.Q: The training loop samples transitions with faded PER and updates encoders using weighted reward, dynamics, and mutual-information losses.The encoder objective is LDR.Q = Σt=1 λrLreward + λdLdynamics + λmLI.
  • C. Experiment Setup: The setup section presents benchmark details and experimental procedures for running DR.Q and baseline methods.The supplied passage identifies the section’s purpose but does not specify individual benchmarks or settings.

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

The paper evaluates DR.Q across standardized MuJoCo, DMC, and HumanoidBench settings, compares it with representative strong baselines, fixes its hyperparameters across tasks, and reports confidence intervals for performance aggregates.

  • C.1. Environment Details: Gym MuJoCo evaluation uses five common v4 locomotion environments, 1M environment steps, no action repeat, and TD3-normalized returns.Normalization follows TD3-Normalized(x) := (x − random score) / (TD3 score − random score).
  • C.1. Environment Details: DMC evaluation covers 28 proprioceptive tasks split into DMC-Easy and DMC-Hard, plus visual tasks using stacked 84 × 84 RGB observations.DMC experiments use action repeat 2 and 500K steps, equivalent to 1M environment steps, while reporting average returns directly.
  • C.1. Environment Details: HumanoidBench evaluates Unitree H1 whole-body manipulation and locomotion tasks with and without dexterous hands, using 500 steps, action repeat 2, and success-normalized returns.The benchmark therefore corresponds to 1M environment steps, and normalization uses random and task-success scores.
  • C.2. Baselines: The comparison includes model-free, model-based, representation-learning, architecture, and sample-efficiency baselines spanning PPO, TD3+OFE, TQC, REDQ, DroQ, DreamerV3, DrQ-v2, TD7, TDMPC2, CrossQ, iQRL, BRO, MAD-TD, MR.Q, SimBa, SimBaV2, and FoG.The paper emphasizes MR.Q, SimBaV2, TDMPC2, and FoG as main baselines because of their strong performance across challenging benchmarks.
  • C.2. Baselines: The study notes that several methods claiming general-purpose use actually change algorithmic configurations or hyperparameters across tasks, motivating explicit comparison with DR.Q.This concern is specifically illustrated by FoG and SimBaV2.
  • C.3. Hyperparameters: DR.Q keeps both its hyperparameters and algorithmic configurations fixed across all benchmark tasks.The paper summarizes these fixed settings in Table 5.
  • D. Full Main Results: Full results compare DR.Q with extended prior algorithms and provide learning curves, focusing on MR.Q, SimBaV2, TDMPC2, and FoG.These methods are selected for strong performance across numerous challenging benchmarks.
  • D. Full Main Results: Performance figures and tables report 95% bootstrapped confidence intervals, while aggregated mean, median, and IQM intervals use n × T benchmark-task samples with rliable.The intervals are computed across random seeds for individual tasks and across evaluated tasks for aggregates.

D.1. Gym MuJoCo Results

On Gym MuJoCo tasks, DR.Q outperforms SimBaV2 in 3 of 5 environments, although its average normalized score is slightly lower. The weaker aggregate result is mainly attributed to poor Hopper-v4 performance, where MR.Q and FoG also perform poorly.

  • Results: DR.Q outperforms SimBaV2 in 3 of 5 Gym MuJoCo environments, but slightly trails it in average normalized score.The gap is mainly due to DR.Q’s poor performance on Hopper-v4.
  • Results: MR.Q and FoG also fail to achieve strong performance on Hopper-v4.This shared weakness helps explain the poor Hopper-v4 result affecting DR.Q’s average normalized score.
  • Evaluation protocol: Table 6 reports final average performance at 1M environment steps for each Gym MuJoCo task.Scores include 95% bootstrap confidence intervals, with aggregate mean, median, and interquartile mean computed over TD3-normalized scores.
  • Learning curves: Figure 5 reports average episode returns across 10 random seeds with shaded 95% bootstrap confidence intervals.The curves cover the Gym MuJoCo tasks evaluated in the comparison.

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

The appendix reports full learning curves and final performance comparisons for DMC-Easy, DMC-Hard, and HumanoidBench with and without dexterous hands. Results use average returns, bootstrap confidence intervals, and aggregate statistics under specified training horizons and action repeats.

  • D.2. DMC Suite Easy Results: DMC-Easy learning curves show average returns for each environment with shaded 95% bootstrap confidence intervals.
  • D.2. DMC Suite Easy Results: DMC-Easy results report final average performance at 500K steps, corresponding to 1M environment steps with action repeat 2.Table 7 includes 95% bootstrap confidence intervals and aggregate mean, median, and interquartile mean (IQM) in units of 1k.
  • D.3. DMC Suite Hard Results: DMC-Hard results report final average returns at 500K steps, equivalent to 1M environment steps under action repeat 2.Table 8 reports bracketed 95% bootstrap confidence intervals and aggregate mean, median, and IQM in units of 1k.
  • D.3. DMC Suite Hard Results: DMC-Hard learning curves present average returns for each environment with shaded 95% bootstrap confidence intervals.
  • D.4. HumanoidBench (w/o Hand) Results: HumanoidBench without dexterous hands reports final average returns at 1M environment steps, equivalent to 500K steps with action repeat 2.Aggregate mean, median, and IQM are computed over the success normalized score, with bracketed 95% bootstrap confidence intervals.
  • D.4. HumanoidBench (w/o Hand) Results: HumanoidBench without hands provides learning curves of average returns for each task with light-colored 95% bootstrap confidence intervals.
  • 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 the success normalized score, and bracketed values denote 95% bootstrap confidence intervals.
  • D.5. HumanoidBench (w/ Hand) Results: HumanoidBench with hands presents task-level learning curves using average returns and light-colored 95% bootstrap confidence intervals.

D.6. Visual DMC Suite Results

DR.Q is evaluated on 12 selected visual-input DMC tasks across 10 seeds, excluding simple tasks whose performance had already saturated. It generally surpasses MR.Q and TDMPC2 in sample efficiency and final performance, particularly on visual-dog-stand and visual tasks.

  • Experimental setup: The study selects 12 visual-input DMC tasks, omitting simple tasks where MR.Q or baseline performance had already saturated.The overall results are reported across 10 seeds.
  • Results: DR.Q generally outperforms strong baselines MR.Q and TDMPC2 in sample efficiency and final performance.The advantage is especially pronounced on tasks such as visual-dog-stand.
  • Results: Final performance is reported as average return at 1M environment steps, equivalent to 500K steps with action repeat 2.Bracketed values denote 95% bootstrap confidence intervals, while aggregate statistics use success-normalized scores.

E. Additional Experiments

Additional experiments omitted from the main paper use the Section C setup, evaluating all variants and baselines across 10 seeds and 1M environment steps. Extended ablations examine InfoNCE loss and sampling strategies, reporting average returns with 95% confidence intervals.

  • All additional experiments follow Section C’s setup and run every variant or baseline for 10 seeds over 1M environment steps.
  • An extended ablation study evaluates the InfoNCE loss, averaging results across 10 seeds.
  • An extended ablation study compares sampling strategies using average return results across 10 seeds.
  • The plotted results use shaded regions to represent 95% confidence intervals.

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

The extended analyses show that DR.Q’s InfoNCE, LAP, forget mechanism, and mutual-information design are important for performance, while visualizations examine its representations and sampling behavior. Additional experiments compare modified hyperparameters and characterize faded prioritized replay.

  • E.1. Extended Ablation Study: DR.Q generally outperforms DR.Q (w/o InfoNCE) in sample efficiency and final performance, with exclusion especially harmful on h1-sit-hard-v0 and h1hand-pole-v0.These ablations demonstrate the necessity and importance of the InfoNCE loss term.
  • E.1. Extended Ablation Study: Removing either LAP or the forget mechanism can cause unsatisfactory performance, especially on dexterous-hand HumanoidBench tasks h1hand-walk-v0 and h1hand-sit-hard-v0.On humanoid-run and humanoid-walk, removing LAP causes a severe performance drop, highlighting LAP’s importance.
  • E.1. Extended Ablation Study: DR.Q removes its latent dynamics loss term to form DR.Q (w/o dyn loss), while MR.Q mainly uses latent dynamics loss and DR.Q additionally uses InfoNCE.The comparison is designed to examine the latent dynamics loss term’s role in DR.Q.
  • E.2. Comparison of DR.Q against MR.Q with Modified Hyperparameters: DR.Q uses modified hyperparameters relative to MR.Q, including a larger encoder learning rate and encoder hidden dimension, and evaluates their effect empirically.The passage identifies these settings as design choices motivated by scaling the network.
  • E.3. On the Mutual Information Loss: The representation analysis visualizes state-action representations zsa with t-SNE after 200K training steps on four tasks using 5000 replay-buffer samples.The tasks are HalfCheetah-v4, humanoid-walk, h1-sit-hard-v0, and h1hand-stand-v0.
  • E.4. Visualization Results of Sampling Strategies: Sampling visualizations compare MR.Q’s LAP-based probabilities, proportional to TD error, with DR.Q’s faded-PER probabilities, proportional to TD error and forget weight.Both agents run for 200K steps on eight continuous-control tasks, using TD errors from the most recent 100K samples.
Loading 2605.11711v1…