Koopman Embedded Equivariant Control

We thank the reviewer for their time and effort in reviewing our work. We appreciate the constructive feedback and suggestions, as well as the recognition of the strengths of our work. Below, we provide detailed responses to your comments, weaknesses, and questions:

(1) Clarification on the Novelty of Equivariance in Koopman-Based Modeling

Thanks for your questions. To clarify, our goal is to address what properties the embedding $g$ should satisfy in order to preserve the control effect in the latent space. We formally answer this question with two key properties: equivariance and isometry.

Equivariance can be expressed as $F^{\text{latent}} \circ g = g \circ F$ (which preserves the properties of flow $F$ in embedding space), and it is a more general concept compared to the Koopman operator. We use the Koopman operator because it naturally satisfies equivariant representation. We do not claim that the derivations in equations (6) and (7) are our core contributions. Instead, our use of the Koopman operator is motivated by its equivariant properties, and to the best of our knowledge, no paper has formally stated why Koopman is equivariant (see Appendix D). It naturally enables the equivariant embedding of both flow and vector fields.

(2) Novelty of Koopman Formalism and Optimal Control + Spectrum in Koopman Dynamics

• We thank the reviewer for providing the reference, and we will include this paper [2] in the our paper. As our answer in weakness 1, we do not claim the equations (6) and (7) are our novleties. And, in general, we have two major differences from [2]:

  • We obtain an analytical control policy from our derived invariant value function, whereas [2] relies on Model Predictive Control (MPC) for optimal control.
  • We integrate deep learning model, greatly improving the scalability of solving control problems compared to [2]. Our method even effectively addresses image-based control problems, which are far more challenging than the low-dimensional systems considered in [2].

• Thanks for your question on the specturm.

Our method can capture the mixed spectrum in deep learning setting. The generator $\mathcal{P}$ of the Koopman operator is a densely defined, unbounded operator with domain $\mathcal{D}(\mathcal{P}) \subset L^2(M)$. In our approach, we approximate $\mathcal{P}$ by constructing a compactified version, $\hat{\mathcal{P}}$, following the compactification procedure described in [1]. Specifically, it is shown in [1] that the operator $\hat{\mathcal{P}} = \Pi \mathcal{P} \Pi$ is a compact operator with a purely atomic spectrum, providing an approximation to the original unbounded generator $\mathcal{P}$. Here, $\Pi$ is a projection operator that maps $L^2(M)$ to the feature function space spanned by $g$, which is dense and countable in $L^2(M)$. The approximated operator $\hat{\mathcal{P}}$ can be expressed as $\hat{\mathcal{P}} = \lim_{t \to 0^+} \frac{\Pi \mathcal{K}_{t} \Pi - I}{t},$ consistent with our learning process as described in Equations (8) and (9) of our work. Moreover, $\hat{\mathcal{P}}$ achieves strong convergence with operator norm to $\mathcal{P}$ as $t \to 0^+$, implying that the spectral properties of $\hat{\mathcal{P}}$ approximate those of $\mathcal{P}$. This convergence also ensures that the spectral measures of $\hat{\mathcal{P}}$ approximate those of $\mathcal{P}$, effectively capturing both the atomic and continuous components of the Koopman spectrum. Consequently, the approximated Koopman evolution operator $\exp(\hat{\mathcal{P}} t)$ achieves strong convergence to $\mathcal{K}_t$, even when the Koopman operator has a mixed spectrum. This result is supported rigorously by Corollary 4 in [1], highlighting the quality of the approximation.

Hope the answer effectively address the reviewer's concern and we look forward to further comments.

• References

  [1] Das, Suddhasattwa, Dimitrios Giannakis, and Joanna Slawinska. "Reproducing kernel Hilbert space compactification of unitary evolution groups." Applied and Computational Harmonic Analysis 54 (2021): 75-136.

  [2] Goswami, Debdipta, and Derek A. Paley. "Bilinearization, reachability, and optimal control of control-affine nonlinear systems: A Koopman spectral approach." IEEE Transactions on Automatic Control 67.6 (2021): 2715-2728.

