Neural Hybrid Automata: Learning Dynamics With Multiple Modes and Stochastic Transitions

1 Introduction . Behaviors emerging from the interaction of continuous and discrete–time dynamics in the presence of uncertainty are described through the language of stochastic hybrid systems ( SHSs ) . Such discrete events can bring along abrupt changes in the state , and in complex multi–mode systems , may also cause a switch between system modes , and corresponding underlying continuous dynamics [ 1 ] . Communication networks [ 2 ] , [ 3 ] , where changes in communication protocol can happen at certain levels of traffic congestion , and biological systems [ 4 ] – [ 6 ] are example domains where the SHS modeling paradigm has proven fruitful . Data–driven identification and learning of hybrid systems are known to be challenging due to the entanglement of continuous flows and discrete events [ 7 ] ; finding a generally applicable technique remains an open problem , particularly in the common scenarios where no a priori knowledge on the number and type of system modes is given . The aim of this work is to apply continuous neural models [ 8 ] – [ 10 ] to the learning SHSs . We introduce a compact descriptive language for this task , decomposing the system into a set of core primitives . Prior work is integrated into the framework , highlighting in the process limiting assumptions and areas of further improvement . ∗Equal contribution . Author order was decided by flipping a coin . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . To address the shortcomings of existing techniques , we introduce Neural Hybrid Automata ( NHA ) as a general procedure designed to enable learning and simulation of SHSs from data . NHAs are comprised of three components : a dynamics module , taking the form of an neural differential equation ( NDE ) [ 8 ] – [ 10 ] capable of approximating a different vector field for each mode , a discrete latent state tracking the internal mode of the target system , and an event module determining the time to next event . In particular , our approach does not require prior knowledge on the number of modes . The synergy among NHA components ensures a broader range of applicability compared to previous attempts , which in example do not directly tackle multi–mode hybrid systems [ 11 ] – [ 14 ] . NHAs are shown to enable mode recovery and learning of systems with stochastic transitions , with additional applications in end–to–end learning of hierarchical robot controllers . 2 Background . We introduce required background on the formalism of stochastic hybrid systems ( SHSs ) and event handling for their numerical simulation . We then provide further contextualization on previous approaches , introducing in the process a unified language for SHS learning tasks . 2.1 Stochastic Hybrid Systems . A stochastic hybrid system ( SHS ) [ 2 ] , [ 15 ] is a right-continuous stochastic process Xt taking values in X ⊆ Rnx with a latent mode process Zt conditioning the dynamics of Xt , where t ≥ 0 . Zt is another right-continuous stochastic process that takes values in a finite set M of size m. In this context , the set M contains identifiers of internal system modes . An event is defined as either a mode switch or a state discontinuity ( a jump in Xt ) , which can in some cases occur simultaneously . We refer to times at which events z → z′ occur as random variables tk ∈ T , with associated intensity functions [ 16 ] λz→z′ ( t|Ht ) ≥ 0. where Ht : = { tk ∈ T : tk < t } is the history of event times . Intensity , as defined in the classical temporal point process ( TPP ) sense , can be interpreted as the expected number of events z → z′ within the time interval [ t , t+ dt ] . The dynamics for Xt when Zt = z is defined by flow dynamics : ẋt = Fz ( t , xt ) . ( z , t , xt ) ∈M× T× X When a jump event z → z′ is triggered , Xt can instantaneously jump according to jump dynamics : x+t = ψz→z′ ( t , xt ) . ( z , z ′ , t , xt ) ∈M2 × T× X Jump maps ψz→z′ and intensities λz→z′ describe the behavior during events z → z′ . 2.2 Event Handling for Hybrid Systems . Following [ 17 ] , to enable forward simulation of SHSs , a convenient mathematical representation of an event is a function g : T× X→ R which nullifies only at any event time t∗ , thus providing the differential equation integration algorithm with a termination or restart condition i.e . t∗ is an event ⇔ g ( t∗ , xt∗ ) = 0 . ( 2.1 ) The particular form of g induces a jump set on D ⊂ T × X , D : = { t , x : g ( t , xt ) = 0 } , and determines transitions from roots of g to regions of the state–space X where g 6= 0 . Notably , this construct enables utilization of root finding methods [ 18 ] in a neighborhood of t∗ to accurately zero in on the event time . The same simulation technique can be extended to the many jump sets case typical of multi–mode systems , by equipping the condition function with identifiers z , z′ ( gz→z′ ) which induces jump sets Dz→z′ . Simulating stochastic events While the event function approach appears to be limited to the deterministic setting , it also subsumes stochastic events whose aleatoric uncertainty is encoded by an intensity λ ( t|Ht ) [ 1 ] , [ 13 ] . Without loss of generality let us consider a single intensity function which is henceforth denoted as λ∗t : = λ ( t|Ht ) . Recalling that the cumulative distribution function ( CDF ) of inter–event times is 1− exp { − ∫ t∗ t0 λ∗tdt } , standard inverse transform sampling [ 1 ] , [ 19 ] yields t∗ : 0 = s− log ∫ t∗ t0 λ∗tdt , s ∼ Uniform ( 0 , 1 ) ( 2.2 ) as a special case of ( 2.1 ) . Approaches developed for learning TPPs , including in the context of Neural ODEs [ 8 ] , [ 11 ] , [ 20 ] , introduce a parametric formulation for the intensity λ∗θ and optimize via direct TPP likelihood objectives . The integral is in general intractable , thus these methods require either a numerical approximation or the augmentation of additional states to compute it alongside the ODE . 2.3 Core primitives for SHS learning . At minimum , a learning model for SHSs necessitates several modules , each mirroring an element of the formulation in 2.1 . More specifically : i. Dynamics module , to approximate continuous dynamics Fz conditioned on mode z. ii . Discrete latent selector , to identify at each event time the latent mode z of the system . iii . Event module , to determine when events happen , and how state x and latent state z are updated after the transition . Prior work considers specific instantiations of SHSs , leading to simplifying choices for each module defined above . In example , switching systems without jumps , where iii . does not require jumps [ 21 ] , single mode systems , where ii . is not required and iii . does not need to determine latent mode transitions [ 11 ] , [ 13 ] , [ 14 ] , systems with known dynamics , where i. is not trained [ 13 ] , or systems with only deterministic events [ 7 ] . We note some of these works suffer from more than a single of these limitations , including additional ones consequence of specific model choices . Direct parametrization of the intensity , while a reasonable choice for single mode systems , requires state augmentation scaling in the worst case as ( m 2 ) [ 22 ] for a SHS with m modes . More importantly , training parameters θ for such direct approaches is affected by the accuracy of the numerical method employed for the solution of the integral in ( 2.2 ) . In this work , we introduce a modelling framework SHSs that does not rely on any of the simplifying assumption on i. , ii . and iii . outlined above . 3 Neural Hybrid Automata . We introduce Neural Hybrid Automata ( NHA ) , a model for learning of SHSs . A NHA is comprised of a dynamics module , a discrete latent state and an event predictor . A general overview of a NHA is depicted in Figure 1 . We start with a description of each module and their interconnections , followed by a step–by–step procedure for NHA training . Event module Intensity–free parametrizations for stochastic mode transitions allows NHAs to sample next event times without solving integrals ∫ λ∗tdt for all target modes z ′ reachable from current z . From event k at time tk , NHAs determine next event times tk+1 through a conditional normalizing flow modeling , for each possible pair of ( z , z′ ) , the density of corresponding inter–event times τkz→z′ = tk+1 − tk , tk ∈ Tz→z′ . Let the intensity be a simple timer i.e λ̇ = 1 . Further , let pz→z′ be the parametrized conditional density obtained by the normalizing flow and let T ( z , z′ , tk ) be a collection of conditional samples from pz→z′ ( one for each pair z , z′ ) , i.e . T ( z , z′ , tk ) = { τkz→z′ ∼ pz→z′ ( τ |Htk ) } z , z′∈Z . Using ( 2.2 ) , we can thus sample an event given the current mode z and the previous event time tk as : tk+1 = tk +min z′∈Z T ( z , z′ , tk ) . ( 3.1 ) Note that the next mode z′ after the event is simultaneously obtained as z′ = argminT ( z , z′ , tk ) . Sampling strategy ( 3.1 ) , differently from ( 2.2 ) , relies on the normalizing flow to explicitly model the density rather than defining it implicitly through ∫ λ∗tdt . When event time tk is reached , a parametrized jump map conditioned on ( z , z′ ) is applied to the state x+ = ψz→z′ ( tk , x ) . Normalizing flow pz→z′ and jump map ψz→z′ together define the full event module of an NHA . It should be noted that ( 3.1 ) always samples the quickest–to–occur event from the normalizing flow , which implies that no other event occurs between tk and tk+1 . While the history can be compressed into a fixed–length vector following [ 23 ] through application of sequence models e , g . RNNs , we note that for hybrid systems equipped with deterministic events , providing ( xtk , tk ) as conditioning inputs for pz→z′ ( τ ) is sufficient since ODE solutions with deterministic transitions are uniquely determined by the initial condition . Finally , deterministic events are a special case of stochastic events [ 2 ] that can be well–represented with a Dirac δ function , of which the normalizing flow learns a smooth approximation with continuous support . Dynamics module To enable approximation of different mode–dependent vector fields , we parametrize the flow map Fz ( t , xt ) of a SHS as a data–controlled neural ordinary differential equation ( Neural ODE ) [ 9 ] with parameters ω , driven between each pair of event times tk , tk+1 by discrete latent mode z ẋt = Fz ( t , xt , ω ) t ∈ [ tk , tk+1 ) ( 3.2 ) Finiteness of admissible values in the latent mode state i.e . m ensures F is capable of approximating a finite number of different vector fields , one for each mode . In particular , we consider one–hot representations for latents z ∈ Rm . In batched data settings , ( 3.2 ) can be integrated in parallel across nb batches of initial conditions xtk ∈ Rnb×nx with different modes , provided the latent is also batched z ∈ Rnb×m . The combination of a given dynamics and event module , applied in turn as depicted in Figure 1 , enables simulation of trajectories of a SHS . We now describe their training procedure .

Summary: This paper introduces a complete pipeline for controlling with multi-modal transition and event handling. The method shown in the paper is based on normalizing flow and neural differential equations. In the experiments, the paper shows the efficacy in a few controlling cases. The method first segments the trajectory into multiple partitions, then performs the reconstruction using normalizing flow.