Paper: Multi-timestep models for Model-based Reinforcement Learning (ICLR-2024)

Reviewer comment (Evaluation):
Sec 4.3 + Fig 2: I am confused. Here $h=2$ ($\alpha=0$) improves performance whereas in Fig 1 it did not?!

Author response:
Sec 4.3 + Fig 2: Here $\alpha = 0$ is not equivalent to $h=2$ in Fig 1. $\alpha = 0$ corresponds to optimizing a one-step model $\hat{s}_{t+1} \sim \hat{p}(s_t, a_t)$ with the loss two-steps ahead:  

$L\Big(s_{t+2},\hat{p}\big(\hat{p}(s_t, a_t), a_{t+1}\big) \Big)$, While $h=2$ is a direct two-steps model $\hat{s}_{t+2}$  

$\sim \hat{p}(s_t,a_t,a_{t+1})$.

Reviewer comment (Presentation):
Fig 4: for h=4 and h=10, the pink lines seems to be missing

Author response:
Fig 4: The pink lines (decay with $\beta=0.9$) for $h=4$ and $h=10$ are missing because the corresponding R2 score is outside the interval $(0,1)$.

Reviewer comment (Presentation):
Fig 5: Do I understand correctly that e.g. the red markers are all the same model trained with n=10 loss but used only for planning with h=1, 5, 10...?

Author response:
Sec 5.3: Indeed, in Fig 5, dots of the same color represent the same model (and therefor have the same return value). We show different manifestations of the same dot along the x-axis (R2 score) to showcase the relationship between the return and $R2(h)$ for different values of $h$.

Reviewer comment (Experiments):
Empirical comparison with similar approaches that are mentioned in the related work is missing. For instance, it is unclear why one should use a weighted loss instead of learning a different model for each time horizon.

Author response:
We thank Reviewer huS2 for their insightful comments. The empirical comparison with similar methods from the literature (particularly the Data Augmentation idea used in several papers [1,2,3]) will be included in a future submission.

[1] Abbeel et. al. (2005) Learning vehicular dynamics, with application to modeling helicopters 

[2] Venkatraman et. al. (2015) Improving multi-step prediction of learned time series models 

[3] Erik Talvitie. (2017) Self-Correcting Models for Model-Based Reinforcement Learning

Reviewer comment (Evaluation):
How do the computational costs for the different horizon training methods compare?

Author response:
Reviewer yoeH’s second point concerns the computational cost of the multi-step loss. This is a primary follow-up idea that we are considering, and we have presented the initial stages of it in Appendix C. Indeed, the analytical form of the multi-step loss gradient simplifies nicely, and we plan to explore approximation schemes that will reduce the current linear temporal complexity.

Reviewer comment (Novelty):
The proposed multistep loss has been proposed and utilized very often in previous works (e.g., references [1–4], with potentially more examples in older system identification literature for linear system identification). For example, [4] provides an ablation study across multiple horizons and systems, showing that going longer in the horizon is not necessarily better in terms of reward and often even worse. Although evaluation metrics and RL algorithms differ (with [4] using a CEM-like planner instead of an actor-critic), this discrepancy requires further investigation.

Author response:
We thank Reviewer XNaE for their insightful comments. We now delve into a discussion on the references cited by Reviewer XNaE and how they intersect with our contribution:

[1] Hafner et. al. (2015). Learning latent dynamics for planning from pixels: This paper presents a derivation of a multi-step $\beta$-VAE bound on latent sequence models trained via variational inference (termed latent overshooting). While the underlying intuition is similar, the context in which we apply the multi-step loss differs, focusing on simple feedforward dynamics models trained via MSE minimization. Moreover, the authors discovered that the latent overshooting variant of their algorithm didn’t yield any improvements and was not adopted in their final agent. This could potentially be due to the fixed value of all $\beta_{>1}$​, which overlooks the exponential growth of the error, a problem we aim to address in our study.

[2] Abbeel et. al. (2005) Learning vehicular dynamics, with application to modeling helicopters: This paper indeed proposes a similar loss function termed “lagged criterion” (equation 4). However, to simplify the non-convex optimization to a linear least squares problem, the authors relax the dependency of intermediate predictions $\hat{s}_{t+h}$ on the model parameters $(A,B)$. This results in a data augmentation-like technique that instructs the model to recover from its own errors without propagating the gradients through the multi-horizon applications of the model.

[3] Venkatraman et. al. (2015) Improving multi-step prediction of learned time series models: Although this paper begins with a theorem (Theorem 1) that illustrates the exponential growth of the multi-step loss, the algorithm idea “Data-As-Demonstrator” is essentially a data-augmentation procedure that “synthetically generates correction examples for the learner to use”. This is akin to [2] and does not involve any gradient-based optimization of the multi-step loss as we propose.

[4] Lutter et. al. (2021). Learning Dynamics Models for Model Predictive Agents: This paper is the closest to our concept as it back-propagates the gradients through the multi-step loss, and we thank Reviewer XNaE for highlighting this. However, as Reviewer XNaE pointed out, the differences in the model architecture, the evaluation metric, and the actor could potentially be confounding factors in our respective analyses. Furthermore, the question of the optimal horizon is strongly tied to the weighting of the intermediate MSE terms as this corrects the differences in scale. This question is not considered in [4], and we aim to demonstrate that it’s a crucial component in maximizing the benefits of the multi-step loss component.

