Continuous latent flow modeling for model-based reinforcement learning using temporal transformer networks

Dynamical models are central to our ability to understand and predict natural and engineered systems. However, real-world systems often show time-varying behavior that is too complex for straightforward statistical forecasting approaches. This is due to the fact that the temporal behavior, while potentially explained by an underlying dynamical model, can show strong, possibly abrupt changes in the observation space. To tackle these fundamental challenges, we propose a new model class which is explicitly aimed at predicting dynamical trajectories from high-dimensional empirical fluid flow data. This is done by combining amortized variational autoencoders and spatio-temporal attention within a framework designed to enforce certain scientifically motivated invariances.

The novel model called LaDID allows inference of the system behavior at any continuous time and generalization well beyond the data distributions seen during training. Furthermore, the model does not require an explicit neural ODE formulation, making it efficient and highly scalable in practice. Moreover, we demonstrate that our model forms an ideal basis for model-based reinforcement learning and present benchmark results for various flow environments of the HydroGym platform.