A CNN encoder for modal phase reconstruction in Adaptive Optics systems

A CNN encoder for modal phase reconstruction in Adaptive Optics systems P. Vermot, D. Gratadour November 24, 2025 ABSTRACT Context: Pyramid wavefront sensing (pWFS) offers some of the best sensitivity for adaptive optics, but suffers from strong non-linearity. Modulation can extend linear range at the expense of sensitivity. Operating the pyramid without modulation is therefore attractive, but remains challenging. Aims: We develop a non-linear CNN reconstructor for pyramid wavefront sensing that maps pWFS images directly to a modal phase representation, and we assess simple hybrid methods combining it with a standard linear least-squares (LS) reconstructor. Methods: Using the end-to-end COMPASS simulator, we generate a large open-loop dataset spanning wide ranges of RMS and power spectra to train a compact CNN encoder. We then compare the CNN against an LS baseline and evaluate hybrid schemes in closed-loop simulations over a grid of guide-star magnitudes and Fried parameters, reporting long-exposure Strehl at λ = 1 . 6 µ m with controller gain re-optimized per method and bin. We also report preliminary offline bench tests on SCExAO. Results: The CNN reduces open-loop reconstruction error and, in closed loop, outperforms LS across most (magnitude, r0 ) conditions, with the largest gains for faint stars where it often closes the loop while LS does not. In very strong turbulence, LS can exceed the CNN; in these cases the hybrid methods are necessary to surpass LS, with a second-stage NN performing best. Inference is real-time capable and the hybrid overhead is negligible. Bench snapshots on SCExAO show successful correction of small static/slow perturbations, with instability for stronger/faster cases. Conclusions: A compact CNN can be trained to perform modal reconstruction for a nonmodulated pWFS and improves performance over a classical linear reconstructor in most regimes. When the CNN alone is not optimal, simple hybrid methods achieve the best performance, suggesting a practical way to exploit the pyramid’s sensitivity without modulation. 1. INTRODUCTION Adaptive optics (AO) can restore diffraction-limited performance on large telescopes by estimating and correcting atmospheric phase in real time. Among wavefront sensors (WFS), the pyramid WFS (pWFS) [Ragazzoni, 1996, V´erinaud, 2004] has become a leading choice for current and future instruments thanks to its high sensitivity [Esposito and Riccardi, 2001]. However, its response is strongly non-linear and in many cases requires modulation to increase the linear range, which is done at the expense of sensitivity [Burvall et al., 2006]. Despite its strong nonlinearities, non-modulated pyramid is thus of high interest for increased sensitivity [Guyon et al., 2011, Agapito et al., 2023]. In this work we directly target the nonlinearity of the pWFS by training a Convolutional Neural Network (CNN) to map WFS images to a modal phase representation. We adopt the Btt modal basis [orthonormal modes derived from the DM influence functions with pure tip–tilt separation Fer, Pou]. We train a compact CNN encoder on a large synthetic dataset produced with the COMPASS end-to-end simulator, spanning wide magnitude and turbulence regimes. We compare its performances against a least-squares (LS) linear reconstructor, which is the pseudoinverse of the interaction matrix [Southwell, 1980, Roddier, 1999]. Motivated by regimes where a linear estimator may still excel, we also test simple hybrid schemes that combine CNN and linear estimates. Our contributions are: (i) a light-weight CNN encoder that performs modal reconstruction from PWFS frames; (ii) extensive closed-loop evaluation across guide star magnitude and Fried parameter r0 , showing that the CNN dominates in most regimes of interest; and (iii) hybrid methods that take Table 1: COMPASS configuration. Pupil Diameter D8 m Central obscuration 0.14 D(1.12 m) Atmosphere Number of layers 3 Layer strengths [0.50,0.35,0.15] Layer altitudes [0,4,10] km Wind speed [15,15,35] m s−1 Wind direction [0◦,20◦,180◦] Outer scale L030 m Wavefront sensor Type Pyramid Modulation rmod = 0 Wavelength 0.85 µm Subapertures 56 ×56 Optical throughput 0.15 DM High-order (PZT) 40 ×40; Gaussian IF Tip-Tilt Controller Type LS; delay 1 frame Timing 1 ms (1 kHz) Science channel Wavelength λ0= 1.67 µm the best of each estimator depending on conditions. We briefly report preliminary semi-successful bench tests on SCExAO to illustrate practical deployment aspects Guyon et al. [2010], Jovanovic et al. [2015], Lozi et al. [2018]. 2. METHODS 2.1 Simulation setup and dataset generation The work presented in this paper is done with end-to-end simulations using the COMPASS tool Fer. We simulate a modern AO system on an 8 m telescope with a 14% central obscuration, a pyramid WFS operated without modulation (56 × 56 subapertures at λ = 0 . 85 µ m), and a 40 × 40 PZT deformable mirror complemented by a tip–tilt stage. The control loop runs at 1 kHz with one-frame delay, and the atmosphere is a 3-layer model with L0 = 30 m. A summary of key parameters is given in Table 1. 2.2 Open-loop data generation We draw a power-law slope α∼ U (0 , 3) and synthesize a random phase in the Fourier domain with amplitude A ( κ ) ∝ ∥κ∥−α and random phase θ ( κ ) ∼ U [0 , 2 π ) (piston removed by setting A (0) = 0): bφ0(κ) = A(κ)eiθ(κ), φ0(x) = ℜF−1[bφ0]. To be able to export this method to a hardware bench later on, we then project φ0onto the DM: φ1(x) = B BTφ0(x). We set a target RMS σφ=r−ln U 2π, U ∼ U(0,1), and normalize over the pupil: φ(x) = σφ φ1(x) stdφ1(x). We then use COMPASS to generate the corresponding WFS frame with the guide-star magnitude drawn uniformly m∼ U (3 , 13). The WFS images are pre-formatted into an array Iby extracting four 60 × 60 patches, stacking them into a 3-channel array, and normalizing this array by its overall standard deviation. Table 2: CNN encoder (default configuration). Stage Op. (kernel / stride / act) Output size (C,H,W) Input — (4,60,60) Block 1 Conv 3×3 / (2,2) + BN + softplus (32,30,30) Block 2 Conv 3×3 / (2,2) + BN + softplus (128,15,15) Block 3 Conv 3×3 / (2,2) + BN + softplus (512,8,8) Block 4 Conv 3×3 / (2,2) + BN + softplus (2048,4,4) Block 5 Conv 3×3 / (2,2) + BN + softplus (8192,2,2) Head Conv 1×1 / (2,2) + linear (1307,1,1) Output Flatten 1307 Finally, we compute the DM modal coefficients by projecting φ(x) onto the modal basis: a=BTφ, We generate 10 6 [input, output] training examples and 1 . 5 × 10 5 validation examples, each consisting of a couple [I,a] with shapes [(4, 60, 60), (1307)]. 2.3 CNN 2.3.1 Architecture The network maps the pre-formatted pWFS tensor I(4 , 60 , 60) to Nmodes = 1307 modal coefficients. It is a strided convolutional encoder with Batch Normalization and softplus activations. We use L = 5 blocks of 3 × 3 convolutions with stride (2 , 2); at each block the channel count grows by a factor of four (starting from 32), followed by a final 1 × 1 convolution with stride equal to the current spatial size to collapse the feature map to 1×1, and a linear activation on the output. 2.3.2 Training We train on 10 6 examples (with additional 1 . 5 × 10 5 validation examples) using Adam optimizer ( β1= 0 . 9 , β2= 0 . 999), batch size 256, for 100 epochs. The learning rate decays exponentially from 10 −4 to 10 −7 over training. The loss is mean squared error between predicted and true modal coefficients. Computations ran on the Jean–Zay (IDRIS) supercomputer on a single NVIDIA H100 GPU with 24 CPU cores (Intel Xeon Platinum 8468): with batch size 256 the per-batch compute time is 22 ms (about 92 s per epoch), totaling slightly under 3 hours for all 100 epochs; at inference, the per-batch compute time is 3 ms (batch size 256). 2.4 Hybridization We combine the CNN and LS estimates in modal space via linear combination of modes. Let acnn, alin be the predicted modal vectors and a⋆the ground truth. We define ahyb i=wiacnn i+(1−wi)alin i, i = 1, . . . , Nmodes, with weights wi∈ [0 , 1]. We consider three weight strategies, fitted on a separate calibration split: 1. Median weights. For each sample n and mode i , the oracle weight minimizing the squared error is w(n) i,orc = clip[0,1]a⋆ i,(n)−alin i,(n) acnn i,(n)−alin i,(n).We set wi= mediannw(n) i,orc. 2. MMSE weights. Assuming zero-mean estimation errors ecnn i = acnn i−a⋆ i and elin i = alin i−a⋆ i , the MSE-optimal static weight per mode is wmmse i=Var(elin i)−Cov(ecnn i, elin i) Var(ecnn i) + Var(elin i)−2 Cov(ecnn i, elin i), estimated from the calibration split and clipped to [0,1]. 3. Second-stage DNN. A small fully connected network (1 hidden layers with 1307 neurons and ReLU activation), taking ( acnn, alin ) as input and outputting w∈ [0 , 1] Nmodes via a sigmoid, trained with MSE loss on the oracle weights wi,orc. 0.5 1.0 1.5 2.0 2.5 PL index 0.2 0.4 0.6 0.8 SR actus_mse_loss_CNN 0.00 0.25 0.50 0.75 1.00 1.25 1.50 1.75 2.00 lossCNN lossRTC Figure 1: Open-loop reconstruction error ratio in the actuator domain: per-bin (by RMS σφ and power-law index α ) median of losscnn/losslinear . Values < 1 indicate lower MSE for the CNN; values >1 favor the linear reconstructor. 3. RESULTS 3.1 CNN Open-loop results To determine the performance level of our CNN, the first test is to compare the quality of its reconstruction on open-loop data. For this purpose, we ran inference on our entire validation dataset (1 . 5 × 10 5 examples) with our CNN and the linear reconstructor and, for each sample, computed the mean squared error (MSE) with respect to the ground truth in the actuator (zonal) domain. In Fig. 1, we present a comparison of the two algorithms binned by the target RMS σφ and the power-law index α of the residual atmospheric phase, showing the per-bin median ratio losscnn/losslinear . These results indicate ratios below 1 in the vast majority of bins, i.e., the CNN outperforms the linear reconstructor across most of the explored grid. 3.2 CNN Closed-loop results As a second step, we included the CNN reconstructor in closed-loop simulations spanning a wide range of magnitude and turbulence conditions. Long-exposure (LE) Strehl ratio is evaluated at the science wavelength ( λ = 1 . 6 µ m) and integrated over the full sequence of 1500 iterations per simulation. For fairness, the loop gain is optimized per method and per bin by testing all gains from 0 . 05 to 1 . 0 in steps of 0 . 05 and retaining the best result. The Fried parameter r0 is specified at 500 nm. We test m∈[3,12] in steps of 0.5 and r0∈[0.05,0.39] m in steps of 0.02 m. In Fig. 2 we compare the LE Strehl obtained with the CNN and with the linear reconstructor (LS) as a function of guide-star magnitude and r0 . The plot shows the per-bin ratio SRCNN/SRLS ; blue indicates CNN better (ratio > 1), red indicates LS better (ratio < 1), and grey marks bins where neither configuration “closes the loop,” defined empirically as SR ≤ 0 . 1. Overall, in the great majority of cases the CNN performs as well as or better than LS, and only in a minority of cases does it underperform: 1. In the faint guide-star regime (mag > 10), the CNN outperforms LS with ratios ≳ 2 (deep blue in Fig. 2); above mag ≈ 11 . 5, LS frequently does not close the loop while the CNN does. 2. In the standard operating regime (mag < 10, r0> 0 . 10 m), CNN and LS perform very similarly, with a slight advantage to the CNN (ratios close to but above 1). 3. In the very high turbulence regime ( r0 = 5 cm), LS performs better than the CNN (ratios <1). 0.05 0.07 0.09 0.11 0.13 0.15 0.17 0.19 0.21 0.23 0.25 0.27 0.29 0.31 0.33 0.35 0.37 0.39 r 0 (m) 3.0 3.5 4.0 4.5 5.0 5.5 6.0 6.5 7.0 7.5 8.0 8.5 9.0 9.5 10.0 10.5 11.0 11.5 12.0 Magnitude LE ratio (masked <0.1 numerator): CNN / Linear 0.00 0.25 0.50 0.75 1.00 1.25 1.50 1.75 2.00 LE Strehl ratio (relative) Figure 2: Closed-loop LE Strehl ratio at λ = 1 . 6 µ m, shown as the per-bin ratio SRCNN/SRLS over the grid m∈ [3 , 12] (step 0 . 5) and r0∈ [0 . 05 , 0 . 39] m at 500 nm (step 0 . 02 m). Blue ( > 1) indicates CNN better; red ( < 1) indicates LS better; grey marks bins where neither closes the loop (empirical threshold SR ≤0.1). 3.3 Hybrid methods closed-loop results We now present closed-loop results for the three hybrid reconstructors (median weights, MMSE weights, and second-stage NN), in addition to the LS and CNN baselines. Like previously, LE Strehl is evaluated at λ = 1 . 6 µ m over the full 1500-iteration sequence, r0 is specified at 500 nm, and for each method and bin the loop gain is optimized. Hybrid weights and the second-stage NN are calibrated on the validation split. In Fig. 3, we show the evolution of the long-exposure Strehl ratio for the five reconstructors (linear, CNN, median weights, MMSE weights, and second-stage NN) in the three characteristic regimes identified earlier (faint star, regular conditions, and high turbulence): 1. In the “faint guide-star regime” (mag > 10), the linear reconstructor provides the lowest performance (SR ∼ 0 . 1), and the CNN performs much better (SR ∼ 0 . 5). More generally, the three hybrids improve over the CNN; sorted by increasing performance: second-stage NN, median weights, MMSE weights. 2. In the “standard operating regime” (mag < 10, r0> 0 . 1 m), all reconstructors have high performance (SR > 0 . 9). The linear reconstructor is significantly below the others (SR ∼ 0 . 92 vs. ∼0.95), while the non-linear and hybrid methods are statistically indistinguishable. 3. In the “very high turbulence regime” ( r0 = 5 cm), the CNN is the worst (SR ∼ 0 . 1). Median and MMSE weights improve it (SR ∼ 0 . 15) yet remain below linear (SR ∼ 0 . 18). The second-stage NN outperforms all others in this regime (SR ∼0.22). Generalizing to the full ( m, r0 ) grid, Fig. 4 shows a color-coded map indicating which of the five methods is best in each bin: 1. In the top-right (faint star with good atmospheric conditions), the MMSE-weight hybrid dominates. 2. In the bottom-left (bright star with bad atmospheric conditions), the second-stage NN provides the best results. 3. In the valley between these regions, the CNN alone is best; note that here all non-linear methods have very close performance and differences are not significant. 0 200 400 600 800 1000 1200 1400 Iteration 0.00 0.05 0.10 0.15 0.20 0.25 0.30 0.35 LE Strehl ratio Full LE vs Iteration @ mag=5.00, r0=0.05 linear CNN median wts MMSE 2nd stage NN 0 200 400 600 800 1000 1200 1400 Iteration 0.2 0.4 0.6 0.8 1.0 LE Strehl ratio Full LE vs Iteration @ mag=7.00, r0=0.21 linear CNN median wts MMSE 2nd stage NN 0 200 400 600 800 1000 1200 1400 Iteration 0.1 0.2 0.3 0.4 0.5 0.6 0.7 LE Strehl ratio Full LE vs Iteration @ mag=11.50, r0=0.29 linear CNN median wts MMSE 2nd stage NN Figure 3: Representative long-exposure Strehl ratio examples at different guide-star magnitudes and turbulence strengths. Each panel shows the five reconstructors (LS, CNN, median weights, MMSE weights, second-stage NN). 0.05 0.10 0.15 0.20 0.25 0.30 0.35 r0 3 4 5 6 7 8 9 10 11 12 magnitude Best Method by (mag, r0) linear CNN median wts MMSE 2nd stage NN None ( 0.1) Figure 4: Winner map over the ( m, r0 ) grid ( m∈ [3 , 12] step 0 . 5; r0∈ [0 . 05 , 0 . 39] m at 500 nm step 0 . 02 m). Color indicates the method with the highest LE Strehl at λ = 1 . 6 µ m (no tie threshold). Grey marks bins where no method closes the loop (empirical threshold SR ≤ 0 . 1). Gains are optimized per method and bin (sweep 0.05:0.05:1.0). 0.05 0.07 0.09 0.11 0.13 0.15 0.17 0.19 0.21 0.23 0.25 0.27 0.29 0.31 0.33 0.35 0.37 0.39 r 0 (m) 3.0 3.5 4.0 4.5 5.0 5.5 6.0 6.5 7.0 7.5 8.0 8.5 9.0 9.5 10.0 10.5 11.0 11.5 12.0 Magnitude LE ratio (masked <0.1 numerator): CNN + Second stage modes / Linear 0.00 0.25 0.50 0.75 1.00 1.25 1.50 1.75 2.00 LE Strehl ratio (relative) Figure 5: Performance ratio over the ( m, r0 ) grid: SRCNN+2nd/SRLS at λ = 1 . 6 µ m. Blue ( > 1) indicates CNN+second-stage NN better; red ( < 1) indicates LS better; grey marks bins where neither closes the loop ( SR ≤ 0 . 1). Colorbar range 0 → 2. Gains are optimized per method and bin. There is not a single method outperforming all others in every regime, but there is no regime in which the linear reconstructor is the best: including the CNN in the phase reconstruction systematically improves performance. 3.4 SCExAO A preliminary version of the method was tested on the SCExAO instrument in an offline laboratory configuration (not on-sky) in April 2025 [Jovanovic et al., 2015, Guyon et al., 2010, Lozi et al., 2018]. The bench used a non-modulated pWFS at 800 nm and a deformable mirror; performance was estimated in the science channel in H-band (broad band). Although the implementation was not yet complete, we summarize the experiment here. The SCExAO instrument is equipped with a pyramid wavefront sensor and a deformable mirror. In the offline setup, we implemented a preliminary version of our algorithm as follows: 1. Generation of random phase screens and projection onto the DM. 2. Acquisition of the associated pWFS frames (non-modulated, 800 nm). 3. Training of a CNN to reconstruct the true modes from the pWFS images obtained on the bench; 200,000 examples). 4. Closed-loop test using the CNN reconstructor (no gain tuning; default controller settings), assessing H-band long-exposure PSFs. The CNN could successfully compensate a small static perturbation in the system and a faint, very slowly evolving turbulent wavefront (see Fig. 6). However, when stronger or more rapidly evolving perturbations were injected, the loop became unstable and diverged. A likely cause is a mismatch between training and inference conditions (e.g., WFS reference slopes). Resolving this will require improved on-bench calibration (reference update) and gain tuning, which we leave to future work. (a) With small static perturbation (before correction). (b) After CNN correction (closed loop). Figure 6: Illustrative SCExAO bench snapshots in H-band. These are screenshots taken during the experiment and are shown for qualitative illustration only; quantitative parameters (e.g., Strehl ratio, exposure time) were not recorded. 4. DISCUSSION Our results show that a CNN reconstructor can close the loop in regions where a linear least-squares (LS) reconstructor cannot. Across the magnitude– r0 grid the CNN provides better performance than the linear reconstructor in the vast majority of cases and most notably for faint stars, where it often closes the loop when LS does not, and also in most standard operating regimes where it yields a few additional percent of Strehl. Due to computational limits we used a single realization per bin; however, the spatial coherence of the winner map across neighboring bins supports the stability of these conclusions. In very high turbulence regimes the linear reconstructor outperforms our CNN. Although this is where the pyramid is most non-linear, we believe the primary cause is an under-representation of highRMS wavefronts in the training set. Increasing the fraction of high-RMS examples is straightforward in principle, but at strong turbulence the pWFS response saturates and the “ground-truth” modal vector becomes impossible to determine, which prevents stable supervised training and hinders convergence. Addressing this will require specific data generation and calibration strategies that remain valid in the saturated regime, to be determined. Hybrid methods improve upon the CNN in several parts of the grid. Although no single hybrid outperforms all others everywhere, including the CNN estimate in the reconstruction is always advantageous: in many cases the CNN alone is sufficient; in others, combining it with the linear reconstructor yields the best performance. The computational overhead of the hybrids is negligible compared to the CNN forward pass. With not fully optimized workflows (but TensorFlow graph compilation) we measured sub-millisecond inference per batch on the SCExAO bench with an NVIDIA A6000, and even shorter inference times on the H100 for the full pipeline (CNN + hybrid), with the hybrid stage being negligible relative to the CNN. Preliminary bench tests on SCExAO were promising but not yet fully successful. The CNN could correct a small static perturbation and a faint, slowly varying turbulent wavefront, but destabilized for stronger or faster perturbations. A likely contributor is a mismatch between training and inference conditions (e.g., WFS reference slopes), compounded by the absence of gain tuning in that run; we plan to revisit the bench with improved reference updates and gain optimization. This work has several limitations and natural extensions. First, only one realization per bin was used; increasing the number of realizations will allow formal significance estimates. Second, a systematic robustness/ablation study (e.g., training-set composition, weighting of high-RMS cases, conditioning on flux/optical gain, registration jitter) remains to be completed. Third, while this study focused on an unmodulated pyramid, applying the same CNN and hybrid approaches to a modulated pyramid is an interesting direction for future work. Finally, scalability to ELT-class systems and broader instrument configurations will be investigated in subsequent work. 5. CONCLUSIONS We introduced a compact CNN reconstructor for the non-modulated pyramid WFS that estimates modal coefficients directly from pWFS frames, and we evaluated simple hybrid methods with the CNN and a linear LS baseline. In open loop, the CNN reduces actuator-space MSE in the majority of (RMS, power-law) bins. In closed loop, the CNN outperforms LS across most of the ( m, r0 ) grid, with the largest gains for faint stars (e.g., m > 10; ratios ≳ 2), while LS often fails above m≈ 11 . 5 and the CNN still closes. In standard conditions ( m < 10, r0> 0 . 10 m) the gains are modest but consistent. Under very strong turbulence ( r0≈ 0 . 05 m), LS exceeds the CNN, likely due to underrepresentation of high-RMS in training and pWFS saturation. Hybrids further improve performance: median/MMSE weights or a second-stage DNN match or exceed the CNN or LS reconstructors alone, with no regime where LS is overall best. The approach is real-time capable (sub-ms inference on A6000; even faster on H100), and hybridization overhead is negligible. Preliminary SCExAO bench tests (offline) showed successful correction of small static and slow perturbations, with instability for stronger/faster cases (likely from WFS reference mismatch and lack of gain tuning). Future work will extend training to high-RMS and saturated regimes, add MMSE/Fourier baselines, test the method on a modulated pyramid, and assess scalability toward ELT-class systems and on-sky operation.