Neural optimal feedback control with local learning rules

1 Introduction . The sensorimotor control system has exceptional abilities to perform fast and accurate movements in a variety of situations . To achieve such skillful control , this system faces two key challenges : ( i ) sensory stimuli are noisy , making estimation of current state of the system difficult , and ( ii ) sensory stimuli is often delayed , which if unaccounted , results in movements that are inaccurate and unstable [ 1 ] . Optimal Feedback Control ( OFC ) has been proposed as a solution to this control problem [ 2 , 3 ] . OFC approaches these problems by building an internal model of the system dynamics , and using 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . this internal model to generate control actions . OFC often employs Kalman filtering to optimally integrate the predictions of this internal model and the noisy/delayed sensory stimuli . Because of the power and flexibility of the OFC framework , biologically plausible neural architectures capable of building such internal models has been under active investigation . Specifically , earlier works used attractor dynamics implemented through a recurrent basis function network [ 4 ] or a line attractor network [ 5 ] to implement Kalman filters . Kalman filtering and control has also been implemented through different phases of estimation , system identification and control [ 6 ] , and more recently , using a particle filtering method for Kalman filtering [ 7 ] . Nonetheless , these works suffer from major limitations . Importantly , none of these considered that sensory feedback is delayed [ 4 , 6 , 5 , 7 , 8 ] , although it has been prominent in the original computational-level OFC proposal [ 2 ] , or merely considered the case of Kalman filtering , and not the combination of it with control [ 4 , 5 , 7 , 8 ] . These works also required knowledge of the noise covariances , either a priori [ 4 , 5 , 7 , 8 ] or obtained in a separate ‘ offline sensor ’ mode [ 6 ] . Moreover , many of these works lack biological plausibility and realism one would expect from a viable model of brain function [ 4 , 5 , 7 , 8 ] . Crucially , biological plausibility requires the network to operate online , ( i.e . receive a stream of noisy measurement data and process them on the fly ) , and also requires synaptic plasticity rules to be local ( i.e . learn using rules that only depend on variables represented in pre and postsynaptic neurons and/or on global neuromodulatory signals ) . Lastly , several of these models suffer from combinatorial explosion as the dimensionality of the input grows [ 4 , 5 ] , require running an inner loop until convergence at each time step [ 8 , 6 ] , or require separate learning and execution phases [ 6 ] , cf . Table 1 . We address these shortcomings and present a complete neural implementation of optimal feedback control , thus tackling an open issue in biological control [ 9 ] . In this model , which we call Bio-OFC , the state space , the prediction error [ 10 , 11 ] ( i.e. , the mismatch between the network ’ s internal prediction and delayed sensory feedback ) , and the control are represented by different neurons , cf . Fig . 1 . The network also receives scalar feedback related to the objective function , as a global signal , and utilizes this signal to update the synaptic connection according to policy gradient method [ 12 , 13 ] . To test the performance of our network , we simulate Bio-OFC in episodic ( finite horizon ) tasks ( e.g. , a discrete-time double integrator model , a hand reaching task [ 1 ] , and a simplified fly simulation ) . Summary of contributions : • We introduce Bio-OFC , a biologically plausible neural network that combines adaptive model based state discovery via adaptive Kalman filtering with a model free control agent . • Our implementation does not require knowledge of noise covariances nor the system dynamic , considers delayed sensory feedback , and has no separate learning/execution phases . • Our model-free control agent enables closed-loop control , thus avoiding the weight transport problem , a challenging problem even in non-biological control . [ 14 , 15 ] 2 Background . We review classical Kalman estimation and control in this section , using boldface lowercase/uppercase letters for vectors/matrices and I for the identity matrix . 2.1 Problem formulation . We model the environment as a linear dynamical system driven by control input and perturbed by Gaussian noise . The true state of the system x is hidden and all the animal has access to are the observations y that are assumed to be linear functions of the state corrupted by Gaussian noise . dynamics : xt+1 = Axt +But + vt ( 1 ) observation : yt = Cxt +wt ( 2 ) Here vt ∼ N ( 0 ; V ) and wt ∼ N ( 0 , W ) are independent Gaussian random variables and the initial state has a Gaussian prior distribution x0 ∼ N ( x̂0 , Σ0 ) . The goal is to estimate the latent state x̂ in order to design a control u that minimizes expected cost expected cost : J = E [ T∑ t=0 c ( xt , ut ) ] ( 3 ) control : ut = k ( x̂t ) = arg min J ( 4 ) where c ( xt , ut ) is the instantaneous cost associated with state xt and action ut . As the environment dynamics is not known to the animal a priori , the parametersA , B , C must be learned online . 2.2 Kalman estimation and control . The Kalman filter [ 16 , 2 ] is an estimator of the latent state x̂t ( and its variance ) via a weighted summation of the current observation and the prediction of the internal model based on prior measurements . However , in biologically realistic situations the sensory feedback is always delayed . This means that the control signal ut has to be issued , and thus the state xt estimated , before yt has been observed . The appropriately modified Kalman filter computes the posterior probability distribution of xt given observations yt−τ , ... , y0 where τ ≥ 1 is the delay . We start here with the case of τ = 1 for which the recursive updates of the mean x̂t and the variance Σt are x̂t+1 = Ax̂t +But +Lt ( yt −Cx̂t ) ( 5 ) Σt+1 = ( A−LtC ) ΣtA > + V ( 6 ) where Lt is known as the Kalman gain matrix which optimally combines the noisy observations yt with the internal model and is given by Lt = AΣtC > ( CΣtC > +W ) −1 . ( 7 ) The Kalman filter is optimal in the sense that it minimizes the mean-squared error E [ et > et ] with prediction error ( innovation ) et = yt −Cx̂t . The output feedback law , also known as policy , ( 4 ) simplifies if the cost J is quadratic in ut and xt : ut = −Kx̂t ( 8 ) and is known as linear-quadratic regulator ( LQR ) . The control gain , K , is found by solving a matrix Riccati equation ( cf . Supplementary Material ) . Linear policies have been successfully applied to a variety of control tasks [ 17 ] . 3 Neural network representation for optimal feedback control . 3.1 Inference . For now , let us assume that the system dynamics , the Kalman gain and the control gainA , B , −C , L , andK are constant and known . Then , the latent state x̂ can be obtained by the Kalman estimator Eq . ( 5 ) . This algorithm naturally maps onto the network in Fig . 1 and Supplementary Fig . S1A with neural populations representing x̂ , e : = y −Cx̂ , and u that are connected by synapses whose weights represent the elements of the matrices A , B , −C and L. The computation of the control variable u according to Eq . ( 8 ) can be implemented by synapses whose weights represent the elements of the matrix −K . If the sensory stimulus delay is τ = 1 , our network implements the Kalman prediction reviewed in the previous section . In case τ > 1 , the latent state must be recomputed throughout the delay period which requires a biologically implausible circuit ( cf . Supplementary Fig . S1C ) . To overcome this , we adapt an alternative solution from online control [ 18 , 19 ] . Specifically , we combine delayed measurements with the similarly delayed latent state estimation . Such inference can be performed by the network of the same architecture and adjusting the synaptic delay associated with matrix C to match with the sensory delay , in accordance with the following expression : x̂t+1 = Ax̂t +But +L ( yt+1−τ −Cx̂t+1−τ ) ︸ ︷︷ ︸ et ( 9 ) As we show in Results ( Fig . 3 ) this reduces predictive performance only modestly compared to the biologically unrealistic scheme . More details on the inference process and the temporal order in which our recurrent network performs the above steps is shown in Supplementary Fig . S1B . 3.2 Learning . 3.2.1 System identification and Kalman gain . Next , we turn our attention to the system identification/learning problem , which was not addressed by some of the previous proposals , cf . Table 1 . Given a sequence of observations { y0 , · · · , yT } , we use a least squares approach [ 20 ] to find the parameters that minimize the mean-square prediction error 1T ∑T t=0 e > t et . We perform this optimization in an online manner , that is at each time-step t + 1 , after making the delayed observation yt+1−τ , we update the parameter estimates Â , B̂ , Ĉ using steps that minimize e > t et , assuming that the state estimate and actions corresponding to prior observations ( e.g . x̂t−τ , ut−τ etc . ) are fixed . To obtain L we would like to avoid solving the Riccati equation ( 6 ) as it requires matrix operations difficult to implement in biology . So , we use the same optimization procedure to update the Kalman gain L. Using Eq . ( 9 ) and explicitly writing out the matrix/vector indices as superscripts yields the following stochastic gradient with respect toA : − ∂ ∂Aij 1 2 ∑ k ( ekt ) 2 = − ∑ k ekt ∂ekt ∂Aij = − ∑ k ekt ∂ ( ykt+1−τ − ∑ l C klx̂lt+1−τ ) ∂Aij = = ∑ k , l ektC kl ∂ ( ∑ mA lmx̂mt−τ + ∑ nB lnunt + ∑ p L lpept ) ∂Aij = = ∑ k , l , m ektC klδliδmj x̂mt−τ = ∑ k ektC kix̂jt−τ ( 10 ) Performing similar derivations for the other synaptic weights , our optimization procedure would rely on the following stochastic gradients : −∇A 12e > t et = C > etx̂ > t−τ −∇B 12e > t et = C > etu > t−τ ( 11 ) −∇L 12e > t et = C > ete > t−τ −∇C 12e > t et = etx̂ > t+1−τ ( 12 ) This yields a classical Hebbian rule between ( a memory trace of ) presynaptic activity x̂t+1−τ and postsynaptic activity et for weights C. However , it suggests non-local learning rules forA , B , L , which runs contrary to biological requirements . We can circumvent this problem by replacing C > with L , which corresponds to left-multiplication of the gradients with a positive definite matrix ( see Supplementary Material Sec . 4 ) . This still decreases the mean-square prediction error under some mild initialization constraints on C and L and yields local plasticity rules , cf . Fig . 1 , ∆Ât ∝ Let x̂ > t−τ ( 13 ) ∆Lt ∝ Let e > t−τ ( 14 ) ∆B̂t ∝ Let u > t−τ ( 15 ) ∆Ĉt ∝ et x̂ > t+1−τ ( 16 ) where the input current Let is locally available at neurons representing x̂ . The first three rules are local , but non-Hebbian , capturing correlations between presynaptic activity x̂t−τ , ut−τ , et−τ and postsynaptic current Let . Note that these updates do not require knowledge of the noise covariances V andW , an advantage over previous work , cf . Table 1 .

Summary: The authors propose a model of optimal feedback control for linear dynamical systems that can be realized via biologically plausible learning rules. The authors recognize the importance of integrating both prediction and control and achieve this by integrating two modules -- a model-based state estimation module that combines noisy measurements with internally generated predictions, and a model-free controller module that selects actions based on the estimated state. They demonstrate that the state estimator module can learn the parameters of the system using Hebbian-like weight updates provided the model parameters are initialized appropriately. They also demonstrate that the effect of sensory delay on the proposed model is comparable to what would be expected from an optimal model. The control module uses a standard policy-gradient method to learn a mapping from states to actions. The authors also evaluate their model in a neuroscience-inspired reaching task as well as a new flying task designed in OpenAI gym.