(1) It appears Fig. 3d and 3e are swapped. Also, in the ablation for isometry loss, please provide the reward for lambda_met=0, so the effect of isometric loss is clearer.

Thank you for pointing out the swapped captions; we have corrected this in the revised manuscript. Additionally, we have included $\lambda_{met}=0$ in Table 1, labeled as KEEC (w/o $\mathcal{E}_{met}$).

(2) Line 511, it is unclear what "optimal state became unobservable" means. It needs clearer definition and better quantification.

Thanks for your question. The upright position of the pendulum represents an optimal state, which corresponds to a saddle point in the state space. When the state space is embedded into a latent space without preserving the underlying metric, this optimal state cannot be observed, and the characteristic properties of the saddle point may no longer hold. We refined the statement about the 'optimal state not being observed' to clarify this.

(3) typos in Appendix F

Thanks for carefully reading our manuscripts. See our response in weakness (4).

(4) Line 1472, What do authors mean by "two strange attractors"?

Thank you for raising this question. We have corrected this in our updated manuscript to refer to "one of the saddle points."

(5) How is the wave equation solved?

The wave equation is integrated using the 4th-order exponential Runge-Kutta method, with the action term as the source term.

(6) Table 2 shows that noise is added to the wave equation, but not others. Why is so? How sensitive is KEEC to noise in the other cases?

We maintained the default experimental settings for each implementation as specified in OpenAI Gym [1]. In the wave equation control problem, our approach demonstrated strong robustness despite the complexity and sensitivity to noise, as shown in Table 1. This suggests that our method would show similar robustness in the other lower-dimensional cases as well.

• References

  [1] Brockman, G. "OpenAI Gym." arXiv preprint arXiv:1606.01540 (2016).

(7) MPPI and PCC use significantly shorter horizons than KEEC. What if the former two use the same longer horizon, or have KEEC using the shorter horizon?

We appreciate the reviewer’s comment. In our experiments, we followed the default settings from the original papers for MPPI and PCC to ensure fair comparisons. Adjusting horizon lengths would require re-tuning and could deviate from their intended use. The shorter horizons for KEEC may hinder its performance. The chosen horizon lengths align with those used in extended dynamic mode decomposition (eDMD) [1], and we will include this discussion in the revised manuscript.

• References

  [1] Kutz, J. Nathan, et al. Dynamic mode decomposition: data-driven modeling of complex systems. Society for Industrial and Applied Mathematics, 2016.

(8) Line 155, what do the authors mean by "didn't comprehensively map the vector field ..."?

Thanks for rasing this questions. We have corrected to comprehensively map the flows and vector field, whereas the "comprehensively" refers to both flows and vector fields of the original dynamics.

We would like to thank the reviewer once again for the valuable time and thoughtful feedback. We look forward to any further comments and will address your further questions if you have.

We would like to thank the reviewer for raising these important questions.

Please see our responses in the corresponding weaknesses.

Thank you for your thoughtful review and valuable questions. We greatly appreciate your time and effort in providing feedback to improve our work. Hope our answers address your questions and concerns effectively. We look forward to your further comments and insights.

(4) The defined equivariance/isometry losses are quite similar to some existing work, such as “DeepMDP: Learning Continuous Latent Space Models for Representation Learning”. Please do a comparison.

Thank you for your question. Our work fundamentally differs from [1] in both research scope and embedding methods.

• continuous setting V.S. probabilistic setting. Our approach is defined in a continuous space with a differential structure, whereas [1] is based on the Markov Decision Process (MDP).

• Embedding methods. Our embedding method first seeks an equivariant representation of the dynamics and then enforces a consistent metric between the original space and the latent space. In contrast, [1] ensures that the latent MDP maintains consistent performance by imposing the Wasserstein-1 metric.

