RoboKoop: Efficient Control Conditioned Representations from Visual Input in Robotics using Koopman Operator

Authors: hemant  kumawat
biswadeep  chakraborty
saibal  mukhopadhyay
Venue: robot-learning_org_CoRL_2024_Conference
Published: 2024-06-07 10:57:08

Abstract: Developing agents that can perform complex control tasks from high-dimensional observations is a core ability of autonomous agents that requires underlying robust task control policies and adapting the underlying visual representations to the task. Most existing policies need a lot of training samples and treat this problem from the lens of two-stage learning with a controller learned on top of pre-trained vision models. We approach this problem from the lens of Koopman theory and learn visual representations from robotic agents conditioned on specific downstream tasks in the context of learning stabilizing control for the agent. We introduce a Contrastive Spectral Koopman Embedding network that allows us to learn efficient linearized visual representations from the agent's visual data in a high dimensional latent space and utilizes reinforcement learning to perform off-policy control on top of the extracted representations with a linear controller. Our method enhances stability and control in gradient dynamics over time, significantly outperforming existing approaches by improving efficiency and accuracy in learning task policies over extended horizons.

Main paper: 
 introduction: 
 Building agents capable of executing intricate control tasks using high-dimensional inputs, like pixels is crucial in many real-world applications. Combining deep learning alongside reinforcement learning [1, 2, 3] , various methods have been developed for learning representations from visual data in robotics for executing intricate control tasks. Most of these learning algorithms can be fundamentally organized into two interconnected yet distinct research directions (refer to Figure 1 and Table 1 ): (1) Developing visual predictive models that learn the underlying dynamics [4, 5, 6, 7] in latent space by predicting the future visual observations. (Figure 1 -A ) (2) Using deep RL to perform tasks with non-linear modeling in an interactive environment [8, 9, 10] (Figure 1-B ). The first approach involves creating predictive models of the environment which are then leveraged for tasks or generating samples, which model-free methods can then use for learning. However, these models could have sub-optimal task performance. They may need either task-visual alignment [11, 12, 13] where pre-trained vision model's parameters are adapted to the task via end-to-end training or manual task and control parameters tuning to achieve good task performance. Alternatively, RL-based methods learn the representation by self-supervised auxiliary tasks but they may need a huge amount of interactions with the environment as the dynamics are learned implicitly in the form of policy networks. In general, all these approaches are not sample efficient i.e., they need to collect a large number of samples from environment interactions to achieve satisfactory task performance. Additionally, these models often use computationally expensive latent dynamics models like transformers, MLP, and RNN.In our work, we aim to learn task-conditioned representations (Figure 1-C ) from visual observations in a sample-efficient way without sacrificing task accuracy by modeling the dynamics of the representations linearly. An important advantage of linear systems is their generalized math-ematical frameworks, in contrast to nonlinear systems which have no overarching mathematical frameworks for general characterization of systems. Linear systems are very well defined by their spectral decomposition, enabling the development of generalizable and efficient control algorithms.  
background: 
 Contrastive Learning: Contrastive Learning has emerged as a powerful framework for learning representations that capture similarity constraints within datasets when there are no data labels. It operates on a principle analogous to performing a search within a dictionary, wherein the task involves identifying matches and mismatches-akin to positive and negative pairs-as if they were keys to a specific query or anchor point [32, 33, 9, 34, 35] . Methods like Curl [9] and To-KPM [16] employ contrastive encoders to directly learn representations from images and control the system. However, these methods require a large number of samples (approximately 1000k) to achieve good task performance, while Curl attempts to learn the task without any explicit dynamic model [9] .Dynamical Models: Prior works model system dynamics from visual observations through deep learning networks, mapping high-dimensional observation spaces to latent embeddings. These networks learn latent embedding evolution using methods like MLP [7] , RNNs [36] or neural ODE [37, 38] , enabling accurate future state predictions by extracting relevant features and understanding state relations. Methods like HNN [28] employ Hamiltonian dynamics, while Lag-VAE [24, 25] utilize Lagrangian dynamics for long-term prediction and control. However, these approaches often face low prediction accuracy and require system state supervision, such as prior object segmentation [25, 39, 27, 26] , to enhance Lagrangian model accuracy.Self-supervised learning and RL: Multiple works [40, 41, 42] have utilized self-supervised learning to develop world predictive models, which are subsequently employed for control tasks or specific tasks. In [7] , the authors learn a linearized latent space, which is then controlled using optimal control methods. However, there is no feedback from the control to guide the task in learning control-oriented features. Reinforcement Learning (RL) based methods [43, 44, 45] for directly controlling systems from pixels have witnessed significant advancements in recent years, leveraging deep learning to interpret complex visual inputs. However, one of the significant challenges in RL, especially when applied to environments with high-dimensional input spaces such as images (pixels), is sample inefficiency. This inefficiency leads to the requirement for an excessively large number of interactions with the environment for the agent to learn an effective policy.3 Proposed ModelThe stochastic gradient ∇θ L cst is an unbiased estimate of the true gradient:E[ ∇θ L cst (θ)] = ∇ θ L cst (θ). 
