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/
Lay Summary
Our goal is to solve offline reinforcement learning, where an agent must learn good decisions only from previously collected data. This setting is important when collecting new interaction data is expensive, unsafe, or impractical, such as in robotics. In this work, we investigate whether using a future predictive model can help offline reinforcement learning. A natural approach is to predict future states step by step, but this can cause small prediction errors to accumulate over long horizons. To address this issue, we propose Universal Horizon Models (UHM), which directly predict the future state after a specified number of steps in a single prediction. Building on UHM, we further introduce an offline reinforcement learning method that uses directly predicted future. We evaluate the proposed method on a wide range of benchmark tasks, including navigation and manipulation tasks. The results show that our method achieves strong performance across diverse environments, especially in tasks that require long-horizon reasoning or involve noisy and suboptimal datasets.