Zero-Shot Transfer of Neural ODEs

Authors: tyler  ingebrand
adam j thorpe
ufuk  topcu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 22:29:42

Abstract: Autonomous systems often encounter environments and scenarios beyond the scope of their training data, which underscores a critical challenge: the need to generalize and adapt to unseen scenarios in real time. This challenge necessitates new mathematical and algorithmic tools that enable adaptation and zero-shot transfer. To this end, we leverage the theory of function encoders, which enables zero-shot transfer by combining the flexibility of neural networks with the mathematical principles of Hilbert spaces. Using this theory, we first present a method for learning a space of dynamics spanned by a set of neural ODE basis functions. After training, the proposed approach can rapidly identify dynamics in the learned space using an efficient inner product calculation. Critically, this calculation requires no gradient calculations or retraining during the online phase. This method enables zero-shot transfer for autonomous systems at runtime and opens the door for a new class of adaptable control algorithms. We demonstrate state-of-the-art system modeling accuracy for two MuJoCo robot environments and show that the learned models can be used for more efficient MPC control of a quadrotor.

Main paper: 
 introduction: 
 Models that are adaptable, generalizable, and capable of learning online from minimal data are essential for autonomy. These models must adapt to unseen tasks and environments at runtime without relying upon a priori parameterizations. For example, consider an autonomous UAV delivery robot navigating in a dense, urban environment through varying wind patterns and carrying uncertain payloads. This scenario requires rapid adaptation to ensure safe and correct operation because conditions can change unpredictably and online model updates are impractical. While prior works can control autonomous systems in a single setting, they fail to adapt to the continuum of real-life scenarios. The key challenge is enabling zero-shot transfer of learned models, where models quickly adapt to new data provided at runtime without retraining.We present a method for modeling differential equations by learning a set of basis functions parameterized by neural ODEs. Our key insight is to learn a space of functions that captures feasible behaviors of the system. By focusing on learning the structure of the space of differential equations, our approach implicitly learns how the dynamics change due to changes in the environment. Our approach is based on the theory of function encoders [14] , a framework for zero-shot transfer that has been applied to task transfer in reinforcement learning contexts.We formulate the space of learned functions as a linear space equipped with an inner product (e.g. a Hilbert space), and learn a set of basis functions over this space, where each basis function is represented by a neural ODE. This structure offers an efficient way to approximate online dynamics via a linear combination of the basis functions.By representing functions in a Hilbert space and pre-training on a suite of functions, we can quickly identify the basis functions' coefficients for a new dynamical system at runtime using minimal data. This is useful, for instance, in scenarios where we can pre-train offline in simulation, but need to 38th Conference on Neural Information Processing Systems (NeurIPS 2024). quickly identify dynamics online from a single trajectory in a zero-shot manner. This strategy greatly reduces the computational overhead associated with adapting or re-training a neural ODE to new tasks at runtime. The ability to efficiently encode the behavior of a system at runtime without retraining is a key component of our approach. Our approach is outlined in Figure 1 .Our approach yields models which generalize to a large set of possible system behaviors and achieves better long-term prediction accuracy than neural ODEs alone. We showcase our approach on a problem of predicting the behavior of a first-order ODE system with Van der Pol dynamics. We demonstrate the scalability of our approach on two MuJoCo robotics environments [28] . Finally, we test the feasibility of using the learned model for downstream tasks such as model-predictive control (MPC). Our results show that our model achieves significantly better long-horizon prediction accuracy compared to the nearest baseline. Additionally, the MPC controller using our model has a lower slew rate, indicating that the improved model accuracy leads to more efficient control decisions.Representing Spaces of Dynamical Systems: We propose a novel framework for representing spaces of dynamical systems, i.e. induced by hidden system parameters, variations in the underlying physics, or changing environmental features. Using a large-scale set of data collected offline, we learn a collection of neural networks that act as a functional basis. This approach is based in the theory of function encoders [14] . Yet the extension to using neural ODEs as basis functions is non-trivial and poses several challenges, mainly from the need to integrate the learned model.A Method for Online Adaptation: We construct a method for adapting neural ODE estimations based on online data without gradient updates, i.e. zero-shot transfer of system models. Our approach overcomes a significant challenge in learning behaviors of differential equations and dynamical systems. By offloading the computational effort to the training phase, we enable rapid online identification, adaptation, and prediction without retraining.Empirical Results: We demonstrate accurate long-horizon predictions in challenging robotics tasks and show these models can be used for online control of quadrotor systems. We assess the quality of the approach in three areas and answer the following questions: 1) How well does our approach adapt to new dynamics online? 2) How does our approach compare to existing approaches for long-horizon prediction tasks? and 3) Does our approach work for downstream tasks such as control? We show that function encoders using neural ODEs as basis functions consistently outperform existing approaches. 