problem formulation: 
 We consider an unknown time-invariant dynamical system of the form: s(t+1) = F (s(t), u(t))+ξ where s(t) ∈ S ⊆ R n , u(t) ∈ U ⊆ R m , and ξ ∼ N (0, Σ ξ ) are the system state, control input and system noise respectively. Function F (s(t), u(t)) : S × U → S governs the transition of the states of the dynamical system and is assumed to be arbitrary, smooth and non-linear. We assume that we can only observe the system through the visual depictions x(t). Our objective is to identify a control sequence u 0:T that minimizes the cumulative task cost function c(x k , u k ) over T time steps by learning an implicit embedding function Φ : X → Z that maps the pixel space to some latent space. Further, the evolution of latent variable ż = g(z, u), g : Z × U → Z is parameterized by a neural network Λ. The network Λ is constrained to follow linear dynamics with a Koopman operator K given by Φ(x t+1 , u t+1 ) = KΦ(x t , u t ) = Az t + Bu t For control problems, we formulate the task cost function as the optimal control problem in latent space given by: minu 0:T -1 T -1 k=0 c(z k , u z ) subject to z t+1 = Az t + Bu t .We aim to learn this Koopman Operator and linear control policy u = π(s) with the Koopman embedding function mapping the high dimensional input to latent space conditioned on the control cost of the system. mentations, such as random cropping [9] , to the state x t , whereas negative samples result from applying analogous augmentations to all other states except x t . These samples are further processed by dedicated key and query encoders, which update their parameters based on the similarity measure between the samples. The output of the query encoder is then used to map the visual space to Koopman embedding space by bifurcating the query vector into real and imaginary components. This introduces a complex-valued Koopman embedding space, characterized by a spectral Koopman operator with learnable eigenvalues expressed as µ i + jω. Utilizing this embedding space and its spectral Koopman operator, our approach focuses on deriving optimal control strategies and their parameters through the iterative optimization of a Linear Quadratic Regulator (LQR) problem for a predefined finite time horizon, guiding the system towards a specified goal state. This goal state is also provided in visual space and subsequently mapped to the Koopman embedding space through the key encoder. Moreover, we construct value functions through dual Q-value functions, serving as critics within the Soft Actor-Critic (SAC) algorithm framework. Model updates are performed following the SAC methodology [15] , grounded on the cost incurred by the LQR controller. For training, we iteratively gather data batches from the environment, applying three specific losses to train our models: 1) SAC loss for the critic and Koopman parameters, 2) Contrastive loss for optimizing contrastive encoder parameters, and 3) Next latent prediction loss for regularizing the Koopman embedding space dynamics.a.) Contrastive Spectral Koopman Representations In our work, we diverge from traditional temporal modeling techniques, which typically process a sequence of consecutive frames to understand scene dynamics [46] . Instead, inspired by contrastive learning approaches [9] [16], we introduce a dual-encoder framework consisting of key and query encoders to extract nuanced visual representations. Each encoder processes high-dimensional inputs to generate corresponding sets of key and query samples. To enrich the diversity of these samples, we employ a random crop data augmentation technique on each input sequence X = {x i |i = 0, 1, 2, . . .}, creating positive samples x + i and negative samples from the augmented states of X \ {x i }. These samples are then embedded via the encoders to produce embeddings z q i , z + i , and z - j . Building on this, we map the query embeddings into a complex Koopman embedding space, by splitting the query embedding vector into real and imaginary components, to model continuous latent dynamics. This is complemented by a control embedding that maps robot control inputs to a complex-valued Koopman control space. We postulate that the latent and control spaces within it are independent, adhering to a model of continuous latent dynamics. This is mathematically represented as dz dt = Kz(t) + Lu(t), where the evolution of the system over time can be resolved through z(s) = e Ks z(0) + s 0 e K(s-t) Lu(t)dt. Here, e Ks denotes the matrix exponential. For practical applications, this continuous model can be converted into a discrete form as xt+1 = Kx t + Lu t , achieved by applying Zero-Order Hold (ZOH) discretization to the original equations, resulting in K = exp(∆tK) and L = (K) -1 (exp(∆tK) -1)L. In scenarios where observations are uniformly sampled over time, this approach incorporates a learnable time step parameter ∆t for adjustment.To enhance efficiency in latent space forecasting, we introduce a diagonalized Koopman matrix, K = diag( λ1 , . . . , λm ), facilitating the use of an m × (τ + 1) Vandermonde matrix V [17] for straightforward row-wise circular convolutions. This method reduces the computational load by avoiding complex matrix operations. EQUATIONand a prediction loss, L pred , focused on Mean-Squared Error (MSE),EQUATIONThese losses drive the learning of model parameters to accurately predict future latent states while learning representations from the visual input.b.) Koopman Operator and Eigenspetrum Initialization: To construct representations that accurately capture the system dynamics, it is imperative to initialize the Koopman operator appropriately.The dynamic system's future state predictions can be expressed as: EQUATIONwhere λi = µ i + iω i denotes the eigenvalues of the diagonal Koopman operator. From a linear stability perspective, we would like to place the real values of these µ i in a negative real plane with large values so that the system reaches asymptotic stability faster. However, this configuration may lead to exploding gradient issues, a direct consequence of the exponential increase in gradient magnitudes proportional to the eigenvalues' real parts.Lemma 1 (Exponential Scaling of Gradient Norms): For the discrete dynamic system depicted in Equation 4, where L pred = L t is defined as the latent prediction loss in Equation 3, let z t represent the latent state and c t = Lu t the control input at time step t. The gradient norms of L t with respect to the j-th components of z j,t and c j,t exhibit exponential scaling compared to the gradient norms at time step t+1, governed by the real parts of the system's eigenvalues, µ j : ∂Lt ∂zj,t = e ∆tµj ∂Lt ∂zj,t+1 , where ∆t denotes the time increment, and µ j the real part of the j-th eigenvalue, underscoring the link between the gradient norms of the loss function, the latent representations, the control inputs, and the system's eigenstructure.To mitigate potential stability issues and constrain gradient magnitudes, we limit the real parts of the eigenvalues, µ j , within the interval [-0.5, -0.01], and initialize the imaginary parts, ω j , in ascending order of frequency as ω j = jπ similar to [17] . This approach ensures comprehensive frequency mode capture within the system dynamics, fostering stable and efficient learning of dynamic representations.c.) Task Conditioning with End-to-End learning : To learn the contrastive spectral representations and spectral Koopman operator, we condition these networks to learn a linear effective controller similar to [16] . Given Spectral Koopman embeddings z = ψ θ k (x) and its associated linear spectral latent system as shown in equation 4, we formulate a finite time horizon LQR problem in Koopman latent space as min u 0:TT k=0 (z k -z ref ) T Q(z k -z ref ) + u T k Ru k subject to z k+1 = Kz k + Lu k (8)where Q and R are state and control cost diagonal matrices and z ref denotes the representation of goal input. We solve the above equation in an iterative procedure to recursively update the solution for a small number of iterations, typically T < 10, which is adequate to obtain a satisfactory and efficient approximation. The control policy is then given by u ∼ π LQR (z|s) ≜ π LQR (z| K, L, Q, R) Further to optimize this controller towards the task, we draw inspiration from [9] [16], we maximize the following objective via two Q value estimators as used in soft actor-critic method [15] :L SAC = E t∼B [min i=1,2 Q i (z, u) -α log π SAC (u|z)]For an end-to-end training of the network, we iteratively collect trajectory data batches from the task environment E and use the three distinct loss objectives to train the networks in an end-to-end fashion parameter. The primary objective, L sac is used to optimize all the task parameters to learn efficient representations for the task. Additionally, we integrate contrastive learning loss L cst and model prediction loss L m to regularize the task parameter learning process and feature extraction.The variance of the stochastic gradient is bounded by a constant σ 2 :E[∥ ∇θ L cst (θ) -∇ θ L cst (θ)∥ 2 ] ≤ σ 2 . 
analytical results: 
 In this section, we aim to provide a theoretical analysis of the convergence behavior of our over network and show that our task network converges to an optimal policy given enough interactions with the environment. Specifically, we state the following theorems Theorem 1: (Convergence of Contrastive Loss via Gradient Descent) Let L cst (θ) be an L-smooth contrastive loss function for encoder parameters θ and assuming stochastic gradient descent (SGD) updates with learning rate α t satisfying Robbins-Monro conditions. If ∇θ L cst is an unbiased estimate of the gradient with bounded variance, thenlim t→∞ E[∥∇ θ L cst (θ t )∥ 2 ] = 0. (Proof)Theorem 2: (Convergence of Koopman Operator Approximations): Given (i) a discrete-time linear dynamical system with states z ∈ R n and control inputs u ∈ R m , evolving according to z k+1 = A true z k + B true u k , where A true ∈ R n×n and B true ∈ R n×m are the true system matrices; and (ii) the Koopman operator approximation estimates A and B such that z k+1 ≈ Az k + Bu k , the minimization of L m (A, B; z k , u k , z k+1 ) with respect to A and B converges to the true system matrices, i.e., lim n→∞ (A, B) = (A true , B true ), where n represents the number of observations. Theorem 3: (Convergence of the LQR Control Policy) Given a discrete-time linear system characterized by state transition matrix A ∈ R n×n and control input matrix B ∈ R n×m and the LQR problem aims to minimize a quadratic cost functionJ = ∞ k=0 (x ⊤ k Qx k + u ⊤ k Ru k ) with Q ≥ 0 and R > 0, the DARE solution P i+1 = A ⊤ P i A -A ⊤ P i B R + B ⊤ P i B -1 B ⊤ P i A + Q,converges to P * , ensuring that the optimal control gains G * yield a stable and optimal control policy.Theorem 4 (LQR within SAC Optimizes Koopman Control Policy): Let L sac be the SAC loss for a given policy π sac (u|z) integrated with the LQR control policy π LQR (z|G) in a latent space Z, derived via the Koopman operator theory for a nonlinear dynamical system. If the SAC loss L sac is Lipschitz continuous with respect to the parameter set Ω = {Q, R, A, B, ψ θ } and L sac is bounded below, then applying gradient descent updates on Ω to minimize L sac guarantees convergence to a stationary point of L sac .Proofs for these theorems are provided in Appendix Section C Building on these theorems, it is established that our networks, given an adequate number of samples from the training data, are guaranteed to converge to an optimal task policy.The learning rate α t satisfies the Robbins-Monro conditions:∞ t=1 α t = ∞, ∞ t=1 α 2 t < ∞. 
