Reward-Consistent Dynamics Models are Strongly Generalizable for Offline Reinforcement Learning

TL;DR

This paper introduces reward-consistent dynamics models that improve generalization in offline reinforcement learning by leveraging a dynamics reward function, leading to significant performance gains across multiple benchmarks.

Contribution

The paper proposes reward-consistent dynamics models and the MOREC method, which enhances offline RL by filtering transitions based on a learned dynamics reward, improving generalization.

Findings

MOREC outperforms previous methods on D4RL and NeoRL benchmarks.

MOREC achieves above 95% online RL performance in multiple tasks.

The method demonstrates strong generalization ability, recovering unseen transitions.

Abstract

Learning a precise dynamics model can be crucial for offline reinforcement learning, which, unfortunately, has been found to be quite challenging. Dynamics models that are learned by fitting historical transitions often struggle to generalize to unseen transitions. In this study, we identify a hidden but pivotal factor termed dynamics reward that remains consistent across transitions, offering a pathway to better generalization. Therefore, we propose the idea of reward-consistent dynamics models: any trajectory generated by the dynamics model should maximize the dynamics reward derived from the data. We implement this idea as the MOREC (Model-based Offline reinforcement learning with Reward Consistency) method, which can be seamlessly integrated into previous offline model-based reinforcement learning (MBRL) methods. MOREC learns a generalizable dynamics reward function from offline data, which is subsequently employed as a transition filter in any offline MBRL method: when generating transitions, the dynamics model generates a batch of transitions and selects the one with the highest dynamics reward value. On a synthetic task, we visualize that MOREC has a strong generalization ability and can surprisingly recover some distant unseen transitions. On 21 offline tasks in D4RL and NeoRL benchmarks, MOREC improves the previous state-of-the-art performance by a significant margin, i.e., 4.6% on D4RL tasks and 25.9% on NeoRL tasks. Notably, MOREC is the first method that can achieve above 95% online RL performance in 6 out of 12 D4RL tasks and 3 out of 9 NeoRL tasks.