Efficient Dynamics Modeling in Interactive Environments with Koopman Theory

We thank the reviewer for their valuable feedback and for identifying the key strengths of our approach. Below, we address the reviewer’s concerns:

Figure 3 seems to undermine the utility of the presented Koopman approach, since the GRU is able to compete with it on all of the examples and metrics. Maybe some explanation in the figure caption explaining why this is expected would be appreciated. GRU-based dynamics modeling is indeed competitive with our approach and is widely used in model-based RL algorithms. However, as shown in Figure 4/Table 3, GRUs are significantly slower during training compared to our method. We aimed to build a simple latent linear model which can be competitive with existing ones while being several times faster and stable while training with a longer horizon. Following the reviewer’s suggestion, we have added more explanation to our captions in the revised manuscript.

Longer/better figure captions would be nice.

We have added further explanation to the captions of Figures 3, 4, 5 and 6. We trust that these explanations will lead to a better understanding of the Figures presented.

This seems like a marginal improvement over existing works cited. Some more comparisons between this approach and existing algorithms could better showcase the additions.

The main focus of our paper is on a latent linear dynamics model using Koopman theory so that we can significantly speed up the training. As shown in Figure 3/Table 1, 2 and Figure 4/Table 3, our model is significantly faster than existing baselines (which include MLP, Transformer, and GRU) with a similar number of parameters while being better/competitive in terms of performance. In the revised version, we additionally provide the results with the Diagonal State Space Model (DSSM)[1,2,3] for offline RL. The addition of DSSM as a baseline further strengthens our results since DSSMs are closely related to our model, as discussed in Appendix D. We demonstrate that while a DSSM-based dynamics model offers faster performance compared to other baseline approaches, it still lags behind our method in terms of speed. Moreover, it also falls short in state prediction accuracy when contrasted with our approach in the offline RL datasets. If the reviewer is suggesting the inclusion of additional baselines in our model-free RL and model-based planning experiments, we are certainly open to exploring additional relevant algorithms. We would appreciate any specific baseline recommendations that the reviewer had in mind. Please note that experiments provided in Sections 4.2.1 and 4.2.2 are to demonstrate that our method can be easily incorporated into existing RL and planning algorithms that use a dynamics model to enhance their performance

$$[1] Gu, Albert, et al. "On the parameterization and initialization of diagonal state space models." Advances in Neural Information Processing Systems 35 (2022). [2] Gupta, Ankit, et al. "Diagonal state spaces are as effective as structured state spaces." Advances in Neural Information Processing Systems 35 (2022). [3] Mehta, Harsh, et al. "Long range language modeling via gated state spaces." arXiv preprint arXiv:2206.13947 (2022).$$

We thank the reviewer for their useful feedback and for appreciating the strengths of our approach.

Below, we address the reviewer’s concerns:

Contributions are not well articulated in the introduction. The presentation needs to be sharpened a bit to highlight the main contribution - the effect of the novel Koopman model on RL in comparison to other generative models.

Following the reviewer’s suggestion, we have now rewritten the contribution section in the revised manuscript to better highlight the main contributions of our work.

It is not clear how this work is different from other Koopman operator-based RL methods.

We have further extended the discussion of related work on Koopman-based RL in Section 5, where we also contrast these with our approach.

Ablation with Hippo model as a backbone for offline RL is missing. The authors mention the relation to the Koopman framework and it would be interesting to see if the proposed model offers an improvement over HIPPO

We thank the reviewer for suggesting a suitable baseline. In the revised submission, we added the DSSM/S4D [1,2] model results as a backbone for offline RL. Note that the authors of [1] have shown that the DSSM with certain initialization of diagonal parameters approximates the convolution kernel derived from the HiPPO framework.[3] Please refer to Figure 3/Table 1, 2 and Figure 4/Table 3 for the new results. Our model is faster than the DSSM baseline with a similar number of parameters while being better/competitive in terms of performance.

Ablation on sequence length for RL tasks is not provided, i.e., do we get an improvement [………..] Therefore, this raises the question of whether we actually need a model capable of long-term predictions. I do not see a discussion on this subject.

We thank the reviewer for suggesting an ablation of sequence length for RL experiments. We provide those experiments for both RL and planning in Appendix H. Now, we address the reviewer’s concern regarding the use of long-term predictions. We would like to point out that the paper (MBPO) [4] the reviewer referred to is on model-based RL, which, as mentioned in the manuscript, is left for future work. In the future, we plan to expand our dynamics model to model stochastic transition dynamics while still ensuring parallel training. This will allow for easy integration of algorithms like MBPO that require uncertainty estimation with model ensembles.