empirical results: 
 In this section, we conduct simulated experiments primarily to explore these two questions: Firstly, Can our methodology effectively achieve desirable task performance using linear spectral Koopman representations across diverse environments exhibiting varying nonlinear dynamics? Secondly, Is it feasible to develop a globally linear model that aligns closely with the latent space, demonstrating resilience to noise and external disruptions, while maintaining efficiency in sampling and computation?Experimental Settings and Baselines In our work, we implement the model that operates within a spectral Koopman space, specifically designed with a 128-dimensional complex space. This model incorporates a control embedding that maps control inputs to a 128-dimensional space as well, and all the models are trained for 100k and 500k timesteps. We use five distinct robotic control tasks from the DeepMind Control Suite [19] , each characterized by varying state and action space dimensionalities: reacher easy, reacher hard, walker, ball in cup, cheetah, and cartpole. For comparative analysis, we implement several baseline models. CURL [9] is a model-free RL algorithm using a contrastive encoder for latent space control learning. Our approach, while also using a contrastive encoder, uniquely learns a complex-valued linear latent space. We also examine the To-KPM [16] , which employs linear dense representations for dynamic modeling, whereas our model adopts a spectral Koopman form for efficient and stable system modeling. Additionally, we include modelbased RL algorithms: TD-MPC [47] and PlaNet [30] use real-time planning with different variants of the Cross-Entropy Method (CEM), while Dreamer [29] employs background planning. SAC-State [15] serves as an upper performance bound, receiving direct state input from the simulator. Our experimental framework and the detailed analysis of these baselines are thoroughly documented in Appendix A.1 and A.2.Evaluation Metric: In RL, agents learn from interactions with their environment, guided by task rewards. These rewards provide scalar feedback, allowing agents to refine policies for maximizing future rewards. For instance, in the cart pole system, the agent receives a fixed reward of +1 for every timestep the pole remains upright, within a predefined angle from the vertical. Therefore, the quicker the cartpole reaches a vertical position, the more rewards it will accumulate. In this work, we set the reward as negative of the LQR cost. The greater the reward, the lower the control cost, enabling a faster achievement of the desired state. We report the mean reward and the variance of the distribution across 5 random experiments for all the models. To operate in a sample-efficient regime, we trained our networks for only 100k steps, in contrast to the 500k-step training regimen utilized for all baseline models. Table 2 provides a detailed comparison of all the models for both 100k and 500k steps. For all the environments except Cheetah, our model achieves the highest reward, which notably corresponds to the negative of the control cost, outperforming all other models. We note that control models utilizing autoencoder-type prediction mechanisms, such as Koopman AE, SAC-AE, Dreamer and Planet, underperform, with SAC-AE, despite its similar RL exploration strategy to ours, performing particularly poorly. To-KPM, with its dense Koopman model, also yields comparatively lower rewards than our model. Furthermore, when comparing our model against state-space model-specifically, state-based SAC -we observe that our approach surpasses the state-based controllers in terms of reward, even with a limited training duration of 100k steps, and exhibits very low trial variance for rewards. This underscores the high expressive power and task-guided representation capabilities of our model. In the cheetah task evaluation, however, our model has the third-best reward in the 100K step evaluation, only behind DrQ and SLACv1. This exception may be attributed to these models' nonlinear dynamics models with non-linear next step prediction and similar encoder to our model, making them optimized for state space exploration through deep RL. For more ablation studies of our model, please refer to Appendix D.Computational Efficiency for Linear vs Non-Linear dynamics In this section, we evaluate the computational efficiency of state-of-the-art predictive models in conjunction with linearized control mechanisms. Specifically, we implement the MLP dynamics model as described in [9] and the Dense Koopman model from [16] for comparison. Additionally, to assess our method against parallelizable and causally structured models, we incorporate the Transformer model [51, 31] and a recurrent model outlined in [29] , which are utilized to simulate system dynamics over more than 100 future steps while concurrently learning a linear control strategy for the system. Figure 3 illustrates the Multiply-Accumulate (MAC) operations required by each model, from which we deduce that our approach of using spectral Koopman necessitates the lowest computational effort for both control and prediction tasks. This comparison underscores the efficiency of our method in managing the computational demands associated with dynamic system modeling and control.Performance under external disturbances: In this section, we introduce an external force to the cart pole system during evaluation, denoted by F , which follows a uniform distribution F ∼ Uniform(a max , 0, a min ). The probability of applying this force is represented by p, where a max and a min correspond to the environment's maximum and minimum allowable control input values, respectively. Figure 4 illustrates the mean reward and Interquartile Mean (IQM) for the models under study. Notably, our models demonstrate robust control capabilities, maintaining high performance even when the probability of external disturbances is increased to 0.25. This resilience is evidenced by a minimal degradation in performance when compared to alternative approaches. 