• References:

  [1] Gelada, Carles, et al. "Deepmdp: Learning continuous latent space models for representation learning." International conference on machine learning. PMLR, 2019.

(5) Please clearly state your assumptions and scopes. For example, the analytical framework is based on the control-affine system in Eq. (1). So, the author must state the application domains.

Thank you for the suggestion. We agree and have explicitly stated our assumptions in Section 2 of the updated manuscript (see Lines 106-107), including that our framework is based on the control-affine system in Eq. (1).

(6) In Fig. 3 (d) and (2), the x-axis doesn’t match the caption. Try to check all figures.

Thank you for carefully reading our manuscript and pointing out the mismatch. We have corrected it (see Figure 3) and reviewed all the figures in the updated manuscript.

(7) Is the computation time in Fig. 3 training or testing time? You may need to compare both.

Thank you for pointing this out. The computation time in Fig. 3 is the testing time. Our focus is on the 'off-line training and online play' scenario, where test-time efficiency is crucial, making training time less critical in this scenario.

We thank the reviewer for the precious feedback, as well as your recognition of the strengths in our work. Below, we provide detailed responses to your comments, weaknesses, and questions:

(1) What is the drawback of traditional Koopman-based analysis? Why do we need to introduce the learning framework?

Thanks for your question. Rather than solely enhancing the traditional Koopman operator, our primary objective is to elucidate what properties the embedding function $g$ should satisfy to effectively model the underlying dynamical system for optimal control.

In the previous embedding methods for learning dynamics [1, 2, 3, 4, 5], it is insufficient to discuss how to embed with consistent dynamics and control policy. To the best of our knowledge, this work is the first to formally and mathematically investigate the essential properties for learning an optimal deep learning embedding tailored for control applications. Our analysis identifies equivariance and isometry as the two most critical properties for preserving control effects. We utilize the Koopman operator because it naturally serves as an equivariant representation of dynamics (see Appendix D) and is compatible with analytical solutions.

• References:

  [1] Watter, Manuel, et al. "Embed to control: A locally linear latent dynamics model for control from raw images." Advances in neural information processing systems 28 (2015).

  [2] Banijamali, Ershad, et al. "Robust locally-linear controllable embedding." International Conference on Artificial Intelligence and Statistics. PMLR, 2018.

  [3] Matsuo, Yutaka, et al. "Deep learning, reinforcement learning, and world models." Neural Networks 152 (2022): 267-275.

  [4] Hafner, Danijar, et al. "Learning latent dynamics for planning from pixels." International conference on machine learning. PMLR, 2019.

  [5] Levine, Nir, et al. "Prediction, consistency, curvature: Representation learning for locally-linear control." arXiv preprint arXiv:1909.01506 (2019).

(2) Please do a proofreading for all notations and equations. For example, in Lines 187-188, there is no z_t. Why do you need to define it? What is in Eq. (6)?

Thank you for pointing this out. We have carefully reviewed all notations and equations in the updated manuscript to ensure consistency and clarity. Specifically, we have addressed the issue in Lines 187–188, and $z_t=g(s_t)$ is the latent state, the same variable in Eq. (6).

$\mathcal{U}$ is the state-dependent (acutation) operator that maps the latent state $z_t$ to a linear operator acting on the control input $a_t$. The operator $\mathcal{U}$ represents how the control input $a_t$ influences the time evolution of $z_t$.

(3) How did the past methods learn $\mathcal{K}$?

Thank you for the suggestion. Existing methods [1,2,3] typically learned the Koopman operator by using a parameterized fully connected (FC) layer. Instead, in [4,5,6], using a Dynamic Mode Decomposition (DMD) -based approach that adaptively fits the $\mathcal{K}$ non-parametrically. Notably, in our work, we take a different approach by learning the operator $\mathcal{P}$, the generator of $\mathcal{K}$, instead of directly learning the $\mathcal{K}$. This key difference enables us to derive an analytical control policy, avoiding the need for numerical optimisation as required in methods like MPC, as in [7,8].

