Double Horizon Model-Based Policy Optimization

AI summary

Overview

Research area: Model-based reinforcement learning (MBRL), specifically the design of synthetic model rollouts for policy optimization on continuous-control benchmarks.

Technical level: Intermediate to Advanced. The paper builds directly on model-based actor-critic methods, model-based value expansion (MVE), the reparameterization trick, and entropy-regularized MDPs, and it includes full objective functions and an algorithm listing.

Scope: The paper proposes Double Horizon Model-Based Policy Optimization (DHMBPO), which splits model rollouts into a long "distribution rollout" (DR) and a short "training rollout" (TR), and evaluates it on Gymnasium (GYM) and DMControl (DMC) continuous-control tasks.

What This Paper Is About

Model-based RL reduces costly real-environment interaction by generating synthetic trajectories from a learned dynamics model, but the length of those rollouts creates two competing problems: long model rollouts drift into model bias and distribution shift, while long differentiable rollouts used for policy gradients inflate gradient variance. The paper argues that the horizon best suited to fixing the first problem is not the same as the horizon best suited to fixing the second, and proposes using two different horizons at once — a long one to produce on-policy states and a short one to compute accurate, stable value gradients.

Key Contributions

1. Diagnosis of a horizon conflict. The paper formalizes why the optimal rollout length for maintaining an on-policy state distribution (to reduce distribution shift) differs from the optimal length for stable value-gradient estimation (to limit gradient variance), and states that these two optima need not coincide.

2. The DHMBPO algorithm. A method that first performs a long DR (horizon 20 steps in their setup) from replay-buffer states to populate a model buffer of on-policy-like samples, then runs a short TR (horizon 5 steps) from those states to build an MVE-based value gradient for the actor update.

3. Empirical gains in sample efficiency and runtime. DHMBPO outperformed the compared MBRL methods on GYM tasks without per-task tuning, and on the DMC suite it showed significant improvements over Dreamer v3 with less variation than TD-MPC2 according to the interquartile mean (IQM).

4. Ablation and sensitivity analysis. An ablation separating DR and TR (including the special cases corresponding to SAC-SVG(H) and MBPO), a sweep over DR horizons (0, 10, 20, 40) and TR horizons (1, 3, 5, 7, 9), and a critic-learning experiment using a normalized root median squared error (RMedSE) against Monte Carlo returns.

Main Findings

• DHMBPO was the most sample-efficient on the GYM tasks tested. On five GYM tasks it achieved the highest sample efficiency across all tasks even without per-task tuning, with results reported over eight random seeds (MBPO used five seeds).

• Large runtime advantage over MACURA. Runtimes until 500K environment steps on an 8 × NVIDIA RTX A4000 16GB GPU system were 3.6, 3.3, 3.7, 5.2 and 4.0 hours for DHMBPO on Ant, HalfCheetah, Hopper, Humanoid and Walker2d, versus 58.3, 32.6, 73.3, 104.3 and 63.0 hours for MACURA, and 5.2, 6.7, 6.7, 6.9 and 6.6 hours for SAC-SVG(H). The mean ratio to DHMBPO's runtime was 1.0 for DHMBPO, 16.8 for MACURA and 1.6 for SAC-SVG(H).