conclusion and future work: 
 In this work, we introduce RoboKoop, a novel approach that integrates a contrastive encoder with spectral Koopman operators to learn visual representations guided by task learning. We demonstrate that RoboKoop surpasses current state-of-the-art methods, achieving superior performance while maintaining sample efficiency and robustness to noise and external disturbances. Looking forward, we aim to further enhance sample efficiency by exploring alternatives to Soft Actor-Critic (SAC) methods, potentially reducing the number of samples required for effective learning.A Model Details and Experimental Settings In our research, we introduce RoboKoop, an algorithm distinguished by its sample efficiency, which processes pixel-based inputs to simultaneously learn linear dynamics and develop an effective control policy. This algorithm demonstrates versatility across a broad spectrum of environments. We have rigorously tested RoboKoop against continuous control challenges within the DeepMind Control Suite. Our selection of these particular tasks is grounded in several critical considerations:1. Existing baseline methods exhibit suboptimal performance on these tasks, highlighting a gap that RoboKoop aims to fill.2. Recent advancements have introduced both model-free and model-based strategies aimed at enhancing the sample efficiency of similar algorithms. Our work contributes to this ongoing dialogue by presenting an alternative approach.3. The performance metrics obtained from these simulated tasks are highly indicative of realworld applicability, underscoring the practical relevance of our findings in broader contexts.Cartpole Swingup This task is centered around the goal of swinging up a pole, initially in a downward orientation, attached to a moving cart, and then maintaining its upright position. Success in this task requires the precise application of forces to the cart, navigating through a 4D state space that represents the cart-pole system's kinematics, complemented by a 1D control space for force application. 
 
 Appendix: The objective here is to orchestrate the movements of a simulated planar cheetah to achieve rapid and stable running. This involves managing an 18D state space that encapsulates the kinematics of the cheetah's entire body, including its joints and limbs, while employing 6D torques as controls to manipulate the joints for optimal locomotion.Reacher The Reacher task is designed to test precise motor control by requiring an agent to maneuver a simulated two-joint robotic arm to a target location in a 2D plane. This task involves navigating through an 11D state space that includes the positions and velocities of the arm's joints, as well as the position of the target. The control space is 2D, representing the torques applied at each joint. Success in this task is measured by the agent's ability to accurately and efficiently move the arm to the target position and maintain it there.In the Ball in Cup task, the objective is to control a simulated robot arm to swing and catch a ball attached to a string in a cup. This task is particularly challenging due to the nonlinear dynamics involved in swinging the ball and the precision required to catch it in the cup. The environment's state space is 8D, capturing the positions and velocities of the ball and the robot arm, as well as the angular position of the cup. The control space is 3D, representing the forces applied to the robot arm to achieve the desired swing motion. Success in this task requires a combination of dynamic coordination and precise timing.Walker The Walker task involves controlling a bipedal robot to achieve stable and efficient locomotion. The state space for this task is 17D, encompassing the kinematic properties of the robot's body and legs, including joint positions and velocities. The control space is 6D, corresponding to the torques applied to the robot's joints. The objective is to navigate the robot through various terrains, maintaining balance and forward motion. Success in this task is determined by the robot's ability to move swiftly and stably without falling.Table 3 provides a comprehensive enumeration of the hyperparameters employed in our model, along with detailed descriptions of each parameter. For To-KPM [16] also, we use the same hyperparameters as our model for a fair evaluation.For an end-to-end training of the network, we iteratively collect trajectory data batches from the task environment E and use the three distinct loss objectives to train the networks in an end-to-end fashion. The primary objective, L sac , is used to optimize all the task parameters to learn efficient representations for the task. Additionally, we integrate contrastive learning loss L cst and model prediction loss L m to regularize the task parameter learning process and feature extraction. Algorithm 1 describes the loss computation and parameter updates for the model. Collect roll-outs τ η by executing policy π LQR , store in D.Sample batch B from D.Compute L cst : Update ψ θ to minimize L cst .Compute L sac : Update Q, R, K, L to minimize L sac (excluding ψ θ ).Compute L m : Update K, L to minimize L m . 0: end for=0 man operator. These limitations are substantiated by the instability of poles (refer to Figures 4 and 5 in Section 5 of our paper) and underscore the efficiency of our approach.Planet is a model-based agent that discerns environment dynamics directly from pixels, facilitating action selection through online planning within a compact latent space. The latent space is structured around a recurrent state-space model, which is computationally intensive, as evidenced in Section 5 (Figure 6 ). Additionally, its emphasis on multi-step prediction in pixel space compromises sample efficiency, necessitating extensive interactions with the environment.The Koopman AE methodology leverages a soft actor-critic policy, underpinned by a regularized autoencoder (AE), to learn a latent space model atop AE features. Unlike Planet, this approach also explicitly models dynamics using a Koopman operator. In contrast, our method eschews the use of VAEs or AEs for pixel reconstruction, opting instead to learn features via contrastive learning alone. This strategy ensures the prioritization of task-relevant features over the reconstruction of pixel space, enhancing task efficiency and model performance.Definitions and Assumptions 1. Smoothness: The function L cst is assumed to be L-smooth with respect to θ, meaning it has Lipschitz continuous gradients:∥∇L cst (θ 1 ) -∇L cst (θ 2 )∥ ≤ L∥θ 1 -θ 2 ∥, ∀θ 1 , θ 2 .Theorem 1.: Let L cst (θ) be an L-smooth contrastive loss function for encoder parameters θ and assuming stochastic gradient descent (SGD) updates with learning rate α t satisfying Robbins-Monro conditions. If ∇θ L cst is an unbiased estimate of the gradient with bounded variance, thenlim t→∞ E[∥∇ θ L cst (θ t )∥ 2 ] = 0.Proof: Given the Lipschitz continuity of ψ θ , and assuming the loss L cst inherits this property with respect to θ, the Descent Lemma can be applied. The lemma states that for a Lipschitz continuous function f with Lipschitz constant L,f (x + ∆x) ≤ f (x) + ∇f (x) ⊤ ∆x + L 2 ∥∆x∥ 2 .Given the L-smoothness of L cst , we have for any θ 1 , θ 2 :L cst (θ 2 ) ≤ L cst (θ 1 ) + ∇L cst (θ 1 ) ⊤ (θ 2 -θ 1 ) + L 2 ∥θ 2 -θ 1 ∥ 2 .Substituting the gradient descent update θ t+1 = θ t -α t ∇θ L cst (θ t ):L cst (θ t+1 ) ≤ L cst (θ t ) -α t ∇L cst (θ t ) ⊤ ∇θ L cst (θ t ) + Lα 2 t 2 ∥ ∇θ L cst (θ t )∥ 2 .Taking expectations on both sides, and using the fact that E[ ∇θ L cst ] = ∇ θ L cst (unbiased gradient estimates) and the bounded variance assumption:E[L cst (θ t+1 )] ≤ E[L cst (θ t )] -α t ∥∇ θ L cst (θ t )∥ 2 + Lα 2 t 2 (σ 2 + ∥∇ θ L cst (θ t )∥ 2 ).Rearranging the terms, we aim to show that:α t (1 - Lα t 2 )∥∇ θ L cst (θ t )∥ 2 ≤ E[L cst (θ t )] -E[L cst (θ t+1 )] + Lα 2 t σ 2 2 .Given α t satisfies the Robbins-Monro conditions and 1-Lαt 2 > 0 for sufficiently small α t , summing both sides over t and applying the law of total expectation give:∞ t=1 α t (1 - Lα t 2 )E[∥∇ θ L cst (θ t )∥ 2 ] ≤ L cst (θ 1 ) -L cst (θ * ) + ∞ t=1 Lα 2 t σ 2 2 ,where θ * is a local minimum of L cst .Given the right-hand side is bounded (due to the boundedness of L cst and the conditions on α t ), and∞ t=1 α t (1 -Lαt 2 )= ∞, it follows from the quasi-martingale convergence theorem and the Robbins-Monro conditions that:lim t→∞ E[∥∇ θ L cst (θ t )∥ 2 ] = 0.This implies that, in expectation, the gradient norm converges to 0, indicating convergence to a stationary point. Now using the Polyak-Łojasiewicz condition, it can be shown that this is a local minimum.The exact form of L cst and its gradient ∇ θ L cst . The Lipschitz constants for ψ θ and L cst Conditions under which the stochastic gradient is an unbiased estimate of the true gradient and has bounded variance. A suitable learning rate schedule α t that guarantees convergence.Theorem 2: Convergence of Koopman Operator Approximations: Given (i) a discrete-time linear dynamical system with states z ∈ R n and control inputs u ∈ R m , evolving according to z k+1 = A true z k + B true u k , where A true ∈ R n×n and B true ∈ R n×m are the true system matrices; and (ii) the Koopman operator approximation approach, which seeks to estimate matrices A and B such that z k+1 ≈ Az k + Bu k , based on a loss function L m (A, B; z k , u k , z k+1 ), the minimization of L m with respect to A and B over the observed data converges to the true system matrices, i.e.,lim n→∞ (A, B) = (A true , B true ),where n represents the number of observations.We model the evolution of the system's state as a linear regression problem, where:• Z next is the matrix of next states z k+1 ,• X is the design matrix composed of current states z k and control inputs u k ,• Θ is the parameters matrix to be estimated, combining A and B,• ϵ is the error term.The equation Z next = XΘ + ϵ encapsulates this linear relationship.The objective function to minimize the difference between the predicted next states and the actual next states, quantified by the Frobenius norm of their difference can be written as:L m = ∥Z next -XΘ∥ 2 F ,where ∥ • ∥ F denotes the Frobenius norm. To minimize L m , we calculate the gradient of the loss function with respect to Θ and set it to zero:∇ Θ L m = -2X ⊤ (Z next -XΘ) = 0.Solving this equation for Θ gives: Θ = (X ⊤ X) -1 X ⊤ Z next . This is the least squares solution, providing the best estimate of Θ given the data.With the assumption that the observations X and Z next sufficiently cover the entire state and control input space and as the number of observations n approaches infinity (N → ∞), the matrices X ⊤ X and X ⊤ Z next will converge to their expected values. This ensures that the estimated parameters Θ, which combine A and B, converge to the true system matrices A true and B true that govern the system's dynamics.The solution involves setting the gradient of L m with respect to Θ to zero, leading to:∇ Θ L m = -2X ⊤ (Z next -XΘ) = 0Solving this equation yields the estimate for Θ:Θ = (X ⊤ X) -1 X ⊤ Z nextGiven a sufficiently diverse and large dataset (n → ∞), the estimates converge to the true system dynamics because the matrices X ⊤ X and X ⊤ Z next approach their expected values, ensuring the estimated parameters (A and B) converge to the true parameters (A true and B true ).This proof assumes sufficient data coverage across the state and control input space, which guarantees the convergence of the Koopman operator approximations to the true system dynamics, thereby validating the theorem.Theorem 3: Convergence of the LQR Control Policy Given a discrete-time linear system characterized by state transition matrix A ∈ R n×n and control input matrix B ∈ R n×m and the LQR problem aims to minimize a quadratic cost functionJ = ∞ k=0 (x ⊤ k Qx k + u ⊤ k Ru k )with Q ≥ 0 and R > 0, the iterative solution to the Discrete-time Algebraic Riccati Equation (DARE)P i+1 = A ⊤ P i A -A ⊤ P i B R + B ⊤ P i B -1 B ⊤ P i A + Q,converges to the optimal solution P * for the LQR problem, ensuring that the optimal control gains G * = -(R + B ⊤ P * B) -1 B ⊤ P * A yield a stable and optimal control policy.To prove the convergence of the Linear Quadratic Regulator (LQR) control policy, we focus on the discrete-time setting, where the goal is to design an optimal control policy that minimizes a given cost function. The essence of the proof involves showing that the solution to the Discrete-time Algebraic Riccati Equation (DARE) converges to a unique positive semidefinite matrix, which then defines the optimal control gains.We are given a discrete-time linear system:x k+1 = Ax k + Bu k ,and aim to minimize the infinite-horizon quadratic cost function:J = ∞ k=0 x ⊤ k Qx k + u ⊤ k Ru k ,where Q ≥ 0 (positive semidefinite) and R > 0 (positive definite) are the state and control weight matrices, respectively.The optimal control policy for this problem can be derived using dynamic programming, leading to the DARE:P = A ⊤ PA -A ⊤ PB R + B ⊤ PB -1 B ⊤ PA + Q,where P is the solution that defines the optimal cost-to-go matrix.The convergence of the LQR control policy essentially means proving that the iterative solution to the DARE converges to a unique positive semidefinite matrix P * . Here are the key steps:1. Monotonicity and Boundedness:To prove that the sequence {P i } generated by the Discrete-time Algebraic Riccati Equation (DARE) iterations is monotonically decreasing and bounded below, thus ensuring convergence, let's delve into equations and inequalities that illustrate these properties. Consider the iterative update rule for the DARE:P i+1 = A ⊤ P i A -A ⊤ P i B R + B ⊤ P i B -1 B ⊤ P i A + Q,where:-A and B define the system dynamics, -R is the control weighting matrix, which is positive definite (R > 0), -Q is the state weighting matrix, which is positive semidefinite (Q ≥ 0), -P i is the cost-to-go matrix at iteration i.To show that P i+1 ≤ P i , we need to establish that P i -P i+1 is positive semidefinite for each i. The Riccati update aims to minimize the cost function J i associated with using the control law derived from P i . Therefore, if we define the cost reduction as ∆P i = P i -P i+1 , we seek to show that ∆P i ≥ 0 (i.e., ∆P i is positive semidefinite).Starting from the DARE update rule and rearranging terms gives us:∆P i = P i -P i+1 = A ⊤ P i B R + B ⊤ P i B -1 B ⊤ P i A,Given that R > 0 and P i is positive semidefinite, it follows that the right-hand side of the equation above is positive semidefinite. This is because the term inside the parenthesis, R + B ⊤ P i B, is positive definite, making its inverse also positive definite, and thus ∆P i is positive semidefinite, indicating that P i+1 ≤ P i .The sequence is bounded below by the zero matrix, given that the cost-to-go matrices P i represent quadratic cost functions which are non-negative:P i ≥ 0 ∀i,implying that the sequence cannot decrease indefinitely and is bounded below by a matrix where all elements are greater than or equal to zero. Given the monotonicity and boundedness of the sequence {P i }, it follows from the Monotone Convergence Theorem for matrices that the sequence converges to a limit, say P * , which is the solution to the DARE and represents the optimal cost-to-go matrix:lim i→∞ P i = P * ,where P * satisfies the DARE and thus confirms the optimality and stability of the LQR control policy derived from it.By establishing the monotonic decrease and boundedness below of the sequence {P i }, we have shown that this sequence converges to a matrix P * that minimizes the LQR cost function. This P * is the fixed point of the DARE, providing the optimal cost-to-go estimate and ensuring the stability and optimality of the LQR control policy derived from it.2. Fixed Point Convergence: Under the assumptions that A, B, Q, and R satisfy certain controllability and observability conditions, it can be shown that the iteration converges to a fixed point.To prove that the limit of the sequence {P i }, denoted as P * , satisfies the Discrete-time Algebraic Riccati Equation (DARE) and is thus a fixed point of the iteration process, we employ the properties of convergence and continuity of matrix operations.Given the iterative process:P i+1 = A ⊤ P i A -A ⊤ P i B R + B ⊤ P i B -1 B ⊤ P i A + Q,we aim to show that, as i → ∞, P i → P * and that P * satisfies the DARE:P * = A ⊤ P * A -A ⊤ P * B R + B ⊤ P * B -1 B ⊤ P * A + Q.From previous steps, we have shown that the sequence {P i } is monotonically decreasing and bounded below, which guarantees convergence to a limit P * due to the Monotone Convergence Theorem for matrices.The operations involved in the iterative update rule, including matrix addition, multiplication, and inversion, are continuous functions of their arguments. This means that if a sequence of matrices {X i } converges to X, then the limit of a continuous function f (X i ) is f (X). The update rule can be seen as the application of a continuous function f to P i :f (P i ) = A ⊤ P i A -A ⊤ P i B R + B ⊤ P i B -1 B ⊤ P i A + Q.Given the convergence P i → P * , by continuity, we have:lim i→∞ f (P i ) = f ( lim i→∞ P i ) = f (P * ).This implies:P * = A ⊤ P * A -A ⊤ P * B R + B ⊤ P * B -1 B ⊤ P * A + Q,which is precisely the DARE. By showing that P * satisfies the DARE, we've proven that P * is a fixed point of the iteration process. This fixed point represents the solution to the DARE, establishing the optimality of the limit matrix P * for the LQR problem.Thus, by leveraging the properties of monotonicity, boundedness, convergence, and the continuity of matrix operations, we've demonstrated that the limit of the sequence {P i }, P * , satisfies the Discrete-time Algebraic Riccati Equation, making it the optimal solution and a fixed point of the iterative process.The convergence of the LQR control policy to an optimal solution involves demonstrating that the iterative solution to the DARE converges to a unique matrix that minimizes the cost function and that the corresponding control policy stabilizes the system. The proof relies on algebraic properties of the Riccati equation, control theory, and the system's controllability and observability conditions.Lemma: Given a loss function L that is Lipschitz continuous with respect to the parameters Ω, and bounded below, the sequence {Ω t } generated by the gradient descent updates:Ω t+1 = Ω t -η∇ Ω L(Ω t ),with a sufficiently small, fixed learning rate η > 0, converges to a stationary point Ω * , where ∇ Ω L(Ω * ) = 0.Proof: Given that L is Lipschitz continuous with Lipschitz constant L, we have for the gradient descent update:EQUATION)EQUATIONEQUATION)Choosing η: Select η such that 0 < η < 2 L , ensuring that:L(Ω t+1 ) ≤ L(Ω t ) -η - Lη 2 2 ∥∇ Ω L(Ω t )∥ 2 .Since L is bounded below, and L(Ω t+1 ) ≤ L(Ω t ) for all t, the sequence {L(Ω t )} is non-increasing and bounded. This implies convergence of the loss function values.The reduction of the loss at each step is proportional to the square of the norm of the gradient. If the sequence {Ω t } did not converge to a stationary point, the gradient norm would not approach zero, contradicting the boundedness and convergence of the loss function values. Therefore, the gradient norm must approach zero, i.e., ∇ Ω L(Ω * ) = 0, indicating convergence to a stationary point.Theorem 4: Let L sac be the Soft Actor-Critic (SAC) loss function for a given policy π sac (u|z) integrated with the Linear Quadratic Regulator (LQR) control policy π LQR (z|G) in a latent space Z, derived via the Koopman operator theory for a nonlinear dynamical system. If the SAC loss L sac is Lipschitz continuous with respect to the parameter set Ω = {Q, R, A, B, ψ θ } and L sac is bounded below, then applying gradient descent updates on Ω to minimize L sac guarantees convergence to a stationary point of L sac .Proof: Assume L sac satisfies the Lipschitz condition with Lipschitz constant L > 0, i.e.,|L sac (Ω 1 ) -L sac (Ω 2 )| ≤ L∥Ω 1 -Ω 2 ∥,for any Ω 1 , Ω 2 in the parameter space. Now, the update rule for the parameters Ω via gradient descent is given by:Ω t+1 = Ω t -η∇ Ω L sac (Ω t ),where η > 0 is the learning rate.Using Lemma 1, given L sac is bounded below and Lipschitz continuous, the sequence {Ω t } produced by the gradient descent updates will converge to a stationary point Ω * , characterized by:∇ Ω L sac (Ω * ) = 0.Hence we show the optimality and stability via LQR Integration. The integration of the LQR policy π LQR ensures that within the linear approximation of the dynamical system dynamics in the latent space Z, the SAC framework, enhanced with LQR, converges towards optimal control actions. The LQR component provides an optimal control policy for linearized dynamics around the current state and control, ensuring that the SAC algorithm's policy updates enhance both stability and optimality in control decisions.For a linear system z k+1 = Az k + Bu k , the LQR aims to minimize the cost function:J = ∞ k=0 z ⊤ k Qz k + u ⊤ k Ru k ,where Q ≥ 0 and R > 0. The optimal control law is u * k = -Kz k with K = (R + B ⊤ PB) -1 B ⊤ PA, where P solves the Algebraic Riccati Equation (ARE):P = A ⊤ PA -A ⊤ PB(R + B ⊤ PB) -1 B ⊤ PA + Q.The SAC algorithm seeks to optimize the policy π sac (u|z) by solving:max π E ∞ k=0 γ k (R(z k , u k ) + αH(π(•|z k ))) ,where H denotes the entropy of the policy, promoting exploration, and α is the temperature parameter that balances reward and entropy.Integration means adjusting the SAC optimization to include the LQR solution as a baseline or regularization term. The objective becomes:max π E ∞ k=0 γ k (R(z k , u k ) + αH(π(•|z k )) -λJ LQR (z k , u k )) ,where λ is a weighting coefficient, and J LQR is the LQR cost function introduced above. This formulation explicitly guides the SAC policy towards the LQR's optimal policy within the linear approximation of the dynamics.The optimal policy π * and the corresponding control law u * from this integrated optimization problem are given by (1) the policy π * that maximizes the augmented objective, and (2) the control law that minimizes the LQR cost, ensuring stability as P guarantees the eigenvalues of (A -BK) lie within the unit circle, ensuring the system's stability.The parameter update rule incorporating both SAC optimization and LQR regularization is given by: Ω t+1 = Ω t -η∇ Ω (L sac (Ω t ) -λJ LQR (Ω t )) ,where L sac and J LQR are differentiable with respect to Ω, ensuring that the gradient descent steps move the parameters towards minimizing the SAC loss while adhering to the LQR optimality criteria. Given the Lipschitz continuity and differentiability of L sac -λJ LQR , the updates guarantee convergence to a stationary point Ω * where ∇ Ω (L sac (Ω * ) -λJ LQR (Ω * )) = 0, encapsulating both the optimal policy in the SAC framework and the stability provided by the LQR control law. Thus, we've shown how this combined approach integrating LQR within the SAC framework leverages LQR's optimality and stability, guiding the policy updates in SAC towards enhanced control decisions. The integration explicitly incorporates the LQR's linear control optimality into SAC's nonlinear policy optimization, ensuring convergence towards optimal and stable control actions in the latent space Z. Thus, we show that under the conditions of Lipschitz continuity and boundedness of the SAC loss function, gradient descent optimization of the combined SAC and LQR policies in the Koopman latent space converges to a stationary point, optimizing the overall Koopman control policy. This integration not only leverages the strengths of both SAC and LQR but also ensures that the optimization process is theoretically grounded and guaranteed to reach a point of stability and optimality.In this section, we conduct an ablation study to identify which components of our network contribute to its superior performance with a limited number of training steps. First, to demonstrate the effect of nonlinearity, we use CURL [9] as a baseline. CURL features a contractive encoder similar to ours but employs nonlinear dynamics, unlike our spectral dynamics. For comparison with a linear dense model, we use ToKPM [16] , which relies on dense linear dynamics as opposed to our spectral model. Throughout the ablation studies, we demonstrate that our model outperforms both baselines. For this section, we present the results for models trained for 150,000 steps, as the other baselines showed poor performance when evaluated at 100,000 time steps.In Figures 5 and 6 , we present the eigenspectrum contour plots for the To-KPM model and our proposed model, respectively. Analysis reveals that the eigenvalues of the To-KPM model predominantly reside on the positive real axis, with an average value of approximately 0.4. Conversely, our model exhibits a symmetric distribution of eigenvalues across the imaginary axis, featuring an equal proportion of positive and negative real eigenvalues. This distribution aligns with an increasing trend of eigenvalues as per ω = jπ. Within the framework of the Koopman operator theory, negative eigenvalues signify that the system's observables exhibit exponential decay over time, as these eigenvalues are integral to the exponential term in the solution to the linear system governed by the Koopman operator. Hence, negative eigenvalues are indicative of stable observable behaviors, whereas positive eigenvalues suggest exponential growth in observables, pointing to instability. The presence of positive eigenvalues in the To-KPM model undermines its ability to learn stable representations from images using a finite-dimensional Koopman operator, leading to inferior performance compared to our model, which benefits from a balanced distribution of positive and negative eigenvalues.In this section, we conduct an ablation study to evaluate various initialization strategies for the Koopman spectral method, as detailed in Section 3.2. Our baseline configuration sets the Koopman operator's real value at -0.2, with frequencies arranged in increasing order. To assess the impact of initialization on performance, we explore three additional designs: (a) learnable real values with increasing frequency, (b) constant real values with random frequency, and (c) learnable real values with random frequency. This examination seeks to identify the initialization method that most effectively enhances the accuracy and stability of the spectral Koopman method. Table 4 presents the mean reward for all the models across cheetah and cartpole simulations. Our analysis reveals that the strategy of employing constant real values with increasing frequency for the imaginary component of the initialization yields superior results. This section delves into the resilience of our model when faced with imperfect sensing conditions during evaluation. We specifically examine its performance in two challenging scenarios: a.) Rainy Environment with Structured Noise: Unlike Gaussian noise, rain noise [52] presents a more structured and challenging interference, making it difficult for common denoising techniques to effectively mitigate [53][54] . We evaluate the model's performance under three distinct levels of rain density: Low (0.03), Medium (0.75), and High (0.0125). The comparison encompasses predictive models equipped with dynamic predictors, and the outcomes are depicted in Figure 7 . The results indicate a general degradation in control performance across all models under test, except for ours, which notably excels by achieving a reward of approximately 400. This demonstrates our model's superior capability to maintain effective system control even in the presence of high noise levels. b.) Imperfect Sensing due to Camera Jerk: To further assess our model's robustness, we intro-duce random camera jerks into the video input stream, simulating real-world sensing imperfections. Three levels of camera jerk are considered: Low jerk (SSIM > 0.8), Medium Jerk (SSIM between 0.4 and 0.5), and High Jerk (SSIM < 0.3). Our findings from Figure 8 reveal that our model consistently outperforms the others under these conditions as well. However, it's noteworthy that the performance gap between our model and the To-KPM model narrows as the jerk intensity increases, with both models exhibiting similar performance metrics at higher jerk levels. Conversely, methods based on autoencoders demonstrate significantly lower performance across all jerk conditions. These evaluations underscore our model's robustness and adaptability to imperfect sensing scenarios, highlighting its potential for real-world applications where sensing conditions are often less than ideal.In this experiment, we analyze the robustness of our model's control performance under the influence of Gaussian noise. We introduce zero-mean Gaussian noise to the input images with increasing standard deviation. Figure 9 illustrates the comparative performance of our model and to-kpm [16] model in the presence of Gaussian noise. We exclude models that under performed significantly from this figure, as their performance was too low to be meaningful. Notably, our model demonstrates exceptional resilience, maintaining high performance even under substantial Gaussian noise in the visual input.8th Conference on Robot Learning (CoRL 2024), Munich, Germany. 


