Reinforcement Learning State Estimation for High-Dimensional Nonlinear Systems

In high-dimensional nonlinear systems such as fluid flows , the design of state estimators such as Kalman filters relies on a reduced-order model ( ROM ) of the dynamics . However , ROMs are prone to large errors , which negatively affects the performance of the estimator . Here , we introduce the reinforcement learning reduced-order estimator ( RL-ROE ) , a ROM-based estimator in which the data assimilation feedback term is given by a nonlinear stochastic policy trained through reinforcement learning . The flexibility of the nonlinear policy enables the RLROE to compensate for errors of the ROM , while still taking advantage of the imperfect knowledge of the dynamics . We show that the trained RL-ROE is able to outperform a Kalman filter designed using the same ROM , and displays robust estimation performance with respect to different reference trajectories and initial state estimates . 1 INTRODUCTION . Active control of turbulent flows has the potential to cut down emissions across a range of industries through drag reduction in aircrafts and ships or improved efficiency of heating and air-conditioning systems , among many other examples ( Brunton & Noack , 2015 ) . But real-time feedback control requires inferring the state of the system from sparse measurements using an algorithm called a state estimator , which typically relies on a model for the underlying dynamics ( Simon , 2006 ) . Among state estimators , the Kalman filter is by far the most well-known thanks to its optimality for linear systems , which has led to its widespread use in numerous applications ( Kalman , 1960 ; Zarchan & Musoff , 2015 ) . However , continuous systems such as fluid flows are governed by partial differential equations ( PDEs ) which , when discretized , yield high-dimensional and oftentimes nonlinear dynamical models with hundreds or thousands of state variables . These high-dimensional models are too expensive to integrate with common state estimation techniques , including the Kalman filter or its numerous extensions . Thus , state estimators are instead designed based on a reduced-order model ( ROM ) of the underlying dynamics ( Barbagallo et al. , 2009 ; Rowley & Dawson , 2017 ) . A big challenge is that ROMs provide a simplified and imperfect description of the dynamics , which negatively affects the performance of the state estimator . One potential solution is to improve the accuracy of the ROM itself through the inclusion of additional closure terms ( Ahmed et al. , 2021 ) . In this paper , we leave the ROM untouched and instead propose a new design paradigm for the estimator itself , which we call a reinforcement-learning reduced-order estimator ( RL-ROE ) . The RL-ROE is constructed from the ROM in an analogous way to a Kalman filter , with the crucial difference that the linear filter gain function is replaced by a nonlinear stochastic policy trained through reinforcement learning ( RL ) . The flexibility of the nonlinear policy enables the RL-ROE to compensate for errors of the ROM , while still taking advantage of the imperfect knowledge of the dynamics . We describe how we frame the problem as a stationary Markov decision process in order to enable RL training , which is non-trivial since the RL-ROE must be able to estimate time-varying states . Finally , we show that the trained RL-ROE is able to outperform a Kalman filter designed using the same ROM , and displays robust estimation performance with respect to different reference trajectories and initial state estimates . The RL-ROE is the first application of reinforcement learning to state estimation for high-dimensional systems . Under review as a conference paper at ICLR 2022 2 PROBLEM FORMULATION . 2.1 SETUP . Consider the discrete-time nonlinear system given by zk+1 = f ( zk ) , ( 1a ) yk = Czk , ( 1b ) where zk ∈ Rn and yk ∈ Rp are respectively the state and measurement at time k , f : Rn → Rn is a time-invariant nonlinear map from current to next state , and C ∈ Rp×n is a linear map from state to measurement . In this study , we assume that the dynamics given in ( 1 ) are obtained from the numerical discretization of a nonlinear partial differential equation ( PDE ) , which typically requires a large number n of state dimensions . Note that we do not account for exogenous control inputs to the system , which will be studied in future extensions of the present work . 2.2 REDUCED-ORDER MODEL . Because the high dimensionality of ( 1 ) makes online prediction and control impractical , it is instead customary to formulate a reduced-order model ( ROM ) of the dynamics ( Rowley & Dawson , 2017 ) . First , one chooses a suitable linearly independent set of modes { u1 , . . . , ur } , where ui ∈ Rn , defining an r-dimensional subspace of Rn in which most of the dynamics is assumed to take place . Stacking these modes as columns of a matrix U ∈ Rn×r , one can then express zk ' Uxk , where the reduced-order state xk ∈ Rr represents the coordinates of zk in the subspace . Finally , one finds a ROM for the dynamics of xk , which is vastly cheaper to evolve than ( 1 ) when r n. There exist various ways to find an appropriate set of modes U and corresponding ROM for the dynamics of xk ( Taira et al. , 2017 ) . In this work , we employ the Dynamic Mode Decomposition ( DMD ) , a purely data-driven algorithm that has found wide applications in fields ranging from fluid dynamics to neuroscience ( Schmid , 2010 ; Kutz et al. , 2016 ) . Starting with a collection of snapshots Z = { z0 , . . . , zm } collected along a trajectory of ( 1a ) , the DMD seeks a best-fit linear model of the dynamics in the form of a matrixA ∈ Rn×n such that zk+1 ' Azk , and computes the modes U as the r leading principal component analysis ( PCA ) modes of Z . The transformation zk ' Uxk and the orthogonality of U then yield a linear discrete-time ROM of the form xk+1 = Arxk +wk , ( 2a ) yk = Crxk + vk , ( 2b ) whereAr = UTAU ∈ Rr×r andCr = CU ∈ Rp×r are the reduced-order state-transition and observation models , respectively . In order to account for the neglected PCA modes of Z as well as the unmodeled dynamics incurred by the linear approximation zk+1 ' Azk , we add ( unknown ) nonGaussian process noise wk and observation noise vk . Additional details regarding the calculation ofAr and U are provided in Appendix A . 2.3 REDUCED-ORDER ESTIMATOR . This paper uses reinforcement learning ( RL ) to solve the following estimation problem : given a sequence of measurements { y0 , · · · , yk } from a reference trajectory { z0 , · · · , zk } of ( 1 ) and knowing the ROM ( 2 ) defined byAr , Cr andU , we want to estimate the high-dimensional state zk at current time k. To this effect , we design a reduced-order estimator ( ROE ) of the form x̂k = Arx̂k−1 + ak , ( 3a ) ak ∼ πθ ( · |yk , x̂k−1 ) , ( 3b ) where x̂k is an estimate of the reduced-order state xk , and ak ∈ Rr is an action sampled from a stochastic policy πθ which depends on the current measurement yk and the previous state estimate x̂k−1 . The subscript θ denotes the set of parameters that defines the stochastic policy , whose goal is to minimize the mean square error E [ zk − ẑk ] over a range of reference trajectories and initial reduced-order state estimates . Here , ẑk = Ux̂k denotes the high-dimensional state estimate reconstructed from x̂k . A Kalman filter is a special case of such an estimator , for which the action in ( 3b ) is given by ak = Kk ( yk −CrArx̂k−1 ) , ( 4 ) Under review as a conference paper at ICLR 2022 with Kk ∈ Rr×p the optimal Kalman gain . Although the Kalman filter is optimal when the statetransition and observation models are known exactly , its performance suffers in the presence of unmodeled dynamics . In our case , such model errors are unavoidable due to the ROM ( 2 ) being an inherent approximation of the high-dimensional dynamics ( 1 ) , which motivates our adoption of the more general form ( 3b ) . This form retains the dependence of ak on yk and x̂k−1 but is more flexible thanks to the nonlinearity of the stochastic policy πθ , which we train with deep RL in an offline stage . The stochasticity of πθ forces the RL algorithm to explore different actions during the training process , in order to find eventually an optimal θ∗ such that E [ zk − ẑk ] is minimized for various reference trajectories and initial estimates . We call the estimator constructed and trained through this process an RL-trained ROE , or RL-ROE for short . Thus , the methodology we propose consists of two steps . In a first offline stage , a ROM of the form ( 2 ) is obtained using high-dimensional snapshots zk from a single trajectory of ( 1 ) . The RL-ROE ( 3 ) is then constructed based on this ROM , and its policyπθ is trained using high-dimensional snapshots zk from multiple reference trajectories of ( 1 ) . Finally , the trained RL-ROE is deployed online to track a reference trajectory of ( 1 ) . In the online stage , the RL-ROE only requires measurements yk from the reference trajectory , and gives an estimate ẑk = Ux̂k for the high-dimensional state . In summary , our contributions in this paper are two-fold : 1 . We propose RL-ROE , a reduced-order state estimator for high-dimensional nonlinear systems . The RL-ROE takes the form ( 3 ) , which combines two unique features : first , the state transition modelAr is a ROM of the high-dimensional dynamics ; second , the term ak that assimilates measurements is sampled from a stochastic policy πθ trained with RL . The training procedure for πθ , which involves a non-trivial reformulation of the time-varying tracking problem as a stationary Markov decision process , is described in Section 4 . 2 . The performance of the RL-ROE is compared in Section 5 with that of KF-ROE , a Kalman filter constructed from the same ROM . The comparison is performed in the context of the Burgers equation using a range of reference trajectories and initial state estimates . 3 RELATED WORK . Previous studies have already proposed designing state estimators using policies trained through reinforcement learning . Morimoto & Doya ( 2007 ) introduced an estimator of the form x̂k = f ( x̂k−1 ) + L ( x̂k−1 ) ( yk−1 −Cx̂k−1 ) , where f ( · ) is the state-transition model of the system , and the state-dependent filter gain matrix L ( x̂k−1 ) is defined using Gaussian basis functions whose parameters are learned through a variant of vanilla policy gradient . Their reward function , however , was calculated using the measurement error instead of the state estimate error , potentially limiting the performance of the trained estimator . Hu et al . ( 2020 ) proposed an estimator of the form x̂k = f ( x̂k−1 ) + L ( xk − x̂k ) ( yk − Cf ( x̂k−1 ) ) , where L ( xk − x̂k ) is approximated by neural networks trained with a modified Soft-Actor Critic algorithm ( Haarnoja et al. , 2018 ) . Although they derived convergence properties for the estimate error , the dependence of the filter gain L ( xk − x̂k ) on the reference state xk limits its practical application . A major difference between these past studies and our work is that they do not construct a ROM of the dynamics and only consider low-dimensional systems with four state variables at most , in comparison with the hundred or more state dimensions that our RL-ROE can handle . Therefore , RL-ROE represents the first application of reinforcement learning to state estimation for high-dimensional systems , which makes it applicable to systems governed by PDEs such as fluid flows .

Summary: This paper proposes using RL to train a policy to output correction terms in the reduced states of reduced order models in order to perform state estimation. In particular the approach is applied to DMD and Burgers equation, but could presumably be applied to other forms of ROMs and dynamical systems. The method was shown to perform well in experiments, and was found to be more robust than baseline methods in the presence of perturbations and noise.