Integrating deep neural networks with COSMIC for real-time control

Abstract

We present results on integrating machine learning (ML) methods for adaptive optics control with a real-time control library: COmmon Scalable and Modular Infrastructure for real-time Control (COSMIC). We test the integration on simulations for the instrument SAXO+. Our proposed solution's pipeline is formed by a two-model ML system. The first model consists of a very deep neural network (DNN) that maps wavefront sensor (WFS) images to phase and is trained offline. The second model consists of predictive control with a more compact DNN. The predictive control stage is trained online, providing an adaptive solution to changing atmospheric conditions but adding extra complexity to the pipeline. On top of implementing the solution with COSMIC, we add a set of modifications to provide faster inference and online training. Specifically, we test NVIDIA's TensorRT to accelerate the DNNs inference, reduced precision, and just-in-time compilation for PyTorch. We show real-time capabilities by using COSMIC and improved speeds both in inference and training by using the recommendations mentioned above.