Decision: Thank you for your submission to CoRL 2024. Reviewers particularly appreciated the computational and accuracy improvements gained by the proposed spectral Koopman embedding network, the detailed theoretic analysis, and the principled approach.

Reviewers raised concerns about the limited scope of experiments, generalizability, and some missing related work and comparisons. They were unclear about aspects of the training method and concerned that it may be too incremental.

Reviewers valued the authors' detailed and thorough responses, which were sufficient to raise at least one reviewer's initial recommendation. However, concerns linger about generalizability and, in particular, the applicability to real-robot deployment. The reviewers and AC encourage the authors to consider extending experiments to address these lingering concerns for a camera-ready revision.

Reviews: review: title: title: More clarifications and experiments may be necessary to make contributions stronger
paper_summary: paper_summary: The paper presents a new model called the Contrastive Spectral Koopman Embedding network, which is designed to learn efficient linear representations from high-dimensional visual data. This model uses reinforcement learning to perform off-policy control based on the extracted representations, employing a linear controller for this purpose.
main_review: main_review: Strength:
1. The Spectral Koopman Embedding network effectively leverages the diagonalized Koopman matrix to enhance the efficiency of Koopman matrix computation and accuracy
in learning task policies.
2.  The paper provides detailed theoretical analysis and proofs of the network's convergence behavior, demonstrating that the task network converges to an optimal policy given sufficient training steps.

