Focal plane wavefront control with model-based reinforcement learning

Abstract: The direct imaging of potentially habitable exoplanets is one prime science case for high-contrast imaging instruments on extremely large telescopes. Most such exoplanets orbit close to their host stars, where their observation is limited by fast-moving atmospheric speckles and quasi-static non-common-path aberrations (NCPA). Conventional NCPA correction methods often use mechanical mirror probes, which compromise performance during operation. This work presents machine-learning-based NCPA control methods that automatically detect and correct both dynamic and static NCPA errors by leveraging sequential phase diversity. We extend previous work in reinforcement learning for AO to focal plane control. A new model-based RL algorithm, Policy Optimization for NCPAs (PO4NCPA), interprets the focal-plane image as input data and, through sequential phase diversity, determines phase corrections that optimize both non-coronagraphic and post-coronagraphic PSFs without prior system knowledge. Further, we demonstrate the effectiveness of this approach by numerically simulating static NCPA errors on a ground-based telescope and an infrared imager affected by water-vapor-induced seeing (dynamic NCPAs). Simulations show that PO4NCPA robustly compensates static and dynamic NCPAs. In static cases, it achieves near-optimal focal-plane light suppression with a coronagraph and near-optimal Strehl without one. With dynamics NCPA, it matches the performance of the modal least-squares reconstruction combined with a 1-step delay integrator in these metrics. The method remains effective for the ELT pupil, vector vortex coronagraph, and under photon and background noise. PO4NCPA is model-free and can be directly applied to standard imaging as well as to any coronagraph. Its sub-millisecond inference times and performance also make it suitable for real-time low-order correction of atmospheric turbulence beyond HCI.

• The paper introduces PO4NCPA, a model-based RL algorithm that directly corrects non-common path aberrations using focal-plane intensity measurements.
• The method leverages sequential phase diversity and neural network-based dynamics to optimize deformable mirror commands, attaining near-optimal Strehl ratios and up to 40x contrast improvement.
• Experimental results in both static and dynamic regimes, including ELT-scale simulations with high noise, demonstrate robust performance and potential for real-time control.

Model-Based Reinforcement Learning for Focal Plane Wavefront Control

Introduction

The paper "Focal plane wavefront control with model-based reinforcement learning" (2604.00993) addresses the fundamental problem of non-common path aberration (NCPA) correction in high-contrast imaging (HCI) instruments, focusing on the direct imaging of exoplanets with ground-based extremely large telescopes. NCPA—both static and dynamic—are principal limitations for modern adaptive optics (AO) systems, with conventional correction strategies often requiring intrusive off-sky calibrations or slow corrective probing. This work proposes a model-based deep reinforcement learning (RL) algorithm, Policy Optimization for Non-Common Path Aberrations (PO4NCPA), for focal-plane wavefront sensing and control, leveraging sequential phase diversity and past telemetry. The algorithm is numerically validated in static and dynamic NCPA regimes, on both classical (standard imaging, SI) and coronagraphic (including vortex coronagraph) optical setups.

Problem Statement and RL Formulation

The challenge is to correct NCPA using only intensity measurements in the science focal plane. The inverse problem is ill-posed due to phase ambiguities: for many pupil-phase aberrations, two distinct phase patterns can produce identical focal-plane images, particularly acute under coronagraphy. The authors recast focal-plane wavefront correction as a Markov Decision Process, with the policy parameterized by a neural network mapping from a sequence of focal-plane images and previous DM (deformable mirror) commands to the next optimal DM command. The learning objective is to maximize a reward function related to PSF sharpness (SI) or post-coronagraphic residual flux (PC/VVC), fully learned from telemetry without prior model assumptions.

For efficient learning, the RL architecture is model-based: the system first learns a dynamics model (NN) to predict focal-plane observations given prior state-action pairs, then this model is used to optimize the control policy via backpropagation through simulated closed-loop rollouts. Key to resolving phase ambiguities is the use of sequential phase diversity—explicitly including the prior observation and corresponding DM command as part of the state.

Observation Model and Data Preprocessing

The observation vector for the RL agent consists of the current and previous focal-plane images, as well as the previous DM command. For non-coronagraphic imaging, the algorithm subtracts the ideal (Airy) PSF and applies a cube-root intensity scaling to flatten the dynamic range and reduce photon noise amplification. For coronagraphic configurations, the ideal image is dark and thus subtracted intensity is directly treated. This normalization ensures the NN models can efficiently learn relevant speckle and aberration features.

Figure 1: Data preprocessing for focal plane observations—background-subtracting the ideal PSF (or dark image under coronagraphy) and flattening via cubic root transformation.

Neural Network Architecture for Dynamics and Policy

The focal-point of the approach is simultaneous dynamics and policy model optimization:

• The dynamics model consists of CNN (for 2D image data) and fully-connected layers (processing DM command vectors), forming a network that predicts the next observation from the tuple (observation history, previous command, current command).
• An ensemble of such models, trained via bootstrapped data subsets, is used to mitigate overfitting and provide robust simulation of stochastic system dynamics.

  Figure 2: Deep neural network (CNN + FC) structure for the dynamics model, simultaneously ingesting focal-plane science images and DM commands.

• The policy model has similar architecture but outputs DM commands to maximize future cumulative reward under the learned dynamics. Training is conducted using simulated horizon rollouts through the dynamics, followed by gradient-based update steps to the policy parameters.

  Figure 3: Policy network architecture for processing state (image + DM command history) and outputting DM corrections.

Experimental Protocol: Simulation Setup