neural odes: 
 Consider an ordinary differential equation ẋ(t) = f (x(t), t), where f is Lipschitz continuous and x(t) ∈ R n is the state at time t. Given an initial condition x(t 0 ), the ODE solution can be written asEQUATIONNote that in general, the explicit dependence of f x(t), t on t can be removed by augmenting the state x to include t. As such, we omit t throughout.Neural ODEs [5] parameterize the function f as a neural network. In particular, neural ODEs solve the ODE using an off-the-shelf integrator and optimize the neural network with respect to a prediction loss. The training procedure requires a dataset D = {(t i , x(t i ))} d i=1 which is used to train the model via a supervised objective, such as mean squared error, back-propagated through the integrator.Neural ODEs have demonstrated impressive accuracy for long-horizon predictions of continuous-time systems. Furthermore, they can be trained on trajectories with irregular time intervals between samples [18] , and generalize better than multi-layer perceptron models. However, they lack adaptability and need to be retrained for every scenario.To achieve zero-shot transfer, we employ the theory of function encoders [14] . While typical machine learning approaches learn a single function, function encoders learn basis functions to span a space of functions. This allows the function encoder to achieve zero-shot transfer within this space by identifying the coefficients of the basis functions for any function in the space at runtime. Once the coefficients have been identified, the function can be reproduced as a linear combination of basis functions. Despite the broad applicability of function encoders, the extension to capture solutions to differential equations via neural ODEs is non-trivial.Formally, consider a function space F = {f | f : X → R m } where X ⊂ R n . Instead of learning a single function, function encoders learn k basis functions g 1 , g 2 , . . . , g k that are parameterized by neural networks in order to span F [14] . Define the inner product of F as ⟨f, g⟩ F = ⟨f (x), g(x)⟩dx. Then, the functions f ∈ F can be represented as a linear combination of basis functions,EQUATIONwhere c ∈ R k are real coefficients and θ i are the network parameters for g i . Let V be the volume of X . For any function f ∈ F, an empirical estimate of the coefficients c can be calculated using data {(x j , f (x j ))} m j=1 via a Monte-Carlo estimate of the inner product,EQUATIONThe basis functions are trained using a set of datasets, D = {D 1 , D 2 , ...}, where each dataset D i = {(x j , f i (x j )} m j=1 consists of input-output pairs corresponding to a single function f i ∈ F. For each function f i and corresponding dataset D i , first compute the coefficients {c 1 , c 2 , ..., c k } via (3) and then obtain an empirical estimate of f i via (2) . Then compute the error of the estimate of f i though the norm induced by the inner product of F. The loss function is simply the sum of the losses for all f i , which is minimized via gradient descent. For more details, see [14] .After training, the basis functions are fixed, and the coefficients c of a new function f ∈ span{g 1 , . . . , g k } are computed via (3) or via least-squares, using data collected online. This is key for efficient, online calculations, since the approximation in ( 3) is effectively a sample mean and requires no gradient calculations.Orthogonality of the Basis Functions: Note that function encoders do not enforce orthogonality of the basis functions explicitly [14] . Using Gram-Schmidt to orthonormalize the basis functions during training can significantly increase the training time and is computationally intensive. Instead, during training the coefficients are computed using (3) presuming that the basis functions are orthogonal. This is key. This causes the basis vectors to naturally become more orthogonal as training progresses since the loss implicitly penalizes the basis functions if they are not orthonormal. See Appendix I.Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: We discuss the key limitations of our approach in §5, including: computational overhead, dependence on training data, and the constraints on integration. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
function encoders with neural odes as basis functions: 
 Consider a space F of Lipschitz continuous dynamical systems f : X → X . The space of dynamical systems can arise, for instance, due to uncertain parameters, minor variations in a first-order physics model, or changing environmental features. Given an initial condition x(t 0 ), our goal is to estimate the state x(t f ) at a future time t f > t 0 . The integral form of the initial value problem is given by (1).Our approach can be separated into two distinct phases: offline training and zero-shot prediction. During offline training, we presume that we have access to an offline dataset D = {D 1 , D 2 , . . .}, where D i is a realization of a trajectory from a function f i ∈ F. During zero-shot prediction, we seek to predict a previously unseen function f and have access to a minimal trajectory D taken from f . We seek to learn a set of basis functions g 1 , . . . , g k that span F, where k is a user-specified hyperparameter. By learning a set of basis functions that span the space F, we obtain a means to represent the behavior of any dynamical system in the space. However, we do not observe measurements of f directly since we cannot typically measure the instantaneous derivative ẋ of a dynamical system. Instead, we will equivalently learn a set of neural ODE basis functions such that the underlying neural networks which are being integrated correspond to g 1 , ..., g k . We then seek to compute a representation of a new function using data collected online.For every f ∈ F, we define the integral term in (1) as a function H : X × T → X , given by,EQUATIONWe model the dynamical system f using a function encoder as in (2) . Using (2) in ( 4), and by the linearity of the definite integral, we have that,EQUATIONwhere g i is neural network parameterized by θ i andG i (x(t 0 ), t f ) := t f t0 g i (x(τ ) | θ i )dτ.One interpretation of the above equation is that we can represent H as a weighted combination of basis functions G i , and the problem of learning a set of basis functions g 1 , . . . , g k can equivalently be viewed as learning a set of neural ODEs G 1 , . . . , G k . Thus, we define the Hilbert space H of functions H as in (4) and equip it with the following inner product,EQUATIONWe then learn basis functions G 1 , . . . , G k spanning H where each basis function is a neural ODE.From (6) , the coefficients of a function H ∈ H are given by c i = ⟨H, G i ⟩ H . However, from [14] , computing the inner product exactly is generally intractable in high-dimensional spaces. We can empirically estimate the coefficients c i using a trajectory of (potentially irregularly) sampled states {x(t) | t = t 0 , . . . , t m } from a dynamics function f ∈ F. Using the trajectory, we form the dataset D = {(x(t j ), x(t j+1 )} m-1 j=0 . We can compute c using D via a Monte-Carlo estimate of the inner product in (6) ,EQUATION)where V is the volume of the region of integration, and following from (1),EQUATIONIn other words, we substitute the difference between states x(t j+1 ) -x(t j ) for H(x(t j ), t j+1 ) in (7) .Let D = {D 1 , D 2 , ...} be a set of datasets, where each D ℓ = {(x(t j ), x(t j+1 ))} m-1 j=0 is collected from a trajectory from a function f ℓ ∈ F. For each dataset D ℓ , we compute the coefficients c 1 , . . . , c k according to (7) . The coefficients can be used to approximate the corresponding H ℓ via (5) . We then evaluate the error of H ℓ using the dataset D ℓ and minimize its loss via gradient descent. This is done c i ≈ V m m-1 j=0 ⟨x(t j+1 ) -x(t j ), G i (x(t j ), t j+1 -t j )⟩ X 9:end for 10:L = L + m-1 j=0 ∥(x(t j+1 ) -x(t j )) - k i=1 c i G i (x(t j ), t j+1 -t j )∥ 2 11:end for 12: θ = θ -α∇ θ L 13: end while for multiple functions f ∈ F at each gradient update to ensure the basis learns the space rather than a single function. We present this as Algorithm 1.Applying Algorithm 1 yields basis functions which span the space of dynamical systems, where each basis function is a neural ODE. This space describes possible behaviors of the system, where variations in environmental parameters, physics, etc. correspond to a particular dynamics function within this space. Therefore, this algorithm represents complicated system behaviors simply as a vector within a Hilbert space. Section 3.2 shows how to use these basis functions for zero-shot dynamics prediction from small amounts of online data.After training, we fix the parameters of the basis functions g 1 , ..., g k , and can compute the coefficient representation c ∈ R k for any function f ∈ span{g 1 , . . . , g k } via (7) . If D is rich enough to capture the various behaviors of the systems in F, then we can estimate the behavior of any dynamics f ∈ F.Given data collected online from a single trajectory, we can compute the coefficients using the Monte-Carlo estimate of the inner product as in (7) . This approximation is a crucial component of the approach. It allows the inner product to be computed from data through an operation that is effectively a sample mean. Therefore, this approach can be computed online quickly even for large amounts of data. Then, given the coefficients, the future states of the system can be predicted using (5) . These properties allow the neural ODE to achieve zero-shot transfer. Identifying the coefficients only requires inner product calculations, vector addition, and scalar multiplication, and so it can be computed online without any gradient updates.The Residuals Method: The zero vector of the coefficients space corresponds to the zero function. Since the feasible dynamics are differentiated by their coefficients, it is numerically convenient if the coefficients corresponding to all feasible systems are centered around zero. Thus, we can re-center the space of coefficients around the center of the cluster of feasible dynamics. This is done by first modeling the average dynamics F avg in the set of datasets D, and then learning the residuals between each function and F avg . In other words, the basis functions are trained to span the function space corresponding to R(x(t 0 ),t f ) = x(t f ) -x(t 0 ) -F avg (x(t 0 ), t f ).This method can achieve better accuracy, but requires learning one additional neural ODE, F avg . Alternatively, an approximate dynamics model based on prior knowledge can be used as F avg . We describe the training procedure for this approach in Algorithm 2.We can account for a zero-order hold (ZOH) control input u ∈ U ⊂ R p with minimal modifications. A ZOH control input is given by a piecewise constant function, meaning it is held constant over the period of integration. Given controlled dynamics, f : X × U → X , we modify the corresponding functionsH to incorporate a constant input, H x(t 0 ), u, t f = k i=1 c i t f t0 g i x(t 0 ), u | θ i dτ .Then, using trajectory data that also includes the controls applied at each time interval, we can estimate =2.17 the coefficients using datasets 5) and (7) . The remainder of the training procedure is unchanged.D = {(x(t j ), u j , x(t j+1 ))} m-1 j=0 , substituting g i x(τ ), u, | θ i in (Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Justification: Our work has positive implications for the design and modeling of autonomous systems. We briefly discuss the broader impacts of our work in §8 Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
numerical experiments: 
 We demonstrate the effectiveness of our approach for predicting and controlling dynamical systems through several numerical experiments. We first demonstrate that the approach can adapt to different dynamics using a Van Der Pol oscillator system. We then show long-horizon prediction accuracy on challenging MuJoCo robotics experiments and compare to neural ODEs (NODE) [5] and function encoders as in [14] using the residuals method (FE + Res). Lastly, we show the learned models are sufficiently accurate for downstream tasks on a difficult control task using a quadrotor system. The source code is available at https://github.com/tyler-ingebrand/NeuralODEFunctionEncoder.Current off-the-shelf integrators with adaptive step sizes do not support efficient batch calculations. Because this algorithm involves training numerous neural ODEs on a large amount of data, the ability to train on data in batches is required. Therefore, we implement an RK4 integrator since it can be efficiently computed for multiple data points in parallel. The Van der Pol visualization uses 11 basis functions while the MuJoCo and Drone experiments use 100. For ablations on how the hyper-parameters affect results, see Appendix G.We first demonstrate that our approach can adapt to a space of dynamics that vary according to a nonlinear parameter. The Van der Pol dynamics are defined as, ẋ = y, ẏ = µ(1 -x 2 )y -x, where [x, y] ⊤ ∈ R 2 is the state, and µ is a hidden parameter. We collect multiple datasets D ℓ where µ is fixed for the duration of the trajectory, but varies between trajectories. We train basis functions using Algorithm 1. We then compute the coefficients via (7) and approximate the dynamics via (5) .We plot the results in Figure 2 . As expected, we observe that our proposed approach can predict the dynamics of a space of Van der Pol systems without retraining. We can also see that a single neural ODE trained on the same data can only fit a single function. Therefore, its prediction corresponds most closely with the behavior of a single Van der Pol system that has the mean µ value. This illustrates that our approach is capable of adapting to different dynamics at runtime. We evaluate the performance of our proposed approach on the Half-Cheetah and Ant environments [28] , shown in Figure 3 . The hidden environmental parameters are the length of the limbs, the friction coefficient, and the control authority. Of the two environments, Ant is more difficult due to its higher degrees of freedom. For training, we collect a dataset of trajectories where the hidden parameters are unobserved, but held constant throughout the duration of a given trajectory. After training, we use 200 datapoints, equivalent to about seven seconds of data for a system running at 30 Hertz, and use only this data to identify the dynamics. Note this online phase is computationally simple, and can be done in only milliseconds on a GPU.Neural ODEs (NODE) perform poorly because they have no mechanism to condition the prediction on the hidden-parameters. Effectively, NODE learns the mean dynamics over all dynamics functions in the training set. Function encoders using the residuals method (FE + Res) can implicitly condition their predictions on the hidden parameters through the coefficient calculation, though they are unable to achieve accurate long horizon predictions on the more challenging Ant problem. This is because it lacks the inductive bias of neural ODEs. Our approach (FE + NODE) can both implicitly condition the predictions on the hidden parameters through data, but also benefits from the inductive bias of neural ODEs. We see that the residuals method performs best out of all approaches in both environments. This is because the average model significantly reduces the epistemic uncertainty and provides a meaningful baseline from which to center the training. The average model acts as a good inductive bias and makes it easier to distinguish between the learned functions during training. We additionally compare against an oracle prediction approach (Oracle), which has access to the hidden parameters as an additional input with a neural ODE as the underlying architecture. While its 1-step prediction accuracy demonstrates good empirical performance, the long-horizon predictions are unstable. This is because Oracle is required to generalize to an entire space of dynamics with one NODE, which is a complex and difficult function to learn.Lastly, we seek to test the accuracy of our approach for use on downstream tasks such as control on a realistic example using a robotic system. We seek to determine if the learned models are sufficiently accurate for model-based control in the presence of hidden parameters. We use a simulated quadrotor system using PyBullet [31] , which is a highly nonlinear control system. We use the quadrotor's mass as a hidden parameter. The goal is to predict the future state of the quadrotor system under any hidden parameters, using 2000 data points collected online to identify the dynamics. Then, given the learned model of system behavior, we seek to control the quadrotor using gradient-based model predictive control (MPC) to reach a pre-specified hover point. The results are plotted in Figure 4 .The results show that FE + NODE + Res outperforms competing approaches at long-horizon predictions. Furthermore, we plot the 10-step MSE as a function of mass in Figure 4 , and we observe that FE + NODE + Res accurately predicts the system behavior across varying masses. We observe a slight decay in performance for low masses which are more sensitive to control inputs during simulation, which causes the simulated trajectories to diverge more from the rest of the observed data.The neural ODE (NODE) performs poorly for different masses, and its performance decays quickly as the dynamics deviate from the mean behavior.Lastly, we see that this prediction accuracy translates to the downstream performance of an MPC controller. While all approaches are sufficiently accurate for control due to the fact that MPC is partially robust to model inaccuracies, the prediction accuracy has a corresponding impact on the task performance. Neural ODEs (NODE) demonstrate a high slew rate, which reflects the need for repeated positional corrections necessitated by taking bad actions. In contrast, FE + NODE and FE + NODE + Res have lower slew rates as they make more accurate decisions. We plot two example trajectories that demonstrate this behavior in Figure 5 . 
scope & limitations: 
 Overhead: Our approach incurs a cost of either increased inference time or memory, depending on if the basis functions are integrated sequentially or in parallel. We integrate them sequentially during training to reduce memory overhead and allow for larger batch sizes, while integrating in parallel at execution to prioritize inference speed. See Appendix D.Data dependency: In order to make efficient, online approximations of a new function f ∈ F without gradient calculations, the basis functions must be trained to span the space of possible dynamics. To do so, there must be sufficient example datasets of possible dynamics under fixed hidden parameters in D. This implies a larger amount of data must be collected to learn a space of dynamics then would be needed to learn a single dynamics function.The training procedure trains the basis functions on short time intervals t f -t 0 , in which x(t) ∈ X . The basis functions have only been trained for inputs in the space of X , where their behavior outside of X is unpredictable. As a result, it is necessary to either integrate each basis function for a short time interval before calculating the state according to (5) , or to integrate the basis functions as described in B. Integrating the basis functions over long horizons without calculating the state of the system during intermediate steps may lead the predicted state to leave X , at which point the behavior of that basis function becomes unpredictable. 
related work: 
 Basis Functions: Function approximation techniques often employ a linear model over a predefined set of basis functions. Techniques such as Taylor series, Fourier series, and orthogonal polynomial systems utilize an infinite set of basis functions, theoretically allowing perfect function representation [6, 23, 7] . However, in high-dimensional spaces, these techniques become impractical due to the exponential growth in the number of basis functions. Additionally, many approaches depend on the choice of a feature map or kernel to define the function space [29, 25] , which imposes structure by selecting the class of functions to learn from (e.g. using radial basis functions [3] ). These design choices necessarily introduce approximation errors through the choice of function class, or may depend on prior domain knowledge, which may not always be available. Methods for identifying system dynamics that use a large library of pre-defined basis functions, such as Koopman operators [2] or nonlinear system identification through sparse regression such as SINDy [4, 26, 15] , have received considerable attention. Yet these approaches typically employ a finely-crafted finite dictionary of basis functions, which requires careful choice to achieve good data-driven performance [20] . Neural network approaches such as functional-link and orthogonal networks [30, 7] omit hidden layers and use gradient descent to learn a linear combination of features, encoding the function class into the network architecture, but fail to generalize well, and are not amenable to zero-shot transfer. In contrast, we compute the coefficients of the model through a well-defined inner product, which scales well with data and can be computed quickly. Furthermore, our basis functions are entirely learned from data during the training phase, similar to representation learning [1] , and thus require no prior assumptions or domain knowledge.Neural ODEs: In existing work, neural ODEs have proven to be a powerful tool for modeling dynamical systems [10, 24, 22, 11, 17] and stochastic differential equations [21, 13] , but generally require an extensive data collection and training phase. While the model training can be enhanced [8] and the models can incorporate prior knowledge [9, 10] to reduce the training time, they inherently focus on a single system at a time. This inherently limits their ability to generalize across different systems without retraining. Notably, parameterized neural ODEs [19] pass the model parameters as an additional input to the neural ODE. This approach has been shown to achieve a form of transfer within the set of allowable parameters, but requires extensive knowledge of the system parameters and the structure of the dynamics, both at training and test time. In principle, our approach can be combined with these existing approaches to incorporate their distinct advantages, making our approach highly generalizable to different systems and modeling frameworks.Deep Learning Techniques: Few-shot meta learning aims to solve a similar problem, where a learned model is adapted given an online dataset [12] . However, meta learning requires gradient updates to the learned models, which may be too slow for real-time control. Transformers are another technique that can adapt given an online dataset by feeding that dataset as input to the encoder side of the transformer. However, transformers have long forward pass times and scale quadratically with the amount of data [16] , and so they are not amenable for model-based control. Domain randomization in reinforcement learning is another technique to generate a policy which is robust to a large set of dynamics [27] . In contrast, our dynamics model adapts to the current dynamics, and thus our controller is adaptive, rather than robust. 
conclusion & future work: 
 We introduced zero-shot neural ODEs, which accomplish both long-horizon predictions and zero-shot transfer. We demonstrated the performance of this approach on two challenging MuJoCo tasks and on the control of a quadrotor system. Our approach makes a significant step towards online adaptability of model-based control and has implications for the safe control of autonomous systems in the presence of uncertainty. In future work, we plan to address safety during training, perhaps using the properties of the Hilbert space to characterize the epistemic uncertainty. We also plan to explore theoretical extensions to stochastic differential equations and Hilbert spaces of probability measures. 
broader impact: 
 This approach demonstrates several clear benefits for enabling same-day adaptation, which is a critical need for autonomous systems that will be deployed in new, unstructured environments. Nevertheless, this approach will require a more thorough theoretical analysis before it can be deployed on actual robotics systems, e.g. to determine confidence or sample bounds to guarantee safety. Notably, this work is a step toward bridging the sim-to-real gap, though it remains unclear how well real-world systems will be represented by a set of basis functions learned in simulation. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We do not provide a model or dataset with a high risk for misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [NA] Justification: Our experiments do not rely upon any restricted licenses or code, and we provide citations to key libraries used.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: We do not release new assets.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: Our paper does not involve crowdsourced data or research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: Our paper does not involve crowdsourced data or research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: All experiments use an Intel 9th Generation i9 CPU and a Nvidia 2060 GPU with 6GB of memory.When computing the approximate inner product between the true system and each basis function, it is necessary to compute G i (x(t 0 ), t f ) for every tuple in the dataset D. However, once the coefficients c have been computed, it is, in theory, no longer necessary to integrate the basis functions separately. From (5), we can approximate H asEQUATIONby the linearity of the integral. In effect, we are summing the gradients of each basis function rather than the basis functions themselves. As a result, inference for a specific set of coefficients can be decreased from requiring k integrations to only a single integration. However, we find this method to make less accurate predictions in practice, which may be due to our choice of integrator. This trick may be better suited to variable step-size integrators, which benefit more from reduced calls to the integrator than RK4 does.We leverage RK4 as the default integrator for this work, as it can run a forward pass in milliseconds.There are more accurate integrators available, such as odeint. However, there is inherently a trade off with respect to both training time and execution time. Integrators such as adaptive step size solvers can potentially make 20 or more calls to the neural ODE during a forward pass, while RK4 makes only 4. The increased number of neural ODE forward passes greatly increases memory usage and compute time. We experimented with more accurate integrators, but ultimately found this tradeoff to be unfavorable. Future work should investigate integrators that are fast, but achieve better accuracy than RK4.While the function encoder algorithm alone has minimal overhead relative to a MLP, this is not the case for neural ODEs due to the need for integration. The k neural ODEs may either be integrated sequentially or in parallel. If they are integrated sequentially, the memory overhead is lower, especially with respect to back-propagation. Thus, we find this useful for offline training, where the sequential method allows us to compute gradients for a larger batch of data at the cost of training time. In contrast, online execution generally favors inference speed over memory overhead. Thus, we use the parallel method for online inference in the drone example, which requires much more memory but only a small overhead of inference time relative to neural ODEs alone.Training a function encoder with the residuals method requires two separate loss functions. The first loss function trains F avg on data from all datasets D ℓ , which means it effectively learns the expectation of F given the training set. This loss function can be skipped if F avg is a fixed function based on prior knowledge.The second loss function trains the basis functions. Unlike in Algorithm 1, the function being learned is x(t j+1 ) -x(t j ) -F avg (x(t j ), t j+1 -t j ). In other words, the residual between the data and the average function. This loss is only used to train the basis functions, it is not used to train the average function. See Algorithm 2.The Residuals Method for all D ℓ ∈ D do 8:L 1 = L 1 + m-1 j=1 ∥(x(t j+1 ) -x(t j )) -F avg (x(t j ), t j+1 -t j )∥ 2 9:end for for i ∈ 1, ..., k do 15:c i ≈ V m-1 m-1 j=1 ⟨x(t j+1 ) -x(t j ) -F avg (x(t j ), t j+1 -t j ), G i (x(t j ), t j+1 -t j )⟩ X 16:end for 17:L 2 = L 2 + m-1 j=1 ∥ (x(t j+1 ) -x(t j ) -F avg (x(t j ), t j+1 -t j )) - k i=1 c i G i (x(t j ), t j+1 -t j )∥ 2 18:end for19: θ = θ -α∇ θ L 2 20: end whileAll baselines use the same training scheme. We use an ADAM optimizer with a learning rate of 1e -3, and gradient clipping with a max norm of 1. NODE baselines uses 4 hidden layers of size 512, while FE + NODE baselines uses 4 hidden layers of size 51 for each basis function. Note this leads to approximately the same number of parameters for both approaches because the number of hidden parameters scales quadratically with the size of the hidden layers. All baselines train on 50 functions per gradient update via gradient accumulation. States are normalized to have 0 mean and unit variance.A random policy is used to collect data for the MuJoCo environments. A PID-based exploratory policy, which moves to random nearby points, is used to collect data for the quadrotor since a random policy collides with the floor. Evaluations are done on a holdout set collected through the same means.All quadrotor baselines use the same MPC controller. The controller optimizes the actions through a combined sampling, gradient descent over 100 iterations. The episode is 100 steps, while the planning horizon is 10 steps. The controller optimizes 100 sample trajectories in parallel, and ultimately chooses the best one. Warm starting is used for following MPC calls to improve performance. The cost function penalizes distance to the objective point, deviance from a stable horizontal position, velocity, and the difference between torques on each rotor. The results indicate the the proposed approach is insensitive to increasing example dataset sizes, which suggests that 200 data points is sufficient for the coefficients to converge. Half CheetahConsider a set of basis functions g 1 , . . . , g k . Suppose that g 1 , . . . , g k is not orthonormal. Now consider a function f , and suppose f happens to be in the span of g 1 , . . . , g k . Then f can be expressed as f = b ⊤ g, where b is a set of coefficients and g is the concatenation of g 1 , . . . , g k . The coefficients are calculated via the inner product,c ⊤ =    ⟨f, g 1 ⟩ . . . ⟨f, g k ⟩    =    ⟨b ⊤ g, g 1 ⟩ . . . ⟨b ⊤ g, g k ⟩    = b ⊤    ⟨g 1 , g 1 ⟩ . . . ⟨g 1 , g k ⟩ . . . . . . . . . ⟨g k , g 1 ⟩ . . . ⟨g k , g k ⟩   The loss functionL = |f -f | 2 = |f -c ⊤ g| 2 .If c = b, then the loss will be 0. Observe that c = b if and only if the Gram matrix is identity, and the Gram matrix is identity only for an orthonormal basis.In other words, the minimizer of the loss function is an orthonormal basis. Thus, in order for gradient descent to decrease loss, the basis functions converge towards orthonormality [14] . This intuition is empirically validated in [14] , Appendix A.5.As a final note, the coefficients can be computed via least squares after training. Least squares does not require an orthonormal basis as it uses the Gram matrix to account for the inner products between basis functions.1. Claims Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: We claim to learn a space of dynamics spanned by a set of neural ODEs. These claims are justified in the methods §3, where we show theoretical results, and empirically validated in §4. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: This paper presents a neural ODEs-based method for modeling differential equations by learning basis functions, focusing on capturing system dynamics of autonomous systems within a structured function space. 

The paper was reviewed by four reviewers who all are inclined to suggest for acceptance. The most reviewers agree that (1) the proposed method, the use of neural ODEs as function encoders, is intuitive and sound, (2) dealing with important problems, zero-shot generalization to dynamic environments, in machine learning and autonomous systems. The reviewers also commented that (3) the paper is well written with good readability and presentation. 

Considering these strengths, this paper is recommended for acceptance. The authors are encouraged to utilize their constructive and detailed discussion with reviewers towards further raising the quality and impact of their work.

Reviews: review: paper_summary: paper_summary: The paper aims to address the challenge of zero-shot transfer and adaptation. The authors propose tackling this challenge by learning a dynamics space spanned by neural ODE basis functions, which can then be used for rapid identification and adaptation to dynamics at inference time without additional training. The paper demonstrates the efficacy of the proposed approach, using both a simpler oscillator system and scaled up simulated robotics environments.
strength_weakness: strength_weakness: - To place the results in context with prior work, it would be additionally helpful for the robotics experiments to show comparisons to other methods that enable adaptation (e.g. training a model free method with domain randomized parameters). 
- It’s not clear how well the method will scale to more complex dynamics, or when trying to generalize beyond the learned basis functions (e.g., if they are not expressive enough or span the relevant spaces for a dynamical system, or if insufficient data is used to learn the dynamics space).
questions: questions: See suggestions in the Weaknesses section above.
limitations: limitations: The authors address method limitations adequately in the paper.
paper_summary: paper_summary: This paper proposes a method to learn the dynamics of autonomous systems in a few shot manner. The core assumption is that the dynamics function dx/dt=f(x) of a new system can be modeled by a linear combination of basis dynamics functions. The method involves two stages. In the offline stage, the method learns a set of basis functions with neural ODE. In the online stage, the method uses incoming dynamics to fit the linear weights, which when combined with the learned basis, can be used to model the new system. The authors evaluated the effectiveness of the method on  Van der Pol Oscillato system (with varying parameters) and Mujoco Ant.
strength_weakness: strength_weakness: 1. My main criticism of the paper is the technical contribution. The problems this system can solve seem to be constrained to systems limited variation in parameter, where the offline dataset & online system share a high level of similarity. While the authors explained how they apply NODE very clearly, I think using NODE is not a fundamental contribution because you can also have other sequence models trained to predict residue. The idea of basis function has also been explored by many prior literatures. 

2. The title is misleading - you are doing quick online adaptation (few shot) instead of a zero-shot setting. A few trajectories of online data is required to adapt to the new system

3. There are a lot of baselines the authors should look into since few shot learning / quick adaptation is a long standing topic of research. For example, other basis function methods, meta learning methods. I understand that the authors picked a setting where time is continuous, but the mujoco Ant environment is also traditionally studied as discrete too, so a meta learning + some sequence prediction model should be applicable here. If you can provide a convincing argument about why those priors works aren't applicable it's also fine.

4. As I suggested in questions section, there could be many perspectives the author should dive deeper to improve the paper


minor:
1. In line 114, the reasoning doesn't provide grounding to the orthogonality question - the training objective make them span the space assuming diverse enough dynamic systems - but this is unrelated to orthogonality! Are you saying that non orthogonal is okay as soon as they span the space?

2. It would be good to give more intuitive visualizations of estimation error in mujoco Ant. e.g. render the predicted trajectories as video

3. ablation about the number basis functions used would be helpful to understanding
questions: questions: 1. In line 165, the authors said "Given data collected online from a single trajectory". I am curious to see ablations how the method improves as the amount of online data grow

2. In figure 3, error goes up as the number of look-ahead steps increases. While this is partially expected due to compounding error, how much of it shall be attributed to limited expressiveness of linear basis. e.g. I can do an experiment where I train a NODE with many many data, not just 200 example points, and roll it out. You will witness how MSE loss increases and gain insights about a upper bound for prediction accuracy that's not due to limited data / limited number of basis
limitations: limitations: Please see my points in weakness
paper_summary: paper_summary: The paper presents a novel framework for the zero-shot transfer of neural ODEs by leveraging function encoders to represent a space of dynamical systems. It demonstrates the method's effectiveness in adapting to unseen environments without retraining, using MuJoCo and quadrotor experiments.
strength_weakness: strength_weakness: The paper presents a promising framework for zero-shot transfer of neural ODEs, but there are several areas for improvement. Firstly, the reliance on a large and diverse dataset for training is a significant limitation. The approach requires extensive data that spans the entire function space of possible dynamics, which may not always be feasible. This dependency on comprehensive datasets should be addressed by exploring data-efficient learning methods or leveraging transfer learning techniques.

Secondly, the computational overhead involved in training multiple neural ODEs is substantial. This might hinder the scalability and real-time applicability of the proposed method. The authors could investigate more efficient training algorithms or consider approximations that reduce computational costs without compromising accuracy.

Thirdly, the paper does not explicitly enforce the orthogonality of basis functions, relying instead on implicit regularization through the loss function. This might lead to suboptimal representation of the function space, affecting the model's performance. 

Moreover, the experiments lack diversity in evaluating the model's robustness to completely unseen environments or significantly different conditions. Including more varied test scenarios or stress tests could provide a deeper understanding of the model's adaptability and limitations.
questions: questions: 1. The computational overhead for training multiple neural ODEs might be significant. How does this affect the scalability and real-time applicability of your method?

2. The paper does not enforce orthogonality explicitly for the basis functions, which might lead to suboptimal representation. Have you evaluated the levels of orthogonality (and their relationship with performance) in your method?

3. The experiments lack diversity in evaluating the model's robustness to completely unseen environments or significantly different conditions. How confident are you in your model’s adaptability in such scenarios?

4. The evaluation metrics used in the MuJoCo experiments are primarily focused on prediction accuracy. Have you considered additional metrics that might better capture your model's practical performance?  For example,  task-specific performance such as scores or cumulative rewards.

5. It seems to lack a sufficient variety of baselines for comparison. Have you considered additional zero-shot / few-shot or multitask / meta-learning studies?
limitations: limitations: Refer to Weaknesses and Questions.
paper_summary: paper_summary: The work explores the use of neural ODEs as basis functions for function encoders. This requires dealing with the additional integration step, with the weighted combination of obtained ODEs representing the behaviour of the function to approximate via its integral. With an inner product and its tractable Monte Carlo estimation scheme, the authors derive an algorithm to train basis functions which span the space of dynamical systems, where each basis function is a neural ODE. The FE + ODE approach is put to the test against ODE and FE schemes for fitting a space of Van der Pol systems, long term prediction for RL, as well as quadrotor MPC control with improvements in zero-shot performance against the baselines used.
strength_weakness: strength_weakness: - Some clumsiness in the math notation: in equation (3) $N$ is used at the denominator when it should be $m$ instead.
- The assumption of orthogonality for the validity of the coefficients is vital. However, both the authors in [13] and in this paper do not share any analysis on the validity of this assumption. It would be interesting to see how well this holds in practice (computing inner product of obtained basis functions/ODEs should be easy) and how many iterations are required to reach orthogonality in the examples provided.
- The introductory paragraph of section 3 where the link between a new unknown function $f$, its observed trajectory $\mathcal{D}$ and ODE basis functions is made, could benefit from more motivation and explanation. It is not immediately clear why reasoning at the derivative level while dealing with integral trajectory data is more advantageous. Perhaps a figure would help as the manuscript is text heavy.
- Explanations of the MuJoCo results are not very clear to me, they seem more like observations than interpretations. Why is the oracle unstable (does it require more data and training given the conditioning on hidden params)? What explains the differences between both experiments: FE + RES does well on half cheetah (better than FE + NODE) and terrible on ant? Where does the lack of inductive bias intervene in this case ?
questions: questions: - The term zero-shot does indeed refer to the ability of networks to perform in novel circumstances without retraining, and does apply here. However, it can be slightly misleading as adaptability here requires analysis of new data from the new setting. Can the authors elucidate the scales involved in the tradeoff? In other words, when does it become interesting to train 100 base models (on the 100 datasets) to gain a deployment advantage that still requires new data to function, versus fine-tuning one model or retraining?
- Do the authors have any insights on changes in levels of performance whether the new ODE is in the convex hull of the available basis or outside it? (for example in the Van der Pol example, if the basis datasets contain trajectories for values of $\mu$ between 0.1 and 3, can the system perform for $\mu = 5$ ?
limitations: limitations: - A major numerical challenge working with continuous time neural networks is that of integration, with a plethora of solvers and schemes available. The authors touch upon this topic in their limitations sections and the impact of the integration horizon selection on behaviour predictability as well as the compute overhead involved (also in appendix C). It would have been interesting to give readers a better sense about the tradeoffs with numerical comparisons that go beyond verbal description.

Comments: comment: The authors' rebuttal and general response have improved my understanding and addressed many of my concerns. As a result, I have revised my score to 5.
I thank the authors for the elements presented in the rebuttal and general response. They clear up some points of confusion and improve in my opinion the understanding of the work, although they do not shed light on new strengths of the approach that might have gone unnoticed in my initial review. Thus, the current score still represents my appreciation of the work.
We thank the reviewer for their insightful comments and discussion, which has improved our central arguments and clarified several details. We respectfully maintain that the ability to adapt and generalize in real time, given a small online dataset, has significant impact for the field of autonomy and robotics. Online adaptability is clearly necessary to bring learning-based systems from laboratory settings to the unstructured nature of real world environments. To the best of our knowledge, we provide the first scalable algorithm for learning models of dynamical systems that can adapt near-instantaneously at runtime without any gradient updates and without leveraging prior information. We believe this will have significant impact on the field of autonomy and robotics.
Overall the author did a good job during the rebuttal. The added ablations addressed most of my concerns. I believe that giving readers more insights via ablation is one direction the paper could further improve. I am changing my score to "technical solid paper, accept outweigh reasons to reject", but I'd still expect more impact from a paper for any higher scores.

I strongly encourage the authors to revise their paper to include a detailed discussion about orthonormality, number of basis functions & data in main paper, and probably NeuralODE speed in the appendix..
Many prior works have addressed this question, and leverage neural ODEs for model-predictive control. For example, please see “Taylor-lagrange neural ordinary differential equations: Toward fast training and evaluation of neural odes” (Djeumou, Neary, Goubault, Putot, Topcu, 2022) which proposes a fixed step size method to that achieves a forward pass in 1 millisecond or less. See also “How to Learn and Generalize From Three Minutes of Data: Physics-Constrained and Uncertainty-Aware Neural Stochastic Differential Equations” (Djeumou, Neary, Ufuk Topcu, 2023), “Autonomous Drifting with 3 Minutes of Data via Learned Tire Models” (Djeumou, Goh, Topcu, Balachandran, 2023),  “Learning-enhanced Nonlinear Model Predictive Control using Knowledge-based Neural Ordinary Differential Equations and Deep Ensembles” (Chee, Hsieh, Matni, 2023), “Model predictive control of nonlinear processes using neural ordinary differential equation models” (Luo, Abdullah, Christofides 2023). 
For adaptive step size solvers, the number of forward passes required to integrate a given time horizon depends on the landscape of dx/dt, as you mention. Indeed, this is a serious limitation of adaptive step size solvers, as the solver can require 10-100 forward passes per integration. However, other fixed step size solvers exist, such as higher-order RK methods or Taylor-Lagrange methods, and these methods may achieve better accuracy without incurring significant overhead. In our work, we use RK4, which has only 4 forward passes per integration step, meaning the compute time is much quicker, and we find that this tradeoff is favorable for real-time control. 

Transformer-based MPC is a nascent topic, and has been explored in “Simultaneous multistep transformer architecture for model predictive control”, (Park et. al., 2023), which shows that transformers achieve significant speedups vs. LSTM-based architectures, but are still longer than required for robotics applications (on the order of seconds). In our own observations, transformer-based architectures are too computationally inefficient since they must compute attention on the example data for every forward pass, in a computation that scales quadratically with the amount of data.

We have improved the discussions in Section 3.2 and in the Appendix to highlight these details.
The authors mentioned that transformers are often too computationally complex for model predictive control, but doesn't the same apply to Neural ODE too? From my personal experience, NODE isn't quite fast either, and depending on the landscape of the dx/dt, solver can struggle quite a bit, especially if the authors model a second order system e.g. mujoco locomotion environments' observations are often positions + velocities. Have you tried transformer style architecture for the problems?
**If quality of prediction isn't sensitive to number of basis functions & data after they are increased to a certain degree, does it mean that the method won't be able to scale up further to bring the prediction error to 0, even for a deterministic system. I understand that every predictive model have errors, I am simply looking for some analysis here - what's the source of the remaining error.**

This is a nuanced and complicated question. The biggest thing that likely leads to error is the fact that the system (MuJoCo) is not strictly continuous due to contact forces. Therefore, there will always be some error for any method that uses a neural network due to the contact forces in the environment. Another factor is that the space of functions is infinite-dimensional. Therefore, for any finite number of basis functions, there is inherently going to be error due to the unrepresented dimensions. We do expect to see diminishing returns as the number of basis functions increase, as some of these dimensions are less important for predicting the dynamics than others, and this aligns with the empirical results. We also leverage RK4 as the integrator for the neural ODE. RK4 is a fourth-order method, and higher-order terms are truncated. This choice of integrator therefore imposes error, which increases for larger time horizons. More accurate integrators are available, but come at the cost of compute time. Lastly, numerical precision may play a small role here due to implementation details. Any numerical precision errors will compound as more operations occur. These errors can affect both training, where they effectively appear as noise added to the gradients, and during execution. 

**In my review, I mentioned that “The problems this system can solve seem to be constrained to system limited variation in parameter”. While the authors did an ablation on OOD parameters, it's still limited to the family of OOD that's reflected via parameters. I am wondering whether the system can generalize to OOD dynamics that completely changes its form (e.g. walking to jumping), not just parameters**

No, our approach is not limited to variation in parameters. As stated in the paper in section 3, we are able to handle variations due to changes in the underlying physics model. The only theoretical requirement is that the functions exist in the same space. We believe the MuJoCo examples demonstrate this as our model can generalize beyond simple parameter varying systems, where the robot’s physical shape is changing. However, your question highlights two nuanced questions. The first is a shift in the distribution of states seen between a robot that is walking and jumping. A jumping robot is likely to experience different states than when it is walking. If these states have never been trained on, then any learned model cannot hope to accurately model these dynamics (without leveraging prior knowledge). In other words, a purely learned model which leverages no prior information needs its training set to cover the state space. The second question is if the learned dynamics can generalize across behaviors. Our model has actions as an input, and therefore it can accurately predict a given transition even if the underlying policy was not used to collect training data. Therefore, generalization across behaviors is possible.
I acknowledge that I've read the rebuttal and general response. Thank you for the ablations and they are helpful to understanding. Here are some additional questions: 

1. If quality of prediction isn't sensitive to number of basis functions & data after they are increased to a certain degree, does it mean that the method won't be able to scale up further to bring the prediction error to 0, even for a deterministic system. I understand that every predictive model have errors, I am simply looking for some analysis here - what's the source of the remaining error.

2. In my review, I mentioned that 

> The problems this system can solve seem to be constrained to system limited variation in parameter

While the authors did an ablation on OOD parameters, it's still limited to the family of OOD that's reflected via parameters. I am wondering whether the system can generalize to OOD dynamics that completely changes its form (e.g. walking to jumping), not just parameters