Fig. 9: Depiction of the neural state-space model (NSSM) training methodology.
From: Learning plasma dynamics and robust rampdown trajectories with predict-first experiments at TCV
The NSSM, defined by the dynamics function fθ and observation function Oθ with parameters θ, is simulated forward in time, given an initial state x0 and an action trajectory a0:T, to generate a sequence of simulated observations, \({\hat{{{\bf{o}}}}}_{0:T}\). The simulated observations are compared with experimental observations via the loss functional \({{\mathcal{L}}}\), which is defined as the time-integrated value of an instantaneous loss function l. Adjoint methods in diffrax17 and JAX automatic differentiation then yield the gradient of model parameters with respect to loss, \({\nabla }_{\theta }{{\mathcal{L}}}\), which allows the optimizer to update the parameters θ.
