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

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. * CQL (Kumar et al., 2020) adds penalization to Q-values for the samples out of distribution; * TD3+BC (Fujimoto & Gu, 2021) incorporates a BC regularization term into the policy optimization objective; * EDAC (An et al., 2021) proposed to penalize based on the uncertainty degree of the Q-value. Model-based offline RL. * COMBO(Yu et al., 2021) which applies CQL in dyna-style enforces Q-values small on OOD samples; * RAMBO(Rigter et al., 2022) trains the dynamics model adversarially to minimize the value function without loss of accuracy on the transition prediction; * MOPO (Yu et al., 2020) learns a pessimistic value function from rewards penalized with the uncertainty of the dynamics model's prediction; * MOBILE (Sun et al., 2023) penalizes the rewards with uncertainty quantified by the inconsistency of Bellman estimations under an ensemble of learned dynamics models.