• References

  [1] Lusch, Bethany, J. Nathan Kutz, and Steven L. Brunton. "Deep learning for universal linear embeddings of nonlinear dynamics." Nature Communications 9.1 (2018): 4950.

  [2] Yeung, Enoch, Soumya Kundu, and Nathan Hodas. "Learning deep neural network representations for Koopman operators of nonlinear dynamical systems." 2019 American Control Conference (ACC). IEEE, 2019.

  [3] Weissenbacher, Matthias, et al. "Koopman q-learning: Offline reinforcement learning via symmetries of dynamics." International conference on machine learning. PMLR, 2022.

  [4] J. Morton, A. Jameson, M. J. Kochenderfer, F. Witherden: Deep dynamical modeling and control of unsteady fluid flows, Advances in Neural Information Processing Systems 31, 2018, pp. 9258–9268

  [5] J. Morton, F. D. Witherden, M. J. Kochenderfer: Deep variational Koopman models: Inferring Koopman observations for uncertainty-aware dynamics modeling and control, Proceedings of the 28th International Joint Conference on Artificial Intelligence, 2019, pp. 3173–3179

  [6] Y. Guo, M. Korda, I. G. Kevrekidis, Q. Li: Learning parametric Koopman decompositions for prediction and control. arXiv:2310.01124

  [7] Li, Yunzhu, et al. "Learning Compositional Koopman Operators for Model-Based Control." International Conference on Learning Representations.

  [8] Goswami, Debdipta, and Derek A. Paley. "Bilinearization, reachability, and optimal control of control-affine nonlinear systems: A Koopman spectral approach." IEEE Transactions on Automatic Control 67.6 (2021): 2715-2728.

Thanks for your question. If we understand correctly, you are asking the traditional Koopman analysis instead of deep learning methods.

Infinite-Dimensional Nature. The Koopman operator is fundamentally infinite-dimensional. Approximating this operator with finite-dimensional models can introduce inaccuracies and limitations, hindering the ability to fully capture the system's dynamics [1, 2].

Feature Function Selection. Selecting an appropriate feature basis is a challenging task for nonlinear dynamics. Traditional approaches rely on fixed feature functions, such as polynomials and Gaussian kernels. However, inadequate or poorly chosen feature functions may lead to incomplete or misleading representations of the system, thereby diminishing the effectiveness of the analysis [3, 4].

High-Dimensional Systems. Even when employing finite-dimensional approximations like Dynamic Mode Decomposition (DMD) [5, 6], the computational resources required can be substantial, particularly for high-dimensional or complex systems. This restricts the scalability of traditional Koopman methods for larger or more intricate systems, making real-time or large-scale applications difficult.

Thanks again. We have included it in our main text (see Page 5)

[1] Korda, M., & Mezić, I. (2018). Linear predictors for nonlinear dynamics: Extended dynamic mode decomposition. Proceedings of the National Academy of Sciences, 115(11), 2700–2705. DOI:10.1073/pnas.1706943114

[2] Budišić, Marko, Ryan Mohr, and Igor Mezić. "Applied koopmanism." Chaos: An Interdisciplinary Journal of Nonlinear Science 22.4 (2012).

[3] Brunton, Steven L., et al. "Modern Koopman theory for dynamical systems." arXiv preprint arXiv:2102.12086 (2021).

[4] Lusch, Bethany, J. Nathan Kutz, and Steven L. Brunton. "Deep learning for universal linear embeddings of nonlinear dynamics." Nature communications 9.1 (2018): 4950.

[5] Tu, J. H., Rowley, C. W., Luchtenburg, D. M., Brunton, S. L., & Kutz, J. N. (2014). On dynamic mode decomposition: Theory and applications. Journal of Applied Mechanics, 81(8).

[6] Korda, M., & Mezić, I. (2018). Linear predictors for nonlinear dynamics: Extended dynamic mode decomposition. Proceedings of the National Academy of Sciences, 115(11), 2700–2705. DOI:10.1073/pnas.1706943114

