RL-DAUNCE: Reinforcement Learning-Driven Data Assimilation with Uncertainty-Aware Constrained Ensembles
Pouria Behnoudfar, Nan Chen
TL;DR
RL-DAUNCE addresses the challenge of physically consistent, uncertainty-aware data assimilation for nonlinear, intermittent systems by marrying reinforcement learning with ensemble-based methods. It defines an ensemble of RL agents that mimic EnKF members, trained with constrained EnKF data and guided by a primal-dual optimization framework that dynamically adjusts Lagrange multipliers to enforce energy conservation and positivity constraints. The framework introduces a constraint-augmented Bellman operator and enforces hard action-space bounds, enabling stable, physically meaningful updates while preserving computational efficiency. Applied to the Madden–Julian Oscillation skeleton model, RL-DAUNCE achieves comparable accuracy and uncertainty representation to a constrained EnKF but with dramatically reduced computation time, offering practical benefits for real-time and large-scale assimilation tasks.
Abstract
Machine learning has become a powerful tool for enhancing data assimilation. While supervised learning remains the standard method, reinforcement learning (RL) offers unique advantages through its sequential decision-making framework, which naturally fits the iterative nature of data assimilation by dynamically balancing model forecasts with observations. We develop RL-DAUNCE, a new RL-based method that enhances data assimilation with physical constraints through three key aspects. First, RL-DAUNCE inherits the computational efficiency of machine learning while it uniquely structures its agents to mirror ensemble members in conventional data assimilation methods. Second, RL-DAUNCE emphasizes uncertainty quantification by advancing multiple ensemble members, moving beyond simple mean-state optimization. Third, RL-DAUNCE's ensemble-as-agents design facilitates the enforcement of physical constraints during the assimilation process, which is crucial to improving the state estimation and subsequent forecasting. A primal-dual optimization strategy is developed to enforce constraints, which dynamically penalizes the reward function to ensure constraint satisfaction throughout the learning process. Also, state variable bounds are respected by constraining the RL action space. Together, these features ensure physical consistency without sacrificing efficiency. RL-DAUNCE is applied to the Madden-Julian Oscillation, an intermittent atmospheric phenomenon characterized by strongly non-Gaussian features and multiple physical constraints. RL-DAUNCE outperforms the standard ensemble Kalman filter (EnKF), which fails catastrophically due to the violation of physical constraints. Notably, RL-DAUNCE matches the performance of constrained EnKF, particularly in recovering intermittent signals, capturing extreme events, and quantifying uncertainties, while requiring substantially less computational effort.