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

arXiv: https://arxiv.org/abs/2402.03146 · HTML: https://mlsanity.com/paper/2402.03146/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in RL theory and sample complexity, one of 15 directions in Deep reinforcement learning methods.
- Builds on 20 of its 50 references that are in the map.
- First posted on arXiv 5 Feb 2024.

## Placed in

- Deep reinforcement learning methods / [RL theory and sample complexity](https://mlsanity.com/deeprl/rl-theory-and-sample-complexity/)

## Builds on (most cited first, 20 of 20)

- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm](https://mlsanity.com/paper/1712.01815/) (2017)
- [Neural Network Dynamics for Model-Based Deep Reinforcement Learning with Model-Free Fine-Tuning](https://mlsanity.com/paper/1708.02596/) (2017)
- [Robust Deep Reinforcement Learning against Adversarial Perturbations on State Observations](https://mlsanity.com/paper/2003.08938/) (2020)
- [Professor Forcing: A New Algorithm for Training Recurrent Networks](https://mlsanity.com/paper/1610.09038/) (2016)
- [Robust Deep Reinforcement Learning with Adversarial Attacks](https://mlsanity.com/paper/1712.03632/) (2017)
- [How (not) to Train your Generative Model: Scheduled Sampling, Likelihood, Adversary?](https://mlsanity.com/paper/1511.05101/) (2015)
- [Generalized Teacher Forcing for Learning Chaotic Dynamics](https://mlsanity.com/paper/2306.04406/) (2023)
- [Investigating Compounding Prediction Errors in Learned Dynamics Models](https://mlsanity.com/paper/2203.09637/) (2022)
- [Self-Correcting Models for Model-Based Reinforcement Learning](https://mlsanity.com/paper/1612.06018/) (2016)
- [Evaluation of deep learning models for multi-step ahead time series prediction](https://mlsanity.com/paper/2103.14250/) (2021)
- [Learning Dynamics Models for Model Predictive Agents](https://mlsanity.com/paper/2109.14311/) (2021)
- [Evaluating model-based planning and planner amortization for continuous control](https://mlsanity.com/paper/2110.03363/) (2021)
- [Exploring the Training Robustness of Distributional Reinforcement Learning against Noisy State Observations](https://mlsanity.com/paper/2109.08776/) (2021)
- [Model-based micro-data reinforcement learning: what are the crucial model properties and which model to choose?](https://mlsanity.com/paper/2107.11587/) (2021)

## Built on by (most cited first, 1 of 1)

- [From Data to Rewards: a Bilevel Optimization Perspective on Maximum Likelihood Estimation](https://mlsanity.com/paper/2510.07624/) (2025)

## Abstract (arXiv)

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.