• Lower update-to-data (UTD) ratio explains much of the runtime gap. The paper reports that MACURA's UTD ratio on Humanoid was 20, while DHMBPO's (and SAC-SVG(H)'s) was 1. Increasing the UTD ratio yielded only marginal sample-efficiency improvements while significantly increasing execution time (Appendix D.5, Figure 20).

• Combining DR and TR beat either alone. In the ablation over 10 tasks (five GYM, five DMC) with 8 seeds each, DHMBPO (20, 5) achieved higher returns and better sample efficiency than the DR-only variant (20, 0, corresponding to MBPO) and the TR-only variant (0, 5, corresponding to SAC-SVG(H)). The TR-only variant suffered on challenging tasks such as Humanoid, Walker2d and quadruped-run due to reliance on off-policy samples, while the DR-only variant produced weaker value corrections.

• DR is beneficial but should not be too short; TR should be short. Varying the DR horizon (0, 10, 20, 40) with TR fixed at 5 showed DR benefits from being relatively long, with 20 outperforming 0. Varying the TR horizon (1, 3, 5, 7, 9) with DR fixed at 20 showed TR = 5 balances sample efficiency and gradient stability; very short TR reduces sample efficiency while excessively long TR can cause gradient norms to explode and destabilize learning (Appendix C.3, Figure 10).

• DR improves critic accuracy. In the critic-learning experiment (100K environment steps on five GYM environments, then freezing models and actor and re-initializing critics), ground-truth returns were computed as Monte Carlo returns using 2,048 real-environment rollouts over 256 sampled state-action pairs, and discrepancy was measured with RMedSE against the 5-step MVE estimate.

• Aggregated DMC comparison. Metrics were computed at 50% normalized steps relative to DHMBPO's mean scores on 18 DMC tasks; the IQM showed DHMBPO improving significantly over Dreamer v3 with less variation than TD-MPC2.

• Sample-efficiency cost of removing TR. For DHMBPO without TR (corresponding to MBPO) at different UTD ratios, raising the UTD ratio improved sample efficiency (Figure 11) but lengthened execution time (Figure 12) — i.e., matching DHMBPO's sample efficiency without TR requires extra computation.

Methodology in Plain English

The authors build a standard deep model-based actor-critic setup: a policy network (actor), a Q-network (critic) with a target critic updated by an exponential moving average, and a learned dynamics model and reward model. The dynamics model is a bootstrap ensemble of probabilistic neural networks whose outputs are independent multivariate Gaussian distributions, so individual models capture aleatoric noise while bootstrapping captures epistemic uncertainty.

Their key design decision is to use the model twice, with different horizon lengths:

1. Distribution rollout (DR), long. Starting from states stored in the replay buffer, they simulate 20 steps with the learned model and save the resulting transitions into a separate model buffer. This pushes the training states closer to the current policy's state distribution.

2. Training rollout (TR), short. Starting from states drawn from that model buffer, they simulate 5 steps with differentiable transitions and backpropagate through them under the reparameterization trick to form a model-based value expansion (MVE) estimate — the on-policy rewards of the short rollout plus a terminal critic prediction. The actor's gradient comes from this value gradient; the critic is trained to match the corresponding MVE targets.

Because the actor uses the accurate value gradient from the short TR, the algorithm needs fewer parameter updates per environment step (a UTD ratio of 1), which keeps runtime down. The critic and actor are updated alternately inside a policy optimization loop that runs for the episode length L times, and the model buffer is cleared each iteration. Model predictions, rewards and rollouts were all kept consistent across the ablated variants so that only the DR/TR combination changed.

Evaluation used rliable for aggregated metrics (sample efficiency curves and aggregation scores such as median, interquartile mean and mean, with 95% confidence intervals), plus a task-specific normalization that divides environment steps by each task's maximum step budget and test returns by a designated baseline's final performance in order to combine GYM and DMC results. GYM comparisons are shown as raw, unnormalized returns; DHMBPO's test return is the sample mean of cumulative rewards over 10 episodes.

Why This Matters

Impact on research. The paper reframes rollout length not as one hyperparameter but as two, tied to two distinct failure modes (distribution shift/model bias versus gradient variance). That gives MBRL researchers a concrete decomposition to test, and it shows that a low UTD ratio — often treated as a sample-efficiency handicap — can be paired with model-based value gradients to keep sample efficiency while cutting wall-clock cost, which matters for benchmarking practice.

Real-world applications (potential, given MBRL's continuous-control focus):

• Robot locomotion and legged/mobile robot control, of the kind represented by the Ant, HalfCheetah, Hopper, Humanoid, Walker2d and DMC tasks.
• Simulation-trained controllers for physical systems where real-world interaction is slow, risky or expensive.
• Autonomous systems that must learn from a limited budget of real trials.
• Industrial process control and any setting where a simulator stands in for the plant.

Industry relevance. Runtimes in the tens-to-hundreds of hours per task are a practical barrier to iterating on learned controllers; DHMBPO's reported runtimes of roughly 3 to 5 hours per GYM task and its lower UTD ratio point at lower compute cost per experiment. The single shared hyperparameter set (with no per-task tuning) is also relevant to deployment pipelines where per-task tuning is not affordable.

Future Directions

• Principled horizon selection. The paper sets DR = 20 and TR = 5 empirically and proposes, in Appendix D.2, a practical method to determine an optimal TR horizon from a small number of real samples — an avenue that could be extended to adapt both horizons online.
• Scaling to harder, higher-dimensional tasks. DHMBPO was tested on GYM and DMC continuous-control suites; whether the double-horizon split holds for more complex domains (for example those where latent model-based methods such as TD-MPC2 and Dreamer v3 are preferred) is not established by the presented results.
• Reconciling with latent and recurrent model architectures. The paper notes that SAC-SVG(H) uses a GRU and DHMBPO uses deep ensembles; how the DR/TR split interacts with different model classes is left open.
• Understanding the theory. The authors state that MBPO's formalization and theoretical analysis do not necessarily support the benefits of model-generated on-policy data, and they interpret its success empirically; a formal account of why on-policy DR states help is still missing.

Target Audience

Reinforcement learning researchers and graduate students working on model-based RL, sample-efficient continuous control, or policy-gradient variance; practitioners who train controllers on MuJoCo-style benchmarks and care about wall-clock runtime as much as sample count; and engineers evaluating whether model-based methods are computationally viable for their robot-control or simulation-based control pipelines. Readers without background in actor-critic methods, MVE, or the reparameterization trick will need to consult the cited background works (Janner et al., MBPO; Feinberg et al., MVE; Parmas et al., gradient variance).

Authors’ abstract

Model-based reinforcement learning (MBRL) reduces the cost of real-environment sampling by generating synthetic trajectories (called rollouts) from a learned dynamics model. However, choosing the length of the rollouts poses two dilemmas: (1) Longer rollouts better preserve on-policy training but amplify model bias, indicating the need for an intermediate horizon to mitigate distribution shift (i.e., the gap between on-policy and past off-policy samples). (2) Moreover, a longer model rollout may reduce value estimation bias but raise the variance of policy gradients due to backpropagation through multiple steps, implying another intermediate horizon for stable gradient estimates. However, these two optimal horizons may differ. To resolve this conflict, we propose Double Horizon Model-Based Policy Optimization (DHMBPO), which divides the rollout procedure into a long "distribution rollout" (DR) and a short "training rollout" (TR). The DR generates on-policy state samples for mitigating distribution shift. In contrast, the short TR leverages differentiable transitions to offer accurate value gradient estimation with stable gradient updates, thereby requiring fewer updates and reducing overall runtime. We demonstrate that the double-horizon approach effectively balances distribution shift, model bias, and gradient instability, and surpasses existing MBRL methods on continuous-control benchmarks in terms of both sample efficiency and runtime.