Weakness:
1. The experiments conducted are limited in scope, as they only involve 6 of the DM control environments and lack real-world testing. This limitation raises questions about the generalizability and robustness of the conclusions drawn.
2. The explanation of model training in Section 3.2 is somewhat unclear. It seems that the loss function in Equation (3) is used to train both the Koopman matrix parameters ($\lambda_i$) and the encoder parameters. If this is correct, explicitly providing labels about which parameters are being optimized with each of the losses or providing an algorithm overview similar to To-KPM [16] will help reducing ambiguity.
3. The implementation of the spectral Koopman method and the overall algorithm seems incremental, building upon the To-KPM approach [16]. The primary difference introduced in this paper appears to be the diagonalized Koopman matrix, which reduces computational load. However, it is unclear why the convergence theoretical analysis presented cannot be applied to previous work such as To-KPM. The paper lacks a detailed explanation of the specific factors contributing to the observed improvements over To-KPM, which should be addressed to strengthen the findings.
questions: questions: 1. In Equation (2), the W seems to be introduced without prior context or explanation. Could you please clarify where it originates from? 

Also see weaknesses summaries in reviews above.
limitations: limitations: 3
review_summary: review_summary: The paper is mostly clear and provides thorough theoretical and empirical analysis. However, the limited scope of test environments and the similarities between algorithms used in this work and the previous work (specifically To-KPM [16]) makes it appear to be an incremental effort focusing mainly on reducing computation loads. These weaknesses limited the contribution, hence the weak reject recommendation.
title: title: extension of prior work on using koopman spectral representation
paper_summary: paper_summary: The authors consider control of nonlinear systems through visual scene feedback. They propose to employ spectrally linearized state representation (using Koopman operator) and then RL to perform off-policy control. The claim is that through the choice of complex eigenvalue representation one can achieve sample efficiency through this such task-specific and structure-exploiting representation.
main_review: main_review: Overall, the paper reads like a combination of known techniques to extend previous architectures in order to improve performance as a function of training iterations. The “global linearization” Koopman approach is promising and could better approximate the system dynamics (especially if there are any periodic motions which could otherwise be difficult to encode even with high-dimensional methods). The idea though is not new, as noted below. 

