Multi-timestep models for Model-based Reinforcement Learning

We thank Reviewer pXQz for their insightful comments. Here we address some of their questions, and leave the open ones for a future submission.

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})$.

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)$.

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$.

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.

We thank Reviewer yoeH for their insightful comments. The correlation between action sequences and the underlying state transitions, as highlighted by Reviewer yoeH, is indeed a significant issue. Our current evaluation method is an approximation that downplays the role of the controller in the dynamics function. We argue that this remains a fair comparison method as it is consistently applied across models. However, we agree that a thorough examination of the interaction between the controller and the visited states necessitates a reconsideration of the current evaluation methodology, beyond just the final return of the agent.

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.

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