A Multi-step Loss Function for Robust Learning of the Dynamics in Model-based Reinforcement Learning

TL;DR

This paper introduces a multi-step loss function for training one-step models in model-based reinforcement learning, improving robustness to noise and reducing compounding errors in trajectory predictions.

Contribution

It proposes a novel multi-step weighted loss for training dynamics models, demonstrating its effectiveness especially under noisy conditions and in real-world scenarios.

Findings

Multi-step loss improves future prediction accuracy.

Models trained with this loss perform better in noisy environments.

Multi-step models outperform one-step models in noisy, real-world settings.

Abstract

In model-based reinforcement learning, most algorithms rely on simulating trajectories from one-step models of the dynamics learned on data. A critical challenge of this approach is the compounding of one-step prediction errors as the length of the trajectory grows. In this paper we tackle this issue by using a multi-step objective to train one-step models. Our objective is a weighted sum of the mean squared error (MSE) loss at various future horizons. We find that this new loss is particularly useful when the data is noisy (additive Gaussian noise in the observations), which is often the case in real-life environments. To support the multi-step loss, first we study its properties in two tractable cases: i) uni-dimensional linear system, and ii) two-parameter non-linear system. Second, we show in a variety of tasks (environments or datasets) that the models learned with this loss achieve a significant improvement in terms of the averaged R2-score on future prediction horizons. Finally, in the pure batch reinforcement learning setting, we demonstrate that one-step models serve as strong baselines when dynamics are deterministic, while multi-step models would be more advantageous in the presence of noise, highlighting the potential of our approach in real-world applications.