Offline Reinforcement Learning with Universal Horizon Models

Abstract: Model-based reinforcement learning (RL) offers a compelling approach to offline RL by enabling value learning on imagined on-policy trajectories. However, it often suffers from compounding errors due to repeated model inference on self-generated states. While geometric horizon models (GHM) alleviate this issue through direct prediction over a discounted infinite-horizon future, they remain challenged in accurately modeling distant future states. To this end, we introduce universal horizon models (UHM), a generalization of GHM that directly predicts future states under arbitrary horizons. Leveraging this flexibility, we propose a scalable value learning method that employs a winsorized horizon distribution to stabilize training by capping excessively large horizons. Experimental results on 100 challenging OGBench tasks demonstrate that the proposed method outperforms competitive baselines, particularly on tasks with highly suboptimal datasets and those requiring long-horizon reasoning. Project page: https://rllab-snu.github.io/projects/UHM/

• The paper introduces Universal Horizon Models (UHMs) that extend geometric horizon models to allow arbitrary multi-step future state predictions.
• The methodology leverages winsorized horizon distributions and a unified value backup operator to stabilize off-policy learning in offline RL.
• Empirical results on OGBench tasks show significant improvements over baselines, with up to 69% gains and reduced computational cost.

Overview

"Offline Reinforcement Learning with Universal Horizon Models" (2605.15603) addresses core challenges in scalable offline RL, particularly the stability and accuracy of model-based value estimation under distributional shifts and long-horizon reasoning. The paper introduces Universal Horizon Models (UHM), a generalization of geometric horizon models (GHM), capable of direct state prediction over arbitrary future horizons and proposes an associated value learning methodology leveraging winsorized horizon distributions to further stabilize training. Extensive empirical benchmarks on OGBench demonstrate numerical superiority over established baselines, especially in suboptimal and long-horizon regimes.

Problem Formulation and Limitations of Existing Approaches

In the offline RL setting, agents learn from fixed static datasets with no further access to the environment. Model-free actor-critic algorithms based on TD learning often suffer from bias propagation due to inaccuracies in bootstrapped value targets. Multi-step returns (as in n-step TD or TD($\lambda$)) partially mitigate bias, but these approaches must revert to limited dataset trajectories, resulting in significant distributional mismatch and reduced effective horizon, especially critical for long-horizon reasoning tasks.

Model-based RL theoretically mitigates this by synthesizing trajectories via learned dynamics, facilitating on-policy or near on-policy value expansion. However, the iterative composition of model predictions leads to compounding model errors, making multistep synthetic rollouts unreliable, especially under covariate shift. GHMs alleviate this by modeling the normalized discounted distribution over future states directly but are intrinsically constrained to geometric horizon sampling, which complicates accurate modeling of long-horizon dynamics.

Universal Horizon Models: Formulation and Learning

UHM generalizes prior future-state generative models by directly modeling $$m^\pi(x|s,a,n) = \Pr(S_n = x | s_0 = s, a_0 = a, \pi)$$ for any discrete $n \geq 1$. This framework subsumes existing single-step models (when $n=1$) and GHMs (when $n\sim\text{Geom}(1-\gamma)$) as special cases. Crucially, UHM allows arbitrary sampling of horizon $n$, making it possible to flexibly control the distribution of future predictions during training and value estimation.

The off-policy learning protocol replaces recursive rollout-based targets with direct bootstrapping:

• For $n=1$, model one-step transitions from the empirical dataset.
• For $n>1$, recursively sample next states/actions, ultimately matching UHM's $n$-step prediction to synthetic samples generated with the policy and the learned model for horizon $n-1$.

This bootstrapping scheme, supported by theoretical contraction arguments, provides a stable target for model learning, even under varied horizon distributions.

Value Learning with Winsorized Horizon Distributions