(1) Point (3) in the Weaknesses section. How would you justify the advantage of using the Koopman generator (instead of the discrete-time operator)? For example, a comparison to a variant of the proposed method constructed with a discrete-time setting would be helpful if any.

Please see our responses in Weakness points 3.

Thanks for your thoughtful review and valuable questions. We greatly appreciate your time and effort in providing feedback to improve our work. Hope our answers address your questions and concerns effectively. We look forward to your further comments and insights.

(2) As mentioned above, learning neural network observables for embedding dynamics states has been widely studied, not only by Li et al. ICLR 2020 (which the authors have cited), but also by many other researchers.

Thanks for providing these references. The learning-based Koopman operator has been used to solve dynamical systems. Most of them directly learn a next-step prediction (discrete-time) similar to E2C and PCC, as discussed in the second paragraph in the Introduction, such as [1, 2, 3, 4, 5]. However, our method aims to leverage the Koopman operator to preserve the ODE form (continuous-time) of the dynamics rather than the direct next-step prediction. Directly leveraging the ODE form can lead an analytical-from to improve the control performance, which other methods cannot achieve.

• References

  [1] J. Morton, A. Jameson, M. J. Kochenderfer, F. Witherden: Deep dynamical modeling and control of unsteady fluid flows, Advances in Neural Information Processing Systems 31, 2018, pp. 9258–9268

  [2] J. Morton, F. D. Witherden, M. J. Kochenderfer: Deep variational Koopman models: Inferring Koopman observations for uncertainty-aware dynamics modeling and control, Proceedings of the 28th International Joint Conference on Artificial Intelligence, 2019, pp. 3173–3179

  [3] M. Han, J. Euler-Rolle, R. K. Katzschmann: DeSKO: Stability-assured robust control with a deep stochastic Koopman operator, Proceedings of the 10th International Conference on Learning Representations, 2022

  [4] Y. Guo, M. Korda, I. G. Kevrekidis, Q. Li: Learning parametric Koopman decompositions for prediction and control. arXiv:2310.01124

  [5] D. Uchida, K. Duraisamy: Extracting Koopman operators for prediction and control of non-linear dynamics using two-stage learning and oblique projections. arXiv:2308.13051

(3) Although the authors claim that the proposed method is different from previous methods in terms of the treatment of the vector field (Lines 90-91), there seems to be no direct empirical comparison from this perspective. E2C may be the most relevant of the examined baselines but is not necessarily a valid reference to investigate the particular advantage of the proposed method. Elaborating more on this point would be helpful.

Empirical comparison with the discrete-time operator is beyond the scope of this paper, as our primary aim is to propose a formal and theoretically robust framework for comprehensive deep learning embedding. Embedding the vector field is crucial because the original dynamics of an affine-control system are represented as an ODE, which makes our framework highly generalizable. Another advantage of our approach is that there is no need to explicitly compute the actuation operator $\mathcal{U}$ in Eq. 16 after embedding the ODE, as it is learned during model training. As noted in our abstract, this result leads to an analytical control policy.

Thanks for your valuable comments. We have elaborated on the discussion in the introduction to better emphasize the advantages of our approach and clarify its distinction from existing methods.

We thank the reviewer for the precious feedback and comments. Below, we provide detailed responses to your comments, weaknesses, and questions:

(1) It is unclear which aspects should be evaluated as novelty. Continuous settings (vector fields) and isometry loss seem trivial. The optimal control part seems unclear.

Thank you for your insightful question.

Our contributions and novelties are as the following points:

• We are the first work to study what properties the embedding function $g$ should satisfy.

  Our primary objective is to elucidate what properties the embedding function $g$ should satisfy to effectively model the underlying dynamical system for optimal control. To the best of our knowledge, this work is the first to formally and mathematically investigate the essential criteria for learning an optimal deep learning embedding tailored for control applications. We identify two pivotal properties that the embedding $g$ must satisfy: equivariance and isometry (see the description in Section 2.2).

