Using neural ODE to estimate dynamics of forced systems

[saihv]

Let's say I have a system that I'd like to describe as x_dot = x+u and given information about when and for how long u was applied; I wish to predict the evolution of x_t+k given x_t and u_t. Ideally, I want the model to be able to generalize for any value of u. For example, a cartpole with initial conditions x and theta, and an input force F.

Can the neural ODE framework deal with x_dot being f(x, u)? How do I go about including this input parameterization in the neural ODE framework? My first thought was to just augment the input state with the time dependent value of u when passing it to the neural network, under the hope that the NN will resolve the relationship between the x_t and u_t when predicting x_t+1, but I haven't had much success with that yet.