For model-based planning, we observed that the results improved with higher rollout length, i.e., a horizon of 20 compared to 5. We have included the ablation of length for model-based planning in Figure 9 in the revised manuscript. However, for model-free RL with dynamics modelling as an auxiliary task, we notice that the performance drops when we increase the sequence length from 5 to 20. This is consistent with the ablation on the horizon provided in Figure 4 (and Section 5) of the SPR [5], where the authors show that the performance starts dropping after a sequence length of 7. Notably, even with a sequence length of just 5 our Koopman-based method achieves better performance.

$$We also thank the reviewer for all the suggestions provided as points 1-6 in the Questions section. We have incorporated them all in the revised manuscript.$$

[1] Gu, Albert, et al. "On the parameterization and initialization of diagonal state space models." Advances in Neural Information Processing Systems 35 (2022).

[2] Gupta, Ankit, et al. "Diagonal state spaces are as effective as structured state spaces." Advances in Neural Information Processing Systems 35 (2022).

[3] Gu, Albert, et al. "Hippo: Recurrent memory with optimal polynomial projections." Advances in neural information processing systems 33 (2020).

[4] Janner, Michael, et al. "When to trust your model: Model-based policy optimization." Advances in neural information processing systems 32 (2019).

[5] Schwarzer, Max, et al. "Data-Efficient Reinforcement Learning with Self-Predictive Representations." International Conference on Learning Representations. 2020.

We thank the reviewer for their positive assessment of our work and valuable suggestions.

Below, we address the reviewer’s concerns:

"While the same dynamics model can also be used for (III) model-based RL, we leave that direction for future work." Do you mean you would improve the model K further? or g_theta perhaps?

In the current submission, we have explored the integration of our Koopman-based model in model-based planning and model-free RL with dynamics modeling. Although our model can be applied to model-based RL, most state-of-the-art methods require stochastic dynamics modeling. As mentioned in Section 7, in the future, we plan to extend our framework to stochastic dynamics models and integrate it with model-based RL algorithms.

Can your approach fail even if eigenvalues are initialized well? After correct initialization, will the real part of the eigenvalues always stay negative? Are there any other failure scenarios?

We thank the reviewer for raising an interesting question. We perform ablation studies of eigenvalue initialization in Appendix G (Tables 4, 5, and 6): a) constant $\mu_i$s, where we fix $\mu_i$s to be a real negative number throughout training and b) learnable $\mu_i$s, where each $\mu_i$ is bounded between two real negative numbers, such as [-0.3, -0.1]. In both cases, $\mu_i$s remains negative once initialized, and our theory suggests that as long as that stays negative, the training should be stable. Additionally, we have not seen any failure cases in training with these initialization settings.

"we apply our dynamics model in model-free RL, " What does it mean to apply a model in the model-free setting? [.....] Instead you could mention more failure cases as part of the ablation studies perhaps.

We appreciate the reviewer’s suggestions. Current state-of-the-art model-free RL methods [1] utilize a dynamics model to improve the sample efficiency of existing model-free algorithms. These methods use a dynamics model to design a self-supervised auxiliary loss for representation learning by ensuring consistency in forward dynamics in the latent space. We have rewritten the first two sentences of Section 4.2.2 and have clarified this. Regarding the code, as pointed out by Reviewer 3syb, having the main code snippet for the Koopman dynamics modeling part can ease the understanding of our algorithm for practitioners. As part of the revised submission, we have also added a link to an anonymous GitHub repository with the code base for the offline RL experiments with the baselines in Section J. Since we are constrained by the page limit, we have attempted to provide the most important experimental details in the main text and refer the readers to the appendix for additional details.

Would be nice to have figures or a section explaining what you're solving, rather than how, e.g. inputs are in the pixel space, encoder is an MLP and the TD-MPC tries to solve X ... So some part of Section E from the appendix could go to the main text.

We have used Figure 1 to explain what we are solving in existing dynamics modeling methods used in RL and planning. We modified the starting paragraph of Section 4.2 to describe that the model-based planning and model-free RL with dynamics modeling is to provide data-efficiency in control tasks. In the original submission, we have also mentioned in Sections 4.2.1 and 4.2.2 whether the environment observations are in pixel space or state space. We believe the reviewer meant Appendix C. Due to page limitations, we moved the detailed background and explanation of these methods to the Appendix. We consider these details to be supplementary rather than central to the paper's main contributions. However, in the first paragraph of Section 4.2, we encourage interested readers, especially those in the reinforcement learning (RL) community, to consult Appendix C for more details.

It is not clear to me which network architecture you used for the encoder g_{\theta}, is it another MLP? As a follow-up to the above: how does the encoder quality effect your results? How do you come up with the encoder structure? Can the theory guide you here?