• We propose an embedding to satisfy the properties: Koopman-Operator-Based Auto-Encoder, and a value-based method leveraging this embedding.

  Guided by the principles of equivariance and isometry, we propose a Koopman-operator-based auto-encoder designed to satisfy these critical properties. This approach is comprehensively summarized in our abstract and elaborated upon in Sections 2.2 to 3.2 of the manuscript. To demonstrate the non-trivial nature of our contributions, we highlight the following key aspects:

  • Equivariance (Flow and Vector Fields):

    For continuous dynamical systems, the flow describes the system's evolution over time, while the vector field defines the instantaneous rate of change at each point in the state space. These two components are intrinsically linked, as the flow is generated by the vector field. The equivariant embedding approach is $F^{latent} \circ = g \circ F$, and Koopman naturally satisfies this property. On the other hand, the infinitesimal generator $\mathcal{P}$ of the Koopman operator automatically embeds vector fields into the latent space. By leveraging the exponential map, we estimate the embedding flow map as $\mathcal{K}_t = \exp(\mathcal{P}t)$, thereby ensuring that the latent representation accurately captures both the flow and vector field dynamics. From a theoretical perspective, a comprehensive embedding must contain the two components.

  • Isometry (Control Effect):

    Introducing an isometry loss in the learning process is a novel aspect of our approach for optimal control embeddings. In many control systems, control costs are defined using a quadratic form. Without preserving the metric information through isometric embeddings, the integral costs or value functions become distorted, leading to suboptimal control policies. By enforcing isometry, we ensure that the value function remains invariant under the embedding $g$, thereby preserving the integrity of control costs and enabling effective policy optimization.

  • Optimal Control:

    Our framework allows for the direct learning of a parametric quadratic latent value function. Utilizing this latent value function, we can derive an analytical solution for the control policy. It should be noted the benefits of the analytical solution are attributed to learned dynamics with its vector fields. Crucially, implementing the policy within the latent space produces effects equivalent to executing it in the original state space, ensuring consistency and reliability in control actions.

• Experimental Contributions

  Beyond the theoretical advancements, our work makes significant experimental contributions. We demonstrate that our deep learning-based embedding approach outperforms existing methods, particularly in handling image-input problems. By leveraging the properties of equivariance and isometry, our embedding facilitates more accurate and efficient state representations, leading to superior performance in tasks involving high-dimensional sensory inputs.

Thank you for the rebuttal. I maintain my score for the following reasons:

(1) Novelty

To the best of our knowledge, this work is the first to formally and mathematically investigate the essential criteria for learning an optimal deep learning embedding tailored for control applications. We identify two pivotal properties that the embedding must satisfy: equivariance and isometry (see the description in Section 2.2).

I do not think the equivariance loss can be claimed to be a part of the novelty. It has been used in most NN-based Koopman operator learning studies, for example in the papers I listed in my initial review and in many others. The authors might claim that the novelty lies in the treatment of continuous time, Koopman generators, but the definition in Eq. (10) based on $\exp(\mathcal{P} \Delta t)$ is a straightforward variant of the discrete-time version. In this sense I agree with Reviewer jz5b's comment in this thread.

As commented in my initial review, the isometry loss may comprise some sort of novelty.

I am still not sure what kind of novelty is claimed in the control part. However this is probably because my expertise is slightly off, I have not been extensively following studies involving OC/RL. I would withhold specific judgement here.

Beyond the theoretical advancements, our work makes significant experimental contributions. We demonstrate that our deep learning-based embedding approach outperforms existing methods, particularly in handling image-input problems.

In my understanding, image-input problems have already been addressed with DNN-Koopman-based control methods, for example as early as Morton et al. (2018):

• J. Morton, A. Jameson, M. J. Kochenderfer, F. Witherden: Deep dynamical modeling and control of unsteady fluid flows, Advances in Neural Information Processing Systems 31, 2018, pp. 9258–9268.