Reviewer comment (Novelty):
There is a lack of clarity on how the proposed multi-step loss differs from previous works, apart from the different weighting approach.

Author response:
We thank Reviewer XNaE for their insightful comments. We now delve into a discussion on the references cited by Reviewer XNaE and how they intersect with our contribution:

[1] Hafner et. al. (2015). Learning latent dynamics for planning from pixels: This paper presents a derivation of a multi-step $\beta$-VAE bound on latent sequence models trained via variational inference (termed latent overshooting). While the underlying intuition is similar, the context in which we apply the multi-step loss differs, focusing on simple feedforward dynamics models trained via MSE minimization. Moreover, the authors discovered that the latent overshooting variant of their algorithm didn’t yield any improvements and was not adopted in their final agent. This could potentially be due to the fixed value of all $\beta_{>1}$​, which overlooks the exponential growth of the error, a problem we aim to address in our study.

[2] Abbeel et. al. (2005) Learning vehicular dynamics, with application to modeling helicopters: This paper indeed proposes a similar loss function termed “lagged criterion” (equation 4). However, to simplify the non-convex optimization to a linear least squares problem, the authors relax the dependency of intermediate predictions $\hat{s}_{t+h}$ on the model parameters $(A,B)$. This results in a data augmentation-like technique that instructs the model to recover from its own errors without propagating the gradients through the multi-horizon applications of the model.

[3] Venkatraman et. al. (2015) Improving multi-step prediction of learned time series models: Although this paper begins with a theorem (Theorem 1) that illustrates the exponential growth of the multi-step loss, the algorithm idea “Data-As-Demonstrator” is essentially a data-augmentation procedure that “synthetically generates correction examples for the learner to use”. This is akin to [2] and does not involve any gradient-based optimization of the multi-step loss as we propose.

[4] Lutter et. al. (2021). Learning Dynamics Models for Model Predictive Agents: This paper is the closest to our concept as it back-propagates the gradients through the multi-step loss, and we thank Reviewer XNaE for highlighting this. However, as Reviewer XNaE pointed out, the differences in the model architecture, the evaluation metric, and the actor could potentially be confounding factors in our respective analyses. Furthermore, the question of the optimal horizon is strongly tied to the weighting of the intermediate MSE terms as this corrects the differences in scale. This question is not considered in [4], and we aim to demonstrate that it’s a crucial component in maximizing the benefits of the multi-step loss component.

Reviewer comment (Experiments):
Evidence from [4] shows that results and conclusions from cartpole experiments do not transfer to more complex systems, where longer horizons may perform worse.

Author response:
We thank Reviewer XNaE for their insightful comments. We now delve into a discussion on the references cited by Reviewer XNaE and how they intersect with our contribution:

[1] Hafner et. al. (2015). Learning latent dynamics for planning from pixels: This paper presents a derivation of a multi-step $\beta$-VAE bound on latent sequence models trained via variational inference (termed latent overshooting). While the underlying intuition is similar, the context in which we apply the multi-step loss differs, focusing on simple feedforward dynamics models trained via MSE minimization. Moreover, the authors discovered that the latent overshooting variant of their algorithm didn’t yield any improvements and was not adopted in their final agent. This could potentially be due to the fixed value of all $\beta_{>1}$​, which overlooks the exponential growth of the error, a problem we aim to address in our study.

[2] Abbeel et. al. (2005) Learning vehicular dynamics, with application to modeling helicopters: This paper indeed proposes a similar loss function termed “lagged criterion” (equation 4). However, to simplify the non-convex optimization to a linear least squares problem, the authors relax the dependency of intermediate predictions $\hat{s}_{t+h}$ on the model parameters $(A,B)$. This results in a data augmentation-like technique that instructs the model to recover from its own errors without propagating the gradients through the multi-horizon applications of the model.

[3] Venkatraman et. al. (2015) Improving multi-step prediction of learned time series models: Although this paper begins with a theorem (Theorem 1) that illustrates the exponential growth of the multi-step loss, the algorithm idea “Data-As-Demonstrator” is essentially a data-augmentation procedure that “synthetically generates correction examples for the learner to use”. This is akin to [2] and does not involve any gradient-based optimization of the multi-step loss as we propose.

[4] Lutter et. al. (2021). Learning Dynamics Models for Model Predictive Agents: This paper is the closest to our concept as it back-propagates the gradients through the multi-step loss, and we thank Reviewer XNaE for highlighting this. However, as Reviewer XNaE pointed out, the differences in the model architecture, the evaluation metric, and the actor could potentially be confounding factors in our respective analyses. Furthermore, the question of the optimal horizon is strongly tied to the weighting of the intermediate MSE terms as this corrects the differences in scale. This question is not considered in [4], and we aim to demonstrate that it’s a crucial component in maximizing the benefits of the multi-step loss component.