Scalable Offline Model-Based RL with Action Chunks

TL;DR

This paper introduces MAC, a scalable offline model-based RL method using action chunks to reduce model errors and rejection sampling to prevent exploitation, achieving state-of-the-art results on complex long-horizon tasks.

Contribution

The paper proposes action-chunk models and rejection sampling techniques to improve the scalability and performance of offline model-based RL in complex tasks.

Findings

MAC outperforms existing offline model-based RL algorithms on long-horizon tasks.

Action chunks reduce compounding errors in long-term predictions.

Rejection sampling prevents model exploitation from out-of-distribution actions.

Abstract

In this paper, we study whether model-based reinforcement learning (RL), in particular model-based value expansion, can provide a scalable recipe for tackling complex, long-horizon tasks in offline RL. Model-based value expansion fits an on-policy value function using length-n imaginary rollouts generated by the current policy and a learned dynamics model. While larger n reduces bias in value bootstrapping, it amplifies accumulated model errors over long horizons, degrading future predictions. We address this trade-off with an \emph{action-chunk} model that predicts a future state from a sequence of actions (an "action chunk") instead of a single action, which reduces compounding errors. In addition, instead of directly training a policy to maximize rewards, we employ rejection sampling from an expressive behavioral action-chunk policy, which prevents model exploitation from out-of-distribution actions. We call this recipe \textbf{Model-Based RL with Action Chunks (MAC)}. Through experiments on highly challenging tasks with large-scale datasets of up to 100M transitions, we show that MAC achieves the best performance among offline model-based RL algorithms, especially on challenging long-horizon tasks.