Building on UHM's ability to sample $$m^\pi(x|s,a,n) = \Pr(S_n = x | s_0 = s, a_0 = a, \pi)$$0-step future states, the authors construct a general value backup operator parameterized by a sub-probability measure over horizons. This enables unified modeling of standard $$m^\pi(x|s,a,n) = \Pr(S_n = x | s_0 = s, a_0 = a, \pi)$$1-step TD, TD($$m^\pi(x|s,a,n) = \Pr(S_n = x | s_0 = s, a_0 = a, \pi)$$2), and their arbitrary weighted analogues. To address instability induced by rare and excessively long horizons, the paper introduces the winsorized geometric measure, capping the effective horizon at $$m^\pi(x|s,a,n) = \Pr(S_n = x | s_0 = s, a_0 = a, \pi)$$3. This strategy prevents the statistical noise and poor model coverage associated with rare, long-range predictions.

Algorithmically, value targets for the critic are constructed via importance-weighted returns from UHM samples, and learning is further stabilized with EMA targets and actor-critic updates that include behavior policy mixing, reward modeling, and explicit handling of terminal states.

Empirical Evaluation

Experiments on a suite of 100 OGBench tasks benchmark the algorithm in three distinct regimes: standard, noisy (suboptimal), and long-horizon reasoning. Key empirical findings include:

• On standard tasks, UHM surpasses all model-based baselines, achieving a 14% higher average success rate than the strongest alternate methods, especially outperforming uncertainty-penalized approaches (e.g., MOPO, MOBILE) that are ineffective under sparse reward or high-dimensional action settings.
• For noisy datasets where trajectory distributions are highly suboptimal, UHM exhibits robust performance, evidencing the utility of on-policy imagination over dataset-reliant value estimation.
• On tasks requiring extended, compositional reasoning over long horizons, UHM demonstrates marked numerical superiority, with up to 38% improvement over GHM and 69% over dataset-based TD($$m^\pi(x|s,a,n) = \Pr(S_n = x | s_0 = s, a_0 = a, \pi)$$4) on average, firmly establishing the benefit of flexible, horizon-capped synthetic prediction.

Ablation studies verify the necessity of key components such as $$m^\pi(x|s,a,n) = \Pr(S_n = x | s_0 = s, a_0 = a, \pi)$$5 scheduling, behavior mixing, explicit terminal state handling, and horizon winsorization in achieving stable and performant learning. Update time analyses further demonstrate that UHM attains an order of magnitude reduction in computational cost compared to iterative rollout models, with update times near those of model-free approaches.

Theoretical and Practical Implications

The introduction of UHM enables arbitrarily flexible, single-shot prediction for $$m^\pi(x|s,a,n) = \Pr(S_n = x | s_0 = s, a_0 = a, \pi)$$6-step futures under any given policy, decoupling model-based value expansion from the limitations of recursive state propagation and fixed geometric horizons. This flexibility has profound implications for both theoretical RL—through its generalization of Bellman operators and value learning schemes—and for practical large-scale offline RL, where stability and computability of model-based approaches have been significant barriers.

The results highlight the importance of controlling horizon distributions in synthetic imagination, emphasize robust handling of terminal states, and suggest that future model-based methods should further integrate advanced generative modeling for richer state distributions, especially in high-dimensional or visually rich domains.

Future Directions

Outstanding challenges include the integration of UHM with high-capacity generative architectures for visual observations, the development of domain-adaptive mechanisms for restricting prediction to in-distribution states under covariate shift, and the combination with action-chunking or hierarchical abstractions to further compress effective decision horizons. Given the demonstrated computational efficiency, UHM-based approaches are promising candidates for scaling offline RL to increasingly complex and long-horizon practical domains.

Conclusion

This work establishes Universal Horizon Models as a theoretically grounded and empirically validated paradigm for scalable, stable, and flexible model-based offline RL. By generalizing future state modeling and introducing robust horizon control, the proposed approach achieves substantial improvements in both performance and computational tractability across a wide spectrum of challenging offline RL benchmarks. The framework provides a compelling foundation for subsequent research in general-purpose offline RL and long-horizon autonomous decision making.