We thank the reviewer for pointing this out. The network architecture of $g_{\theta}$ depends on the input state. Section 3.1 of the modified submission has been updated to explain this. In theory, $g_{\theta}$ needs sufficient expressivity for the prediction errors to go down as we have offloaded the task of finding eigenfunctions to it. In practice, we have observed that empirically having a standard encoder architecture like an MLP for state spaces and a CNN for pixel spaces suffices.

I have read all of the points raised by the reviewers as well as the authors' replies. I would like to thank the authors for their careful rebuttal. I think this paper can be accepted, as it provides an incremental but solid improvement over the existing Koopman-operator-inspired-approaches, as far as I can tell.

However I keep my score the same, because of the incremental nature of the paper and the conceptual/theoretical lack of understanding of various issues that were raised in the reviews (e.g., lack of theoretical underpinning of "Koopman-approximations", all-too-experimental explanations of several design decisions etc.)

We thank the reviewer for their comments. It appears to us that the overall rating is inconsistent with the review, which mentions connections to the existing Koopman theory literature and other relatively minor issues. It would be helpful if the reviewer could emphasize any (remaining) major issue to give us the chance to respond during the remaining part of the discussion period.

I've had the opportunity to delve into Koopman operator theory in my past research. [......] Some recent works on Koopman operator theory have also showcased how a control-affine nonlinear system can be transitioned to a more generalized input-separable Koopman system, with bilinear and linear forms being special instances of these separable Koopman formats.

We thank the reviewer for pointing us to the additional relevant paper. However, our discussion of the control-affine setting and the resulting bilinear form is for the sake of completeness only. As discussed in Section 2.2, we assume a decoupling of state and control observables, leading us to equation (2). Moreover, we have cited the paper the reviewer suggested and have added more clarification on the control-affine setting assumption in Section 2.2.

There are alternative NN methods, such as Neural ODE, that have consistently demonstrated promising results for long-horizon predictions, in my experience. It might be valuable to consider or reference them in the context of this paper

It would be interesting to see the results with NeuralODE

While a Neural ODE is an interesting option for sequence modelling, it is not a common baseline in RL. The reason is that the strength or advantage of a Neural ODE is in modelling “irregular’’ sequences, where time steps are not equally distanced. Existing RL benchmarks use regular sequences. Moreover, for high-dimensional data, a “latent” neural ODE is used, which is comparable to an RNN but is generally slower; for example, see [1,2]. Even in [2], where they use a Neural ODE in RL, the resulting model is “solved” to produce RNN-like updates. In our paper, we compare to GRU, which is a stronger baseline that outperforms RNN for long-range sequence models.

The problem of long horizon prediction for both controlled and non-controlled dynamical systems using Koopman operator has been explored before;[...] I am aware that you already cited this reference but “twice” (please check)

We have corrected this double-citation error in the revised article. Indeed, as the reviewer points out, we did mention this use of the Koopman operator in past work. See Sections 2.1, 2.2 and 5 for mention of relevant work to our paper. A point of contrast is that our formulation enables efficient parallel training using convolution and the diagonal Koopman matrix.

While controlled systems are a primary focus in this paper, an additional section dedicated to control design could potentially improve the paper further.

We agree and would like to direct the reviewer to Section 2.2, which discusses control design in the context of Approximate Koopman. However, this section admittedly does not provide comprehensive coverage, and due to the page limit, we refer the reader to existing surveys such as [3].

On the topic of stability in Koopman-based learned matrices, there are several contributions, such as "Learning Stable Koopman Embeddings", and its subsequent related studies. It might be advantageous to discuss or incorporate the stability guarantees they present in the context of this paper.

We appreciate the reviewer’s suggestion to look into [4]. We have now cited this in Section 2.1 of Preliminaries, where we discuss how Koopman theory can be used for stability analysis of dynamical systems. We want to emphasize that, as discussed in Sections 3.2.1 and 3.2.2, stability in the context of our work refers to the stability while training these dynamics models and backpropagate error from further away in time. This is different from [4], which focuses on contraction analysis for modelling stable non-linear systems and derives a Koopman formulation that respects that. Moreover, we focus on dynamical systems with control input and extend our framework to solve vanishing or exploding gradients for long-range predictions in decision-making frameworks like RL and planning.

Dear reviewer xZWp,

As we approach the conclusion of the discussion phase today, we kindly request an acknowledgement of our response. We would be happy to address any remaining major concerns the reviewer might have.

Given the reviewer's initial negative assessment, it will be really valuable to know if they still have any major concerns after our rebuttal response. This will give us the chance to respond.

Thank You,

Authors