(2) Continuous time

However, our method aims to leverage the Koopman operator to preserve the ODE form (continuous-time) of the dynamics rather than the direct next-step prediction.

As mentioned above and pointed out by Reviewer jz5b, the training process of the method uses the Koopman generator only through the exponential map $\exp(\mathcal{P}\Delta t)$, which is almost the same as the next-step prediction; only the difference is whether $\Delta t$ is fixed or variable.

For example, Bevanda+ (2021) (picked up in terms of relevance, i.e., NN-based observable learning) deals with the continuous-time setting, assuming the time derivative of the state is available as data. I do not think such a setting is notably different from the common discrete-time setting (particularly in this context) either. Still, your model-training method is even closer to the discrete-time.

• Bevanda et al., Diffeomorphically Learning Stable Koopman Operators, arXiv:2112.04085

Directly leveraging the ODE form can lead an analytical-from to improve the control performance, which other methods cannot achieve.

So my point (3) in the initial review is about this thing. To support this claim ("to improve the control performance"), it seems important to compare the proposed method with a discrete-time variant, which would be a kind of ablation study. That is, to do a variant of the proposed method where only the continuous-time consideration is dropped.

(3) Experiment

Empirical comparison with the discrete-time operator is beyond the scope of this paper, as our primary aim is to propose a formal and theoretically robust framework for comprehensive deep learning embedding. Embedding the vector field is crucial because the original dynamics of an affine-control system are represented as an ODE, which makes our framework highly generalizable.

I see the control method is based on the continuous-time, vector-field-based formulation, but I do not think such a fact makes the comparison to the discrete-time variant of the proposed method out of the scope. Moreover, in my understanding, the examined baselines are based on discrete-time setting.

Thank you for your insightful question.

We employ a continuous-time Koopman bilinear form to theoretically derive and generalize the control law using the operators $\mathcal{P}$ and $\mathcal{U}$. For practical implementation in discrete-time environments, we calculate the flow using the exponential form, ensuring consistency with our continuous-time framework. Since truly continuous control laws are not feasible in reinforcement learning settings, we implement the control policy in a discretized manner.

(1) In equation (3), what are the values taken by $\tau$ in the sum? Is it at discrete time points?

Thank you for your question. In Equation (3), $\tau$ represents discrete time steps starting from the current step $t$, summing future rewards, as is the standard equation for value functions in reinforcement learning.

(2) In Figure 3(d-e), the magnitude of $λ_{met}$ is between 32 and 256, and the latent dimension is between 0.1 and 1.0. Should this be switched?

Thank you for pointing out the swapped captions; we have corrected this in the revised manuscript, see Figure 3.

(3) Question about operators $\mathcal{K}$ and $\mathcal{P}$

See our responses in weakness 1.

(4) The KEEC model architecture is given in Table 3. What architectures are used for the comparison methods (SAC, CQL, etc.)? It would be good to compare the number of parameters needed for each.

We appreciate the reviewer’s suggestion. The architectures for SAC, CQL, and other methods follow their official implementations with default parameter counts, as detailed in Appendix H.2 (lines 1690–1698). Table 3 outlines the architecture and parameters of KEEC. We will include a parameter comparison in the updated manuscript for greater clarity.

---

Thank you for your thoughtful review and valuable questions. We greatly appreciate your time and effort in providing feedback to improve our work. We hope our answers address your questions and concerns effectively. We look forward to your further comments and insights.

We thank the reviewer for the valuable feedback and your recognition of the strengths of our work. Below, we provide detailed responses to your comments, weaknesses, and questions:

(1) There is some confusion regarding the theory, particularly regarding the operator $\mathcal{P}$. Is the infinitesimal generator $\mathcal{P}$ an infinite or finite dimensional operator? The Koopman operator $\mathcal{K}$ is an infinite-dimensional operator and $\mathcal{K}= \exp(\mathcal{P})$, so the generator should be infinite-dimensional. However, in equations (6-8), $\mathcal{P}$ seems to be a finite-dimensional matrix.

