Paper: Efficient Dynamics Modeling in Interactive Environments with Koopman Theory (ICLR-2024)

Reviewer comment (Novelty):
Contributions are not well articulated in the introduction. It is not clear how this work is different from other Koopman operator-based RL methods. For example, I find the tools to compute long-term predictions quite interesting and a bit more detail on this in the introduction would strengthen the paper.

Author response:
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. 

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

Reviewer comment (Experiments):
Ablation on sequence length for RL tasks is not provided, i.e., do we get an improvement from the model being able to predict several steps ahead? This is an important ablation since it was shown by Janner et al that using SAC as a policy learning algorithm does not benefit from rollouts longer than a one-time step. One potential reason for this is the algorithm itself that learns only from one-step-ahead predictions. 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.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Presentation):
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. I think the authors can pivot to argue that their Koopman model should be considered as the base model for most of the RL problems (offline RL, planning, model-based RL etc). But this a matter of choice.

Author response:
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.

Reviewer comment (Evaluation):
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.

Author response:
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.

Reviewer comment (Presentation):
Longer/better figure captions would be nice.

Author response:
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.

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

Author response:
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

Reviewer comment (Writing):
Sometimes the discussions in the paper are not very clear, for instance: "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?

Author response:
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.

Reviewer comment (Theory):
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?

Author response:
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.

Reviewer comment (Writing):
"we apply our dynamics model in model-free RL," — What does it mean to apply a model in the model-free setting? Please explain not just in the appendix but also in the main text briefly. In general, much of the discussion in the Appendix could be streamlined with the text. Please organize the appendix and reduce its scope; code should be put in GitHub and linked, no need for printing your code in the appendix. Instead, you could mention more failure cases as part of the ablation studies perhaps.

Author response:
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.

Reviewer comment (Presentation):
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. Some part of Section E from the appendix could go to the main text.

Author response:
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.

Reviewer comment (Reproducibility):
It is not clear to me which network architecture you used for the encoder g_{\theta}, is it another MLP?

Author response:
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.

Reviewer comment (Theory):
This is not a problem only concerning this work but in general applies to all 'Koopman-operator-inspired-approaches': the connection to the theory is very very loose and hence these methods could instead be simply called linear latent-dynamics modelling. Do you agree? If not, please mention how the Koopman-operator theory guides you analytically and intuitively.

Author response:
We agree with the reviewer that Koopman theory in the context of modeling dynamics using deep neural networks usually boils down to using a linear latent dynamics model. However, there can be many formulations under different assumptions, as mentioned in Section 2. The theory can help us understand which ones we are making for the given formulation.

Reviewer comment (Theory):
As a follow-up to the above: how does the encoder quality affect your results? How do you come up with the encoder structure? Can the theory guide you here?

Author response:
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.

Reviewer comment (Evaluation):
In the optimized Koopman matrix, are the eigenvalues sparse, or are they all nonzero/complex? Can they be made sparse and is that interesting?

Author response:
This is an interesting question. Please note that having sparser eigenvalues is the same as operating in a lower dimensional latent space as we use a diagonal Koopman structure. We provide an ablation on the dimension of latent space in Appendix I. We observe that a generally higher dimension of the latent space leads to better performance but comes with an increase in the number of parameters. This is inline with the Koopman theory. However, the performance saturates after a point. Moreover, we have observed that the learned complex eigenvalues are nonzero.

Reviewer comment (Theory):
Please mention the limitation to deterministic environments in the introduction. Is this a limitation that all Koopman-operator approaches in the literature face?

Author response:
Stochastic dynamics indeed presents a challenge to the traditional formulation of Koopman theory. However, there are extensions of classical formulations that extends it to approximately model stochastic dynamics. As mentioned in Section 7, we leave this direction for future work.

Reviewer comment (Novelty):
I've had the opportunity to delve into Koopman operator theory in my past research. It's noted that control-affine systems transform to a bilinear control Koopman based system in the lifted space under particular conditions, as exemplified by Theorem II.1 in the paper titled "Advantages of Bilinear Koopman Realizations for the Modeling and Control of Systems with Unknown Dynamics". It's also worth mentioning that these bilinear forms may not be universally applicable to all nonlinear systems, especially those that do not adopt the control affine form. 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.

Author response:
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.

Reviewer comment (Novelty):
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.

Author response:
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.

Reviewer comment (Novelty):
The problem of long horizon prediction for both controlled and non-controlled dynamical systems using Koopman operator has been explored before, please see the Nature publication titled "Deep learning for universal linear embeddings of nonlinear dynamics", and in subsequent studies. I am aware that you already cited this reference but “twice” (please check).

Author response:
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.

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

Author response:
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].

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Novelty):
Considering the above, a more comprehensive literature review on the Koopman operator for controlled/non-controlled dynamical systems might enhance the paper's breadth.

Author response:
We have referenced relevant Koopman-based literature for dynamical systems in Sections 1 and 2.1. We have also referenced related literature on Koopman operators for controlled dynamical systems in Sections 2.2 and 5. In the revised manuscript we have added a few more relevant citations based on the reviewer’s suggestions.  Due to limited space, for a more thorough discussion, we now refer the readers to comprehensive literature on Koopman theory for controlled/non-controlled dynamical systems such as [3] in Section 2.2.

Reviewer comment (Reproducibility):
For the purpose of replication and validation, making the code available as open-source (while preserving the authors' anonymity) could be beneficial for the broader community.

Author response:
We have added a link to an anonymous GitHub repository with the code base for the offline RL experiments with the baselines in Appendix J. This is in addition to the pseudocode of the “Diagonal Koopman Dynamics model” in Appendix J of the original submission.

Reviewer comment (Experiments):
It would be interesting to see the results with NeuralODE.

Author response:
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.

Reviewer comment (Novelty):
Please do a more thorough literature review on Koopman operator theory for control and modelling.

Author response:
We have referenced relevant Koopman-based literature for dynamical systems in Sections 1 and 2.1. We have also referenced related literature on Koopman operators for controlled dynamical systems in Sections 2.2 and 5. In the revised manuscript we have added a few more relevant citations based on the reviewer’s suggestions.  Due to limited space, for a more thorough discussion, we now refer the readers to comprehensive literature on Koopman theory for controlled/non-controlled dynamical systems such as [3] in Section 2.2.