The difference from paper: Task-Oriented Koopman-Based Control with Contrastive Encoder from last year’s CoRL should be better clarified. Is the key point that the authors focus on the spectral representation? If so, it is surprising that that paper from last year did not already consider such representation since this is the natural choice to perform the “global linearization”.

The relative impact on performance of the “contrastive” part vs the “spectral” part of the approach is not clear. It would have been useful to show the role of each one through some type of ablation study. Better yet, if the core idea of the paper is the spectral representation, then it would have been clearer and to just focus on that. Otherwise, there are too many complicating details in the paper. 

There are key missing references, for instance a Nature Communications paper: Deep learning for universal linear embeddings of nonlinear dynamics, Lusch et al, 2018, which also contains a number of highly relevant references to Koopman representation for nonlinear dynamical systems. Note that the paper by Lusch also explains how to employ a spectral representation using complex eigenvalues. 

The video shows two simple examples, and one more complex one – the “Walker”, but it is not clear what the Walker portion is trying to show. Is this the final trained policy (doesn’t look great)?   It would have been more useful to show how a trained policy by the proposed method compares to another competitive existing method at several stages of the training process. 

Many references are missing details such as conference or journal names.
questions: questions: Explain the impact on performance of contrastive vs spectral part
Explain differences from those prior missing references
Explain novelty and contributions in view of above
limitations: limitations: 2
review_summary: review_summary: The novelty of the contribution is not clear; in addition, the source of improved performance is also not clear.
title: title: Review
paper_summary: paper_summary: This paper introduces RoboKoop, a novel contrastive spectral Koopman representations for visual RL; RoboKoop outperforms baseline visual RL algorithms on several DeepMind Control Suite tasks.
main_review: main_review: Strengths:
1. The paper is theoretically sound and proposes a principled algorithm for learning better visual representations for visual RL