The approach is evaluated on numerically simulated 39.3-m telescope data with both circular and ELT pupils, in both static and dynamic NCPA regimes. Both standard imaging (SI) and coronagraphic regimes (PC, vector vortex coronagraph VVC) are examined. The effect of dynamic NCPA is modeled using a frozen flow (Taylor hypothesis) with Kolmogorov turbulence, matched to physically realistic water vapor seeing parameters for mid-infrared AO.

The DM is controlled via Zernike modes (55 modes), with the RL agent observing $33 \times 33$ pixel science images per time step. For dynamic NCPA, the typical loop rate is 10 Hz, without AO residuals, and full photon/background noise is injected for ELT-METIS scenarios.

Training and Convergence

Training proceeds episode-wise, combining a warm-up phase (random actions) with alternating dynamics and policy updates per episode. For dynamic NCPA, an initial phase has the DM flattened at each episode start (to prevent open-loop saturation), followed by continuous closed-loop operation from the previous episode's end—mimicking on-sky operations.

Figure 4: Episode-wise convergence of PO4NCPA for both static and dynamic NCPA correction in SI and PC configurations; reward convergence traces (cumulative, smoothed) indicate rapid and stable learning behavior.

Results

Static NCPA Correction

PO4NCPA achieves near-optimal Strehl ratio for SI (99.4%, close to the theoretical fitting error limit of 99.6%) and for PC achieves lower or equal focal-plane flux than DM modal least-squares projection, indicating additional suppression of starlight via optimal DM superpositions not accessible to standard projection. Wavefront RMS error in the control region for SI is 9–26 nm, and 2–5 nm for PC.

Figure 5: Static NCPA correction traces: reward convergence and residual RMS per episode for PO4NCPA, compared with least-squares fitting error.

Figure 6: Average PSF and radial contrast curves (500 runs) with static NCPA for SI and PC; PO4NCPA approaches or exceeds the performance of reference correction.

Dynamic NCPA and Predictive Control

For dynamic NCPA driven by wind (WV seeing), PO4NCPA demonstrates residual error and Strehl ratio comparable to the best physically realizable control (least-squares with one-step delay). Performance aligns with delay-integrator limits for ambiguous modes, while showing enhanced prediction for unambiguous modes (e.g., odd Zernikes in SI).

Figure 7: Wavefront RMS time traces (100-step window) under dynamic NCPA: PO4NCPA matches the reference integrated correction.

Figure 8: Modal RMS decomposition in dynamic NCPA: PO4NCPA outperforms delay control for unambiguous low-order modes.

Figure 9: Long-exposure PSF contrast for dynamic NCPA in SI and PC: PO4NCPA matches delay-integrator performance.

The methodology remains robust for open-loop RMS up to $\sim$1900 nm, corresponding to open-loop Strehl of $\sim$30%. At higher levels, direct training occasionally fails but transfer of a policy trained at lower error yields partial control.

Figure 10: PO4NCPA Strehl response versus increasing NCPA RMS; transfer learning permits partial extension into regimes with failed direct training.

ELT-METIS/VVC/Noisy Case

For a high-noise, ELT-pupil + vortex coronagraph configuration, the method continues to deliver effective NCPA suppression, yielding up to 40x improvement in focal-plane contrast at small separations compared to open-loop.

Figure 11: Residual post-coronagraphic PSF contrast for a METIS-like (ELT pupil + VVC + noise) configuration: PO4NCPA approaches the delay-integrator benchmark.

Discussion and Implications

The results empirically validate the claims that a model-based RL architecture can, with minimal physical prior, provide direct focal-plane based NCPA control at the theoretical performance limits dictated by the number of DM modes and the intrinsic phase ambiguities of the optical system. The framework is agnostic to coronagraph design—a critical requirement for operations with next-generation ELTs where complex phase masks (e.g., VVC) and asymmetric pupils are commonplace.

Key implications:

1. Model-Free Application: The method eliminates the need for explicit optical modeling, calibration, or external WFS in the NCPA correction path. All phase retrieval (including phase sign disambiguation) is performed by NN-based exploitation of sequential telemetry and induced diversities.
2. Low Latency Real-Time Control: Once trained, inference requires less than 1 ms per step, suitable for real-time control at AO/SCI camera frame rates.
3. Versatility and Robustness: The method generalizes without modification to high noise, strongly aberrated regimes, and arbitrary focal plane mask designs. Its performance degrades gracefully as ambiguity (e.g., in even Zernike modes, highly symmetric masks) increases.
4. Limitations and Future Directions: Training relies on large simulated/synthetic datasets due to the slow on-sky (or lab) convergence, indicating the need for pre-training and potential domain adaptation strategies. Modal correction remains limited by phase ambiguity inherent in the SI and symmetric coronagraphs; enhanced phase diversity strategies or asymmetric pupil masks may further break ambiguities. Extending reward regions (e.g., to targeted dark holes) will tailor starlight suppression for direct exoplanet imaging.

Conclusion

This work establishes policy optimization via model-based RL as a competitive paradigm for focal-plane wavefront control, matching or surpassing traditional NCPA correction strategies for both static and dynamic error regimes across diverse optical configurations. The approach leverages the flexibility of neural networks and the efficiency of simulated model rollouts, overcoming the phase ambiguity problem through sequential image-action telemetry. Its applicability to current and future HCI systems, including challenging ELT scenarios, demonstrates its strategic significance for the optical/astronomical AO community. Continued advances in training strategies, hyperparameter optimization, and on-sky pre/post-training integration will be key for routine deployment of RL-driven focal-plane control on high-contrast imagers.