(Along with the Question 3: Is the infinitesimal generator $\mathcal{P}$ an infinite or finite dimensional operator? The Koopman operator $\mathcal{K}$ is an infinite-dimensional operator and $\mathcal{K}=\exp{(\mathcal{P})}$, so the generator should be infinite-dimensional. However, in equations (6-8), $\mathcal{P}$ seems to be a finite-dimensional matrix.)

We thank the reviewers for raising these critical questions regarding the operator $\mathcal{P}$ and $\mathcal{K}$:

• In equations (6) and (7), $\mathcal{P}$ and $\mathcal{U}$ are still infinite-dimensional. We corrected the identity matrix $I$ in Line 238 to identify the operator. The infinitesimal generator $\mathcal{P}$ is inherently an infinite-dimensional operator. As detailed in Appendix D, both the Koopman operator $\mathcal{K}$ and the generator $\mathcal{P}$ operate within an infinite-dimensional function space.

• In practical applications, it is infeasible to represent true infinite-dimensional operators. Consequently, we approximate $\mathcal{P}$ and $\mathcal{K}$ with finite-dimensional operators, denoted as $\hat{\mathcal{P}}$ and $\hat{\mathcal{K}}$, respectively, as seen in the loss function in Equation (8). This finite-dimensional approximation is a common and effective strategy for handling infinite-dimensional operators [1]. By choosing a sufficiently large dimension for the approximated operator $\hat{\mathcal{P}}$, we ensure that the resulting Koopman operator $\hat{\mathcal{K}}$ achieves good convergence properties and that the approximation error remains controlled.

• References:

  [1] Schmüdgen, Konrad. Unbounded self-adjoint operators on Hilbert space. Vol. 265. Springer Science & Business Media, 2012.

(2) There is no justification for Lemma 3.3.

Thanks for your question. Lemma 3.3 is derived based on the findings from references [1] and [2], which we have cited in our paper appropriately. Specifically:

• Reference [1]: establishes the foundation when optimal control problems associated with the reward are equivalent;
• Reference [2]: supports this by demonstrating that the preservation of the metric ensures that the integral of rewards along trajectories in both the original and latent spaces are identical.

The derivation is based on the two references [1,2]

• Invariant Value Function: Since the rewards are the same in both spaces, the value function remains invariant under the embedding function $g$;
• Consistent Policy Execution: Consequently, executing the policy in the latent space yields the same control effect as executing it in the original space.

This invariance ensures that the optimal control policy derived in the latent space is equally effective when applied to the original space, thereby validating the consistency and reliability of our approach.

• References:

  [1] Jean, Frédéric, Sofya Maslovskaya, and Igor Zelenko. "On the projective and affine equivalence of sub-Riemannian metrics." Geometriae Dedicata 203.1 (2019): 279-319.

  [2] Maslovskaya, Sofya. Inverse Optimal Control: theoretical study. Diss. Université Paris Saclay (COmUE), 2018.

(3) The page limit was exceeded.

The page limit is ten this year, and our submission is within this limit. See ICLR call for paper: https://iclr.cc/Conferences/2025/CallForPapers.

(4) While KEEC is faster than MPC and MPPI, it is slower than standard RL methods such as SAC and CQL.

The speed of KEEC control can be easily addressed with engineering tricks. KEEC is also a model-based method in practice, which learns the value function and derives the greedy policy analytically. The cause of the slower control is Auto-Differentiation (auto-diff) $\nabla_{z} V_g$ in Eq. 16. However, according to the previous work [1] Eq. 8 and 9, the auto-diff can be avoided by learning this particular derivative directly.

• References

  [1] Levine, Nir, et al. "Prediction, Consistency, Curvature: Representation Learning for Locally-Linear Control." International Conference on Learning Representations (2020).