Weaknesses:
1. Robokoop is only evaluated on Deepmind Control Suite; it's unclear whether RoboKoop is applicable to more realistic robot manipulation tasks in either simulation or the real world. 

2. Some missing prior works on applying Koopman operator for high-dimensional control, e.g., https://arxiv.org/abs/2303.13446

Additional Questions:

1. I wonder whether RoboKoop could be used in an imitation learning setting? Visual RL may be difficult to apply in the real-world, and the DeepMind control suite tasks are very far away from manipulation tasks in the real world. This paper can be made much compelling if results on manipulation benchmarks, such as MetaWorld or RoboSuite, in either imitation learning/RL setup. A recent work applies Koopman operator to dexterous imitation learning: https://arxiv.org/abs/2303.13446

2. Relatedly, in a lot of manipulation tasks, a single camera observation is not enough to capture all relevant visual information. How to extend RoboKoop to multi-view visual control settings?

3. RoboKoop seems general enough to be used with other RL algorithms. It would be good to show some results on alternative RL algorithms to SAC; this helps demonstrating RoboKoop's generality. 

4. Some background on Koopman operator should be provided before introduction of the method.

Overall, I think this is a sound paper that makes sufficient contribution to robot learning; however, its impact may be limited given the relative toy-ish setting of its evaluation.
questions: questions: My questions for rebuttal are listed in the review section.
limitations: limitations: 1
review_summary: review_summary: A nice work on learning linear latent representation for visual RL; however, its applicability to real-world hardware control may be limited.
title: title: Robotic task verifications are required.
paper_summary: paper_summary: Robot control based on image observation is a challenging problem regarding learning efficiency and robustness. This study proposes a Contrastive Spectral Koopman Embedding network that achieves high computational efficiency and agility to disturbances by explicitly representing the linear state space introduced by Koopman Theory to a diagonal canonical form.  This paper also proves the convergence of the policy learning in the LQR problem on linearized latent spaces based on Koopman operators. It demonstrates its computational efficiency and high performance through standard image-based RL benchmarks.
main_review: main_review: The paper's description is clear, and the proposed method's originality is compared with other methods. The paper is of sufficient quality, as the proposed method's effectiveness is verified theoretically and experimentally. On the other hand, the paper's significance could be more questionable, as there are no benchmarks specific to the robotics field, and it is unclear whether the assumptions in the proof of the theorem are held in a real environment.

Strengths
1. A Contrastive Spectral Koopman Embedding network is proposed, which achieves high computational efficiency and agility to disturbances by explicitly representing the linear system introduced by Koopman Theory to a diagonal canonical form.
1. The convergence of the proposed method's policy in the linearized LQR problem on Koopman operators is proved.
1. Computational efficiency and high performance were demonstrated through standard image-based RL benchmarks.

Weakness
1. It has been evaluated only on standard RL benchmarks in the physics simulator and is unsuitable for robotics verification.
The state-space representation is restricted to a diagonal canonical system, which improves learning and computational efficiency. However, based on current experimental results, judging whether this will work on a real robot is difficult.

Thank you for the author's response.
The authors answered my equation and thoughtfully discussed my concerns. However, my recommendation remains unchanged. The proposed idea of explicit diagonal structure in state space is interesting, and the benchmark score is promising for high-dimensional observation. This research progress is expected by conducting further robotics experiments.
questions: questions: 1. Can the Elapsed time be evaluated about Figure 3? In real robots, the overall computation time, including overhead, needs to be evaluated rather than the number of computational operations.
1. Can the robustness of the physics parameter to disturbances be evaluated about Figure 4?
In robotics applications, the robustness of Sim2Real against environmental changes is also subject to evaluation.
1. Does the assumption hold in a real system, the continuous Objective functions assumed for the proof of Theorem 4?
According to the motivation for which Fazel et al. (2018) derived Lemma 6, it is non-trivial to derive a proof of convergence based on the gradient method for even simple LQR by continuity of the Objective function since the Feedback System (A-BK) can be unstable. Does the proposed method address this point? And did such instability occur in the data of the actual learning process?

Fazel, M., Ge, R., Kakade, S. &; Mesbahi, M.. (2018). Global Convergence of Policy Gradient Methods for the Linear Quadratic Regulator. Proceedings of the 35th International Conference on Machine Learning
limitations: limitations: 1
review_summary: review_summary: Please see above