Deep Reinforcement Learning Data Collection for Bayesian Inference of Hidden Markov Models

Abstract

Hidden Markov Models (HMMs) are a powerful class of dynamical models for representing complex systems that are partially observed through sensory data. Existing data collection methods for HMMs, typically based on active learning or heuristic approaches, face challenges in terms of efficiency in stochastic domains with costly data. This paper introduces a Bayesian lookahead data collection method for inferring HMMs with finite state and parameter spaces. The method optimizes data collection under uncertainty using a belief state that captures the joint distribution of system states and models. Unlike traditional approaches that prioritize short-term gains, this policy accounts for the long-term impact of data collection decisions to improve inference performance over time. We develop a deep reinforcement learning policy that approximates the optimal Bayesian solution by simulating system trajectories offline. This pre-trained policy can be executed in real-time, dynamically adapting to new conditions as data is collected. The proposed framework supports a wide range of inference objectives, including point-based, distribution-based, and causal inference. Experimental results across three distinct systems demonstrate significant improvements in inference accuracy and robustness, showcasing the effectiveness of the approach in uncertain and data-limited environments.

I. Introduction

Hidden Markov Models (HMMs) are widely recognized statistical models employed in numerous domains, including network security, ecology, transportation, image processing, sentiment analysis, autonomous driving, health monitoring, and disaster evaluation [1]–[8]. Inference of HMMs is crucial for understanding and predicting sequential data patterns, offering insights into underlying processes masked by noisy or incomplete observations.

Several inference techniques have been developed for modeling and constructing HMMs. A stochastic variational inference algorithm for learning the parameters of HMMs is introduced in [9], where its efficacy is demonstrated on synthetic experiments and a large genomics dataset. Wiuf et al. [10] suggest a novel reaction network scheme, the Baum-Welch reaction network, which effectively learns parameters for HMMs and demonstrates that its expectation and maximization steps converge exponentially fast. An online expectation-maximization algorithm is proposed in [11] for enabling realtime parameter estimation in HMMs. Markov Chain Monte Carlo (MCMC) methods [12], [13] are also a popular class of methods used for Bayesian inference of HMMs. Moreover, combining approximate Bayesian computation with MCMC techniques is proposed in [14] to improve computational efficiency, especially when dealing with intractable likelihoods and poor prior knowledge about model parameters. The existing inference methods can be categorized broadly into point-based, such as a maximum likelihood or maximum a posteriori inference, and distribution-based, such as Bayesian inference.

The accuracy of the inference methods is limited to the capacity in which the available data represent the underlying complex mechanism of systems. One of the main challenges in the inference of HMMs is the non-identifiability issue [15], which refers to cases where different sets of model parameters can produce the observed data. This can lead to difficulties in model interpretation and inference. The non-identifiability issue often happens in systems under normal conditions or tightly controlled systems, where available data (though large) only represent a small portion of the state space. Therefore, collecting additional data to properly learn the underlying mechanism of systems is essential. Data collection can be achieved by influencing the system state transitions, such as through excitation or perturbation into the system states, intelligently selecting the sensors to monitor the system, or any other actions that may lead to obtaining different information from the systems.

Several data collection approaches have been developed to enhance the inference of HMMs and overcome the non-identifiability issue [16], [17]. For instance, a method for inferring human behaviors in human-robot collaboration tasks is proposed in [18] by collecting and analyzing data such as gestures, eye movements, and head motions, utilizing a Bayesian learning approach within a Partially Observed Markov Decision Process (POMDP) framework. Also, a Bayesian method for learning POMDP observation parameters in robot interaction management systems is introduced in [19], allowing gradual adjustment of the robot’s behavior based on user interactions and an oracle’s optimal policy. These approaches, however, are application-specific and are not applicable to general HMMs. Sensor selection or scheduling approaches are another class of data collection techniques for HMMs [20]–[24]. However, these methods can only optimize the allocation of monitoring resources rather than changing the system transitions to collect informative data. Also, the majority of these methods are developed to enhance state estimation performance and rely on the availability of full knowledge of HMMs. Experimental design techniques are the most relevant approaches for data collection for inference of HMMs. This class of methods aims to sequentially and greedily collect data to reduce model uncertainty (e.g., entropy or related measures) [25]–[27]. For instance in [28], expected information gain is utilized to optimize data selection in neurobiological research, aiming to improve inference with fewer data and potentially enhance the design of clinical trials. Also, a framework for active learning in latent variable models is proposed in [29], leveraging entropy measures to reduce data requirements for fitting mixtures of linear regressions and HMMs with generalized linear model observations. Despite the relative success of these techniques, they become inefficient in domains with a huge amount of uncertainty in state, measurement, and system models, as well as domains with limited data collection resources. This comes from the incapability of these methods to account for the long-term impacts of decisions on system uncertainty and inference accuracy.

In recent years, Reinforcement Learning (RL) techniques have been extensively applied to the inference and control of complex dynamical systems [30], [31]. For example, a fuzzy-model-based approach has been proposed in [32] to optimize nonlinear Markov jump singularly perturbed systems using RL. In [33], an inference method is developed that incorporates expert knowledge into the inference process for HMMs. The expert is modeled as a sub-optimal RL agent, and the method aims to extract its knowledge for inference rather than collecting new data. Furthermore, an RL-based data collection method for offline meta-RL is presented in [34]. This method addresses the identifiability issue and sample-efficient interactions across domains modeled by a Markov Decision Process (MDP). However, this approach focuses on task adaptation in offline meta-RL and does not address the challenge of identifying unknown model parameters. Moreover, in some complex systems, such as cyber-physical and biological systems, non-Markovian behavior can emerge, where the future state depends not only on the current state but also on past states. Recent studies have proposed fractional-order derivatives and master equations to capture these long-range dependencies [35]–[37]. In such cases, data collection strategies may need to distinguish whether Markovian or non-Markovian dynamics are more appropriate for modeling these systems.

This paper proposes a deep reinforcement learning Bayesian data collection policy for inferring unknown HMMs. A Bayesian representation of unknown HMMs is provided through a belief state, which keeps track of the joint distribution of the possible system states and models. The optimal Bayesian data collection process is formulated in belief space, where the belief transition propagates the system and data uncertainty. The offline training is achieved through simulating belief trajectories for pre-training a deep reinforcement learning policy without the need for real data or prior knowledge of the system model. The offline-trained Bayesian policy is executable in real-time, according to the latest data collected from the real systems. Some of the key features of the proposed policy are as follows:

• Approximating the optimal Bayesian lookahead data collection policy, which accounts for uncertainty in the system and data and considers the long-term impact of decisions on the HMM’s inference performance.

• Learnable through offline trajectories synthetically generated from the uncertain model; the trained policy is executable in real-time during data collection from the real system and dynamically adapts to the system as new data is collected.

• Applicability to a wide range of inference criteria, including point-based criteria such as a maximum likelihood or maximum a posteriori, distribution-based criteria, and causal inference.

Unlike existing data collection methods, such as active learning [28] and maximum a posteriori approaches [38], which prioritize short-term or greedy objectives, our approach formulates data collection as a lookahead policy optimization problem. This allows us to consider the future impact of data collection actions, leading to more informed decisions that improve inference accuracy over time. Additionally, our method’s ability to operate offline, using synthetic belief trajectories during training, significantly reduces the amount of real-time data required, making the learning process more efficient and scalable across different domains.

The effectiveness of the proposed data collection approach is evaluated across three distinct systems. These include a general HMM with unknown state and measurement processes, a gene regulatory network characterized by an unknown state process, and a communication network featuring an unknown measurement process.

Fig. 1 illustrates a schematic diagram for the proposed data collection method. The blue arrows denote the already collected data, and the red dashed arrows indicate the future undecided decisions during data collection. The less visible network model in the early steps expresses the uncertainty about the network parameters. Given the already available data, the proposed data collection policy sequentially selects actions that can generate the most valuable data for quick and accurate inference.

Fig. 1:

The schematic diagram of the proposed data collection method. The proposed method takes into account the data collected in the past and makes decisions about the next data to collect to ensure the quickest and most accurate inference.

The rest of the paper is organized as follows: Section II contains a background regarding finite-state HMMs and discusses the inference of unknown HMMs, followed by the formulation of data collection in HMMs. Further, Section III presents the optimal Bayesian lookahead data collection policy for unknown HMMs. The proposed deep reinforcement learning Bayesian data collection policy is explained in detail in Section IV. Section V includes the numerical results for three different HMM models, demonstrating the efficacy of the proposed method. Finally, Section VI consists of the conclusion and a discussion of future work.

II. Background and Preliminaries

A. Finite-State Hidden Markov Models

HMMs characterize the dynamical behavior of systems observed through noisy time-series data. The schematic diagram of HMMs with external inputs is shown in Fig. 2. The time-varying behavior of systems is represented through the state process, where ${\mathbf{x}}_k$ represents the system state at time step $k$. This paper assumes the state takes its value from a finite set $𝒳=\left\{{\mathbf{x}}^1,\dots ,{\mathbf{x}}^n\right\}$. The states are observed indirectly through observations, where ${\mathbf{y}}_k$ represents the measurement at time step $k$, taking in an arbitrary space. The following two processes represent the finite-state HMMs:

$$\begin{gathered} {\mathbf{x}}_k~p\left(.\mid {\mathbf{x}}_{k-1},{\mathbf{a}}_{k-1},\theta \right)\left(\text{state}\text{process}\right) \\ {\mathbf{y}}_k~p\left(.\mid {\mathbf{x}}_k,{\mathbf{a}}_{k-1},\theta \right)\left(\text{measurement}\text{process}\right) \end{gathered}$$ (1)

for $k=\mathrm{1,2},\dots$; where $p(\cdot )$ denotes probability density or probability mass functions, ${\mathbf{a}}_{k-1}\in 𝒜$ is the external input to the system at time $k-1$, and $\theta$ represents the unknown elements of the two processes. The state and measurement processes are stochastic; the state process represents the system dynamics, whereas the measurement process characterizes the data generated from sensors/technology. The parameter $\theta$ takes its value from a finite parameter space $\Theta =\left\{{\theta }^1,\dots ,{\theta }^r\right\}$, where each $\theta \in \Theta$ represents a possible HMM. Therefore, the terms parameter and model are used interchangeably throughout this paper.

Fig. 2:

The schematic diagram of HMMs with external inputs.

B. Formulation of Optimal Inference for HMMs

1). Inference of Unknown HMMs:

The inference in HMMs consists of estimating the unknown parameters of the model in (1) represented by $\theta$ according to the available data. Let $\left\{\left({\mathbf{a}}_0,{\mathbf{y}}_1\right),\left({\mathbf{a}}_1,{\mathbf{y}}_2\right),\dots ,\left({\mathbf{a}}_{k-1},{\mathbf{y}}_k\right)\right\}$ be the available timeseries data, where ${\mathbf{a}}_{i-1}$ is the input or action at time step $i-1$, which has led to the measurement ${\mathbf{y}}_i$ at time step $i$. Depending on the application, the inference objective can be expressed in terms of point-based or distribution-based criteria, as described below.

The point-based inference aims to find a single best estimate of the unknown parameters. Two well-known point-based criteria are maximum likelihood (ML) and maximum a posteriori (MAP), which can be expressed as [39], [40]:

$$\begin{gathered} {\widehat{\theta }}_k^{\text{ML}}=\underset{\theta \in \text{Θ}}{\text{argmax}}p\left(y_{1:k}\mid a_{0:k-1},\theta \right), \\ {\widehat{\theta }}_k^{\text{MAP}}=\underset{\theta \in \text{Θ}}{\text{argmax}}p\left(\theta \mid y_{1:k},a_{0:k-1}\right), \end{gathered}$$ (2)

where ${\hat{\theta }}_k^{\text{ML}}$ and ${\hat{\theta }}_k^{\text{MAP}}$ represent the parameters with the highest likelihood and posterior probability, respectively. Note that MAP and ML become the same under a uniform prior distribution of the unknown parameters.

The confidence of the MAP inference ${\hat{\theta }}_k^{\text{MAP}}$ in (2) can be assessed using the probability of the MAP model as:

$$C_k=p\left({\hat{\theta }}_k^{\text{MAP}}\mid {\mathbf{y}}_{1:k},{\mathbf{a}}_{0:k-1}\right)=\underset{\theta \in \Theta }{\text{max}}p\left(\theta \mid {\mathbf{y}}_{1:k},{\mathbf{a}}_{0:k-1}\right),$$ (3)

where $1/r\le C_k\le 1$ represents the highest posterior probability of the models. The closer $C_k$ to 1, the higher the confidence of the MAP inference is. The confidence rate of $1/r$ corresponds to the poorest performance of the MAP inference. Note that $C_k=1/r$ is a special case where all the models have the same posterior probability given the available data; thus, the true model is not distinguishable from the others.

For distribution-based inference, the objective is to find the posterior distribution of the possible system models, as opposed to a model with the highest likelihood or posterior probability. Given the available data up to time step $k$, this can be expressed as [39]:

$$\begin{gathered} p_k=p\left(\Theta \mid y_{1:k},a_{0:k-1}\right) \\ ={\left[p\left({\theta }^1\mid y_{1:k},a_{0:k-1}\right),\dots ,p\left({\theta }^r\mid y_{1:k},a_{0:k-1}\right)\right]}^T. \end{gathered}$$ (4)

Similar to the expected confidence in (3) for point-based inference, the confidence of the distribution-based inference in (4) can be assessed using the following entropy measure [41]:

$$D_k=-\mathscr{H}\left[p\left(\Theta \mid {\mathbf{y}}_{1:k},{\mathbf{a}}_{0:k-1}\right)\right]=\sum _{i=1}^r{p}_k\left(i\right)\text{log}{p}_k\left(i\right),$$ (5)

where $\mathscr{H}[.]$ represents the Shannon entropy and $\text{log}(1/r)\le D_k\le 0$. The values of $D_k$ close to zero represent cases with posterior distributions peaked over a single model and, consequently, the most confident inference. On the other hand, the values close to $D_k=\text{log}(1/r)$ indicate the poorest performance of distribution-based inference.

2). Formulation of Data Collection in HMMs:

Data can be costly, in terms of time and/or money, in many practical domains, and it is often challenging or impossible to rely on only the available data for accurate inference of complex systems. This is due to the fact that the possible models in $\Theta$ often cannot be easily distinguished from each other through data collected from the system in normal conditions. Therefore, rather than collecting more data, targeted data should be collected to ensure the desired inference objective. Data collection in dynamic systems is challenging, as targeted actions should be made to intelligently change the system dynamics and generate data that reveal the true system model.

Depending on the inference criteria, the expected accuracy of the inference process can be assessed through the measures defined for the point-based and distribution-based inference in (3) and (5). Thus, given the huge cost of data collection, the goal is to select the best sequence of actions that leads to measurements under which the expected accuracy/confidence of inference is maximized. This can be expressed as:

$${\mathbf{a}}_{0:k-1}^{*}=\underset{{\mathbf{a}}_{0:k-1}\in {𝒜}^k}{\text{argmax}}\mathbb{E}\left[C_k\text{or}{D}_k\mid {\mathbf{a}}_{0:k-1}\right],$$ (6)

where the expectation is with respect to stochasticity in the unobserved measurements. Finding the optimal sequence of actions in (6) without full knowledge of the system model or access to large real data is impossible. Meanwhile, the batch selection of the action sequence makes it unsuitable for real-time and sequential data collection, as the latest data are not utilized to change the selections. Methods such as active learning [42] are developed to sequentially choose actions by accounting for possible changes in inference measures in a single step. While these methods make the computation of the expectation in (6) tractable, they often provide non-optimal solutions for data collection due to their greedy nature and often require huge computations in real-time for their selection process. In the next paragraphs, our proposed sequential and lookahead data collection approach to overcome these challenges is described.

III. Optimal Bayesian Lookahead Data Collection

A. Bayesian Modeling of Decision-Making in Unknown HMMs

Let $\left({\mathbf{a}}_{0:k-1},{\mathbf{y}}_{1:k}\right)$ be the available data collected up to time step $k$, where ${\mathbf{a}}_{1:k-1}$ and ${\mathbf{y}}_{1:k}$ represent the sequence of selected actions and observed measurements, respectively. Let $p_k\left(\theta \right)$ also be the posterior probability of the model $\theta \in \Theta$, given the collected data up to time $k$. We represent the information carried by the sequence of data and the posterior probability of models into a single vector called belief state. The belief state captures the joint probability distribution of the possible system states and system models using an $n\times r$ vector. The belief state at time $k$ can be represented as:

$${\mathbf{b}}_k=p\left({\mathbf{x}}_k=𝒳,\Theta \mid {\mathbf{y}}_{1:k},{\mathbf{a}}_{0:k-1}\right)=\left[\begin{array}{c}{\mathbf{b}}_k^{{\theta }^1}\\ ⋮\\ {\mathbf{b}}_k^{{\theta }^r}\end{array}\right],$$ (7)

where

$${\mathbf{b}}_k^{{\theta }^l}=\left[\begin{array}{c}p\left({\mathbf{x}}_k={\mathbf{x}}^1,{\theta }^l\mid {\mathbf{y}}_{1:k},{\mathbf{a}}_{0:k-1}\right)\\ ⋮\\ p\left({\mathbf{x}}_k={\mathbf{x}}^n,{\theta }^l\mid {\mathbf{y}}_{1:k},{\mathbf{a}}_{0:k-1}\right)\end{array}\right].$$ (8)

The belief state is a sufficient statistic for representing the history of information as a compact form in unknown HMMs. The initial belief ${\mathbf{b}}_0=p\left({\mathbf{x}}_0=𝒳,\Theta \right)$ demonstrates the initial joint probability of the system states and unknown parameters, which could be set as uniform if no prior information is available.

In this part, we describe that the belief transitions follow a Markov process and provide a recursive procedure to update the belief state as new action is taken and a new measurement is obtained. Let ${\mathbf{b}}_0,{\mathbf{b}}_1,\dots ,{\mathbf{b}}_k$ be the sequence of belief states, with ${\mathbf{b}}_k=p\left(𝒳,\Theta \mid {\mathbf{a}}_{0:k-1},{\mathbf{y}}_{1:k}\right)$ be the latest belief. The belief state after taking an action ${\mathbf{a}}_k$ and observing a new measurement ${\mathbf{y}}_{k+1}$ can be updated as:

$$\begin{gathered} b_{k+1}=p\left(x_{k+1}=𝒳,\Theta \mid b_0,\dots ,b_k,a_k,y_{k+1}\right) \\ =\frac{p\left(x_{k+1}=𝒳,\Theta ,y_{k+1}\mid b_0,\dots ,b_k,a_k\right)}{p\left(y_{k+1}\mid b_0,\dots ,b_k,a_k\right)} \\ \propto p\left(y_{k+1}\mid x_{k+1}=𝒳,a_k,\Theta \right)\times p\left(x_{k+1}=𝒳,\Theta \mid b_0,\dots ,b_k,a_k\right), \end{gathered}$$ (9)

where the numerator of the second line is expanded in the last line. The first term of the last line of (9) can be computed using the measurement process in (1). The second term can be computed for any ${\mathbf{x}}^i\in 𝒳$ and $\theta \in \Theta$ as:

$$\begin{gathered} p\left(x_{k+1}=x^i,\theta \mid b_0,\dots ,b_k,a_k\right) \\ =\sum _{j=1}^{n}p\left(x_{k+1}=x^i,x_k=x^j,\theta \mid b_0,\dots ,b_k,a_k\right) \\ =\sum _{j=1}^{n}p\left(x_{k+1}=x^i\mid x_k=x^j,a_k,\theta \right)p\left(x_k=x^j,\theta \mid b_0,\dots ,b_k\right) \\ =\sum _{j=1}^{n}p\left(x_{k+1}=x^i\mid x_k=x^j,a_k,\theta \right)b_k^{\theta }\left(j\right), \end{gathered}$$ (10)

where ${\mathbf{b}}_k^{\theta }\left(j\right)$ denotes the $j\text{th}$ element of ${\mathbf{b}}_k^{\theta }$ and $p\left({\mathbf{x}}_k={\mathbf{x}}^j,\theta \mid \right.{\mathbf{b}}_0,\dots ,{\mathbf{b}}_k)$ in the third line is equal to ${\mathbf{b}}_k^{\theta }\left(j\right)$ according to (8). Replacing (10) into (9) leads to the following expression for the numerator of (9):

$$b_{k+1}^{\theta }\left(i\right)\propto p\left(y_{k+1}\mid x_{k+1}=x^i,a_k,\theta \right)\times \sum _{j=1}^{n}p\left(x_{k+1}=x^i\mid x_k=x^j,a_k,\theta \right)b_k^{\theta }\left(j\right),$$ (11)

for $i=1,\dots ,n$ and $\theta \in \Theta$. Equations (9) and (11) enable a recursive update of the belief state and reflect the Markovian transition of the belief states.

B. MDP Formulation of Data Collection

This section provides the belief space and belief transition formulation within the context of the Markov Decision Process (MDP) for the data collection procedure. The belief state takes its value in a simplex of size $n\times r,{\Delta }_{n\times r}$, denoted by $ℬ$ and referred to as the belief space. If $\mathbf{b}$ is the current belief state, upon taking action $\mathbf{a}$, the next belief state before observing the next observation is stochastic, with the transition represented by:

$${\mathbf{b}}^{\prime }~p(.\mid \mathbf{b},\mathbf{a}).$$ (12)

The belief transition in (12) represents the probability of the next belief state after taking action $\mathbf{a}$ from the current belief state $\mathbf{b}$. Therefore, the stochasticity in the next belief state corresponds to probable data that could be collected/observed in the next time step given all the available information reflected in the current belief state. Letting $\mathbf{b}={\left[{\mathbf{b}}^{{\theta }^1},\cdots ,{\mathbf{b}}^{{\theta }^r}\right]}^T$ be the current belief state, the belief state after taking action $\mathbf{a}$ and observing a new measurement $\mathbf{y}$ can be written as ${\mathbf{b}}_{\mathbf{y},\mathbf{a}}={\left[{\mathbf{b}}_{\mathbf{y},\mathbf{a}}^{{\theta }^1},\cdots ,{\mathbf{b}}_{\mathbf{y},\mathbf{a}}^{{\theta }^r}\right]}^T$, where

$${\mathbf{b}}_{\mathbf{y},\mathbf{a}}^{\theta }\left(i\right)=\frac{p\left(\mathbf{y}\mid {\mathbf{x}}^i,\mathbf{a},\theta \right)\sum _{j=1}^{n}p\left({\mathbf{x}}^i\mid {\mathbf{x}}^j,\mathbf{a},\theta \right){\mathbf{b}}^{\theta }\left(j\right)}{\sum _{{\mathbf{x}}^{\prime }\in 𝒳}\sum _{{\theta }^{\prime }\in \Theta }p\left(\mathbf{y}\mid {\mathbf{x}}^{\prime },\mathbf{a},{\theta }^{\prime }\right)\sum _{j=1}^{n}p\left({\mathbf{x}}^{\prime }\mid {\mathbf{x}}^j,\mathbf{a},{\theta }^{\prime }\right){\mathbf{b}}^{{\theta }^{\prime }}\left(j\right)},$$ (13)

for $i=1,\dots ,n$ and any $\theta \in \Theta$ according to (9) and (11). The belief transition probabilities in (12) can be discrete or continuous distributions, depending on the size of the measurement space. Given that the measurement takes its value in a continuous space, the belief transition can be expressed as:

$$p\left({\mathbf{b}}^{\prime }\mid \mathbf{b},\mathbf{a}\right)={\int }_{𝒴}p\left(\mathbf{y}∣\mathbf{b},\mathbf{a}\right)I_{{\mathbf{b}}^{\prime }={\mathbf{b}}_{\mathbf{y},\mathbf{a}}}\mathrm{d}\mathbf{y},$$ (14)

where $𝒴$ represents the space of all potential measurements and $I_{{\mathbf{b}}^{\prime }={\mathbf{b}}_{\mathbf{y},\mathbf{a}}}$ is an indicator function which returns 1 if ${\mathbf{b}}^{\prime }={\mathbf{b}}_{\mathbf{y},\mathbf{a}}$ and 0 otherwise. Notice that for multi-variant measurements, a single integral can be replaced by multiple ones, and in scenarios with discrete/finite measurement outcomes, the integral in (14) should be replaced with summations. (14) shows that when an action is chosen for data collection at any particular belief state, multiple potential subsequent belief states exist before observing the next data. It is important to recognize that the potential subsequent belief states will differ depending on the chosen actions and collected data. This enables us to propagate uncertainty in the belief and evaluate which action choices could result in future beliefs that align better with our data collection goals.

We formulate the gain during the data collection as a reward function in the belief space. The reward function can be expressed as $R\left({\mathbf{b}}_k,{\mathbf{a}}_k,{\mathbf{b}}_{k+1}\right)$, which measures the immediate gain in the inference accuracy after taking action ${\mathbf{a}}_k$ in belief state ${\mathbf{b}}_k$ and moving to the next belief state ${\mathbf{b}}_{k+1}$. Various reward functions can be defined given the inference and data collection objectives. If the goal is enhancing the confidence rate of point-based inference in (2), the reward function can be formulated using (3) as:

$$\begin{gathered} R\left(b_k,a_k,b_{k+1}\right)=C_{k+1}-C_k \\ =\underset{\theta \in \Theta }{\text{max}}{p}_{k+1}\left(\theta \right)-\underset{\theta \in \Theta }{\text{max}}{p}_k\left(\theta \right) \\ =\underset{\theta \in \Theta }{\text{max}}{\Vert b_{k+1}^{\theta }\Vert }_1-\underset{\theta \in \Theta }{\text{max}}{\Vert b_k^{\theta }\Vert }_1, \end{gathered}$$ (15)

where $\Vert .{\Vert }_1$ is the L-1 norm, and $p_{k+1}$ and $p_k$ express the model probabilities at time steps $k+1$ and $k$, computable using belief states at time steps $k+1$ and $k$, respectively. A positive reward characterizes a single-step increase in the maximum posterior probability of models, and a reward of zero represents that the maximum posterior probability is unchanged after observing the next observation as a result of the last taken action.

For the distribution-based inference in (4), the reward function can be formulated using (5) as:

$$R\left(b_k,a_k,b_{k+1}\right)=D_{k+1}-D_k=-H\left[p_{k+1}\right]+H\left[p_k\right]=\sum _{\theta \in \Theta }{\Vert b_{k+1}^{\theta }\Vert }_1\log {\Vert b_{k+1}^{\theta }\Vert }_1-\sum _{\theta \in \Theta }{\Vert b_k^{\theta }\Vert }_1\log {\Vert b_k^{\theta }\Vert }_1,$$ (16)

where the positive values represent a reduction of uncertainty in the posterior distribution of models after the latest action. In the rest of the paper, (15) and (16) are referred to as the point-based reward and entropy reward respectively. Note that more complicated and time-dependent reward functions can be used instead of (15) and (16). For instance, the reward function could include the cost of actions in domains with different action costs and the causality of inference objectives in domains with specific data collection goals.

C. Optimal Stationary Bayesian Data Collection

The belief state provides the MDP representation of the unknown HMMs. The belief transition allows for the propagation of state, measurement, and model uncertainty, given all the available information. The belief state gets updated stochastically according to the belief transition upon taking any action, and a reward is observed according to the next belief state. We define a deterministic stationary data collection policy $\mu :ℬ\to 𝒜$, which associates an action to any given belief state. The optimal Bayesian lookahead data collection policy ${\mu }^{*}$ can be defined as:

$${\mu }^{*}\left(\mathbf{b}\right)=\underset{\mu }{\text{argmax}}\mathbb{E}\left[\sum _{t=0}^{\infty }{\gamma }^{t}R\left({\mathbf{b}}_t,{\mathbf{a}}_t,{\mathbf{b}}_{t+1}\right)\mid {\mathbf{b}}_0=\mathbf{b},{\mathbf{a}}_{0:\infty }~\mu \right],$$ (17)

for all $\mathbf{b}\in ℬ$; where $0<\gamma <1$ is a discount factor indicating the importance of early reward compared to the future ones. The optimal Bayesian lookahead policy ${\mu }^{*}\left(\mathbf{b}\right)$ prescribes the best action that could be taken at the belief state $\mathbf{b}$ to maximize the long-term performance of inference (i.e., achieving the highest expected accumulated reward). The term “Bayesian” comes from the expectation in (17), which is with respect to the future belief states generated from the belief transition. In other words, all probable next belief states propagating the uncertainty in state, measurement, and models from the current belief state $\mathbf{b}$ are taken into account for finding the optimal Bayesian data collection policy ${\mu }^{*}\left(\mathbf{b}\right)$.

The exact computation of the optimal Bayesian policy in (17) is not possible, primarily due to the continuity of the belief space $ℬ$. As previously noted, the belief space forms an $(n\times r)\text{-simplex}\left({\Delta }_{n\times r}\right)$, preventing the use of dynamic programming approaches, such as value iteration or policy iteration [43], for computing the optimal policy. In the subsequent section, the proposed deep reinforcement learning approach is introduced to approximate the optimal data collection policy.

IV. Deep Reinforcement Learning Bayesian Data Collection

The recent advances in deep reinforcement learning (RL) have enabled learning on much larger scales than before. In this paper, we develop a deep RL Bayesian data collection policy to approximate the solution to the optimization problem in (17). To approximate the expectation in (17), the proposed deep RL algorithm uses simulated samples from the belief transition, which propagates the system’s uncertainty without requiring data collection from real systems. The trained offline policy is executable for real-time and sequential data collection. In this section, the matrix-form implementation of the belief state is first described, followed by the proposed deep RL approach for learning the Bayesian data collection policy.

A. Matrix-form Implementation of Belief State

Let $\left\{{\mathbf{x}}^1,\dots ,{\mathbf{x}}^n\right\}$ and $\left\{{\theta }^1,\dots ,{\theta }^r\right\}$ be an arbitrary enumeration of the possible system states and models, respectively. We define the transition matrix of size $n\times n$ associated with the model $\theta \in \Theta$ under action $\mathbf{a}$ as:

$${\left(M^{\theta }\left(\mathbf{a}\right)\right)}_{ij}=p\left({\mathbf{x}}_{k+1}={\mathbf{x}}^i\mid {\mathbf{x}}_k={\mathbf{x}}^j,{\mathbf{a}}_k=\mathbf{a},\theta \right),$$ (18)

for $i,j=1,\dots ,n$. The observation process can be expressed using the update vector $T^{\theta }(\mathbf{y},\mathbf{a})$ as:

$${\left(T^{\theta }(\mathbf{y},\mathbf{a})\right)}_i=p\left({\mathbf{y}}_{k+1}=\mathbf{y}\mid {\mathbf{x}}_{k+1}={\mathbf{x}}^i,{\mathbf{a}}_k=\mathbf{a},\theta \right),$$ (19)

for $i=1,\dots ,n$. Letting $\mathbf{b}={\left[{\mathbf{b}}^{{\theta }^1},\cdots ,{\mathbf{b}}^{{\theta }^r}\right]}^T$ be the current belief state, the next belief state ${\mathbf{b}}_{\mathbf{y},\mathbf{a}}$ after taking action $\mathbf{a}$ and observing $\mathbf{y}$ can be computed according to (13). This can be expressed in terms of the transition matrix and update vector in (18) and (19) as:

$${\mathbf{b}}_{\mathbf{y},\mathbf{a}}=\frac{\left[\begin{array}{c}{T}^{{\theta }^1}(\mathbf{y},\mathbf{a})\circ \left(M^{{\theta }^1}\left(\mathbf{a}\right){\mathbf{b}}^{{\theta }^1}\right)\\ ⋮\\ T^{{\theta }^r}(\mathbf{y},\mathbf{a})\circ \left(M^{{\theta }^r}\left(\mathbf{a}\right){\mathbf{b}}^{{\theta }^r}\right)\end{array}\right]}{\sum _{m=1}^r{‖T^{{\theta }^m}(\mathbf{y},\mathbf{a})\circ \left(M^{{\theta }^m}\left(\mathbf{a}\right){\mathbf{b}}^{{\theta }^m}\right)‖}_1},$$ (20)

where “$\circ$” denotes Hadamard product (i.e., the componentwise multiplication of two vectors), and $\Vert \cdot {\Vert }_1$ denotes the absolute sum of a vector’s elements. The belief transition probability in (14) can also be expressed as:

$$p\left({\mathbf{b}}^{\prime }\mid \mathbf{b},\mathbf{a}\right)={\int }_{𝒴}\left(\sum _{m=1}^r{‖T^{{\theta }^m}(\mathbf{y},\mathbf{a})\circ \left(M^{{\theta }^m}\left(\mathbf{a}\right){\mathbf{b}}^{{\theta }^m}\right)‖}_1\right)I_{{\mathbf{b}}^{\prime }={\mathbf{b}}_{\mathbf{y},\mathbf{a}}}\mathrm{d}\mathbf{y}.$$ (21)

B. Deep Reinforcement Learning Bayesian Data Collection Policy

For an arbitrary data collection policy $\mu :ℬ\to 𝒜$ defined over the belief space, we define the Q-function or the expected discounted reward function at belief state $\mathbf{b}\in ℬ$ after taking action $\mathbf{a}\in 𝒜$ and following policy $\mu$ afterward as:

$$Q^{\mu }\left(\mathbf{b},\mathbf{a}\right)=\mathbb{E}\left[\sum _{t=0}^{\infty }{\gamma }^{t}R\left({\mathbf{b}}_t,{\mathbf{a}}_t,{\mathbf{b}}_{t+1}\right)\mid {\mathbf{b}}_0=\mathbf{b},{\mathbf{a}}_0=\mathbf{a},{\mathbf{a}}_{1:\infty }~\mu \right],$$ (22)

where the expectation is with respect to the stochasticity of the belief transition. The optimal Q-function, denoted by $Q^{*}≔Q^{{\mu }^{*}}$, provides the maximum expected return under the optimal policy ${\mu }^{*}$, where $Q^{{\mu }^{*}}\left(\mathbf{b},\mathbf{a}\right)=\underset{\mu }{\text{max}}{Q}^{\mu }(\mathbf{b},\mathbf{a})$, for all $\mathbf{b}\in ℬ$ and $\mathbf{a}\in 𝒜$.

Since the tabular representation of the Q-function is impossible over the continuous belief space, a deep reinforcement learning approach is employed for approximating the Q-function. This paper adopts the deep Q-network (DQN) approach [44] to learn the Bayesian data collection policy, aiming to approximate $Q^{{\mu }^{*}}(\mathbf{b},\mathbf{a})$ according to neural networks. Note that other deep RL methods can also be used to derive the proposed data collection policy. The DQN method approximates this Q-function using two feed-forward deep neural networks with the same architecture: the Q-network, $Q_{\mathbf{w}}$, and the target network, $Q_{{\mathbf{w}}^{-}}$. The input for both networks is the belief state, and the outputs are $Q_{\mathbf{w}}\left(\mathbf{b},{\mathbf{a}}^1\right),\dots ,Q_{\mathbf{w}}\left(\mathbf{b},{\mathbf{a}}^{\left|𝒜\right|}\right)$, each corresponding to an action. The initial weights for both the Q-network and the target-network are set randomly.

To train the neural networks, we utilize a fixed-size replay memory $𝒟$. This replay memory stores episodes containing belief states that are generated using actions selected by the epsilon-greedy policy [43]. When the replay memory reaches its capacity, the newest experience replaces the oldest. Each episode begins with an initial belief state, which, if unknown, can be randomly chosen from the belief space $ℬ$. At step $t$ of the episode, an action is chosen based on the epsilon-greedy policy defined using the Q-network $Q_{\mathrm{w}}$:

$${\mathbf{a}}_t~\left\{\begin{array}{ll}{\text{argmax}}_{\mathbf{a}\in 𝒜}{Q}_{\mathbf{w}}\left({\mathbf{b}}_t,\mathbf{a}\right)& \text{w.p.}1-ϵ\\ \text{random}\left\{{\mathbf{a}}^1,\dots ,{\mathbf{a}}^{\left|𝒜\right|}\right\}& \text{w.p.}ϵ\end{array},\right.$$ (23)

where $0\le ϵ\le 1$ is a parameter that determines the degree of exploration during the training.

After choosing ${\mathbf{a}}_t$, the subsequent belief state is sampled from $p\left(.\mid {\mathbf{b}}_t,{\mathbf{a}}_t\right)$. This sample can be obtained by initially sampling from the current belief state as follows:

$$\left({\mathbf{x}}_t,{\theta }_t\right)~\text{Cat}\left[{\mathbf{b}}_t^{{\theta }^1}\left(1\right),\dots ,{\mathbf{b}}_t^{{\theta }^1}\left(n\right),\dots ,{\mathbf{b}}_t^{{\theta }^r}\left(1\right),\dots ,{\mathbf{b}}_t^{{\theta }^r}\left(n\right)\right],$$ (24)

where “Cat” stands for categorical (discrete) distribution, and $\left({\mathbf{x}}_t,{\theta }_t\right)$ is a selected sample from the belief state ${\mathbf{b}}_t$, which is a realization of possible system state and model given the information up to time step $t$ (reflected in belief state ${\mathbf{b}}_t$). A sample representing the next measurement can be generated according to (${\mathrm{x}}_t,{\theta }_t$) and (1) as:

$${\mathbf{y}}_{t+1}~p\left(.\mid {\mathbf{x}}_{t+1},{\mathbf{a}}_t,{\theta }_t\right),\text{where}{\mathbf{x}}_{t+1}~p\left(.\mid {\mathbf{x}}_t,{\mathbf{a}}_t,{\theta }_t\right).$$ (25)

Then, one can compute ${\mathbf{b}}_{t+1}$ using ${\mathbf{a}}_t$ and ${\mathbf{y}}_{t+1}$ in (20). The reward function can also be computed according to the transition as: $r_{t+1}=R\left({\mathbf{b}}_t,{\mathbf{a}}_t,{\mathbf{b}}_{t+1}\right)$.

After reaching a pre-determined number of steps, $N_{\mathrm{Q}\_\text{Update}}$, the Q-network $Q_{\mathbf{w}}$ needs to be updated based on a minibatch of experiences of size $N_{\text{batch}}$, randomly chosen from the replay memory $𝒟$ as:

$$Z={\left\{\left({\mathbf{b}}_i,{\mathbf{a}}_i,{\mathbf{b}}_{i+1},r_{i+1}\right)\right\}}_{i=1}^{N_{\text{batch}}}~𝒟,$$ (26)

The target values for the Q-network update can be calculated for the minibatch samples as:

$$y_i=r_{i+1}+\gamma \underset{\mathbf{a}\in 𝒜}{\text{max}}{Q}_{{\mathbf{w}}^{-}}\left({\mathbf{b}}_{i+1},\mathbf{a}\right),$$ (27)

for $i=1,\dots ,N_{\text{batch}}$. Note that the target values are calculated utilizing the target-network $Q_{{\mathbf{w}}^{-}}$. With these values, the weights of the Q-network, represented by $\mathbf{w}$, can be updated as follows:

$$\mathbf{w}\leftarrow \mathbf{w}-\alpha {\nabla }_{\mathbf{w}}\left[\sum _{i=1}^{N_{\text{batch}}}{\left(y_i-Q_{\mathbf{w}}\left({\mathbf{b}}_i,{\mathbf{a}}_i\right)\right)}^2\right],$$ (28)

where $\alpha$ represents the learning rate, and the mean squared error serves as the loss function for updating the weights. The optimization described in (28) can be implemented using a stochastic gradient optimization method like Adam [45]. After updating $\mathbf{w}$, the weights of the target-network, denoted as ${\mathbf{w}}^{-}$, should be updated using the soft update method as:

$${\mathbf{w}}^{-}\leftarrow (1-\tau ){\mathbf{w}}^{-}+\tau \mathbf{w},$$ (29)

where $\tau$ is the soft update hyperparameter. This incremental adjustment of the target network weights toward the Q-network is a critical factor in ensuring the stability of the deep reinforcement learning process. Training can halt either when performance gains become marginal or a predetermined performance threshold is reached. After the offline training, the final Q-network can be used for approximating the optimal Bayesian data collection policy as ${\mu }^{*}\left(\mathbf{b}\right)\approx {\text{argmax}}_{\mathbf{a}\in 𝒜}{Q}_{\mathbf{w}}(\mathbf{b},\mathbf{a})$. This Bayesian policy associates an action to each sample in the belief space $\mathbf{b}\in ℬ$. This policy can be utilized in real-time according to the belief states computed based on data from the real system.

Fig. 3 presents a schematic of the proposed data collection framework, showing both the offline training and real-time execution phases. Our deep RL data collection policy approximates the solution to the optimization problem in (17) by using belief transition samples, eliminating the need for costly real-world data collection needed for training. These offline belief trajectories are generated to train our proposed deep RL data collection policy. Fig. 3(a) shows this process, where rewards from belief transitions are used to update the policy. In real-time, as illustrated in Fig. 3(b), the trained policy selects actions based on the current belief state, while newly collected data updates the belief state, reflecting the Bayesian nature of our approach.

Fig. 3:

Schematic diagram of the proposed data collection policy: a) Offline Training: belief trajectories and rewards are generated to train the deep reinforcement learning data collection policy, where $t$ refers to the iteration of the training process; b) Online/Real-Time Execution: the offline trained data collection policy is executable in real-time according to the obtained belief state, where $k$ indicates the $k\text{th}$ data collected from the real system.

The main computational complexity of the proposed method arises during the offline phase, which occurs before real-time execution. In this phase, simulated belief trajectories are used to train the deep RL policy. Although this training is computationally demanding due to neural network optimization and belief state updates, it is carried out offline, allowing the policy to be fully optimized before real-time use. Once trained, the policy can be applied to real-time data collection without further retraining, enabling efficient operation in practical applications. During real-time execution, the primary computational cost involves updating the belief state, with a time complexity of $O\left(r{n}^2\right)$, where $r$ is the number of possible system models and $n$ is the number of system states. This complexity arises from calculating transition probabilities and updating the joint belief state across all models and states, as detailed in (20).

V. Numerical Experiments

This section contains three sets of experiments investigating the performance of the proposed method. The results are represented in terms of data collection for enhancing the performance of point-based, distribution-based, and causal inference. The following three data collection methods are used for comparison with our proposed policy: the active learning data collection policy [28], the maximum a posteriori (MAP) data collection policy [38], and the Bayesian Offline Reinforcement Learning (BOReL) data collection policy [34]. Active learning and MAP focus on immediate data impact by minimizing uncertainty in a single time step, resulting in a primarily short-term, greedy approach. These methods are effective in scenarios where immediate system feedback is critical but are limited in managing long-term uncertainty, particularly in complex environments. In contrast, BOReL addresses offline meta-reinforcement learning, where the objective is to equip a meta-agent with strategies that generalize across new tasks by leveraging past task data. BOReL uses Bayesian reinforcement learning principles, incorporating task dynamics and environmental uncertainty to balance exploration and exploitation effectively. Through a belief-augmented neural network, BOReL generates a structured belief estimate from offline data, enabling adaptable exploration across diverse scenarios. Unlike BOReL, which addresses data and uncertainty separately, our proposed data collection policy considers the joint distribution of state and parameters, achieving exact Bayesian optimality.

Three HMMs are selected for the numerical experiments: 1) a general HMM featuring unknown state and measurement processes 2) a gene regulatory network (p53-MDM2) characterized by an unknown state process 3) Littman’s network, serving as a communication network example with an unknown measurement process. The hyperparameters for our approach, such as learning rate, exploration rate, minibatch size, and replay memory size, were selected through a combination of grid search and prior research recommendations to balance training stability and performance. Following are the final hyperparameter values: 3 hidden layers, 128 neurons in each hidden layer, learning rate $\alpha =0.0005$, replay memory size $\left|𝒟\right|=100,000$, minibatch size $N_{\text{batch}}=64$, discount factor $\gamma =0.95$, exploration rate $ϵ=0.1$, Q-network update frequency $N_{{\mathrm{Q}}_{-}\text{Update}}=4$, and soft update coefficient $\tau =0.001$. Moreover, the ReLU activation function is used between each of the two layers in the neural networks. Further, the training process is limited to a maximum of 100,000 episodes, at which point the training is stopped.

For each HMM system, the proposed data collection policy is trained offline using synthetically generated belief state trajectories rather than real-world data. These belief trajectories are created by simulating the transitions in each system under different actions, allowing us to train our proposed policy without the need for costly or time-consuming real-time data collection. This offline training phase enables learning effective policies that can be applied to real-time data collection scenarios. The trained policies are finally tested for real-time data collection in the real systems and their performance is compared with MAP, active learning, and BOReL policies. The results are derived from the average of 1000 trials, with $95\%$ confidence bounds also reported where possible.

A. A General HMM Problem with Unknown State and Measurement Processes

Our proposed Bayesian data collection method is first tested on an HMM problem consisting of four states. The state transition diagram of this system is visualized in Fig. 4. The system’s four states and three actions are represented as $𝒳=\left\{{\mathbf{x}}^1,{\mathbf{x}}^2,{\mathbf{x}}^3,{\mathbf{x}}^4\right\}$ and $𝒜=\left\{{\mathbf{a}}^1,{\mathbf{a}}^2,{\mathbf{a}}^3\right\}$. The unknown parts of the HMM include the state transition probabilities from state ${\mathbf{x}}^1$ to the other states under actions ${\mathbf{a}}^1$ and ${\mathbf{a}}^2$. This leads to the following parameters for the state transitions:

$$\psi =\left[{\psi }_1=\left(q_{11}^1,q_{12}^1,q_{13}^1,q_{14}^1\right),{\psi }_2=\left(q_{11}^2,q_{12}^2,q_{13}^2,q_{14}^2\right)\right],$$ (30)

where $q_{ij}^l=p\left({\mathbf{x}}_{k+1}={\mathbf{x}}^j\mid {\mathbf{x}}_k={\mathbf{x}}^i,{\mathbf{a}}_k={\mathbf{a}}^l\right)$. Each $q_{ij}^l$ is assumed to take its value from the set $\{\mathrm{0,1}/\mathrm{3,1}/\mathrm{2,1}\}$, subject to $q_{11}^l+q_{12}^l+q_{13}^l+q_{14}^l=1$, for $l=\mathrm{1,2}$. This results in 14 possibilities for each of the parameter vectors ${\psi }_1$ and ${\psi }_2$ as:

$${\psi }_i\in {\Psi }_i=\left\{\right(1/\mathrm{2,1}/\mathrm{2,0},0),(1/\mathrm{2,0},1/\mathrm{2,0}),(1/\mathrm{2,0},\mathrm{0,1}/2),(\mathrm{0,1}/\mathrm{2,1}/\mathrm{2,0}),(\mathrm{0,1}/\mathrm{2,0},1/2),(\mathrm{0,0},1/\mathrm{2,1}/2),(\mathrm{1,0},\mathrm{0,0}),(\mathrm{0,1},\mathrm{0,0}),(\mathrm{0,0},\mathrm{1,0}),(\mathrm{0,0},\mathrm{0,1}),(1/\mathrm{3,1}/\mathrm{3,1}/\mathrm{3,0}),(1/\mathrm{3,1}/\mathrm{3,0},1/3),(1/\mathrm{3,0},1/\mathrm{3,1}/3),(\mathrm{0,1}/\mathrm{3,1}/\mathrm{3,1}/3\left)\right\}$$ (31)

for $i=\mathrm{1,2}$. Therefore, the parameter set $\Psi ={\Psi }_1\times {\Psi }_2$, consisting of all unknown parameters of the state transition, has a size of 14 × 14 = 196, with the true parameters ${\psi }^{*}=\left[{\psi }_1^{*}=(\mathrm{0,1}/\mathrm{2,1}/\mathrm{2,0}),{\psi }_2^{*}=(\mathrm{0,1}/\mathrm{3,1}/\mathrm{3,1}/3)\right]$ according to Fig. 4. Notice that ${\psi }_1^{*}$ and ${\psi }_2^{*}$ correspond to the 4th and 14th element of the sets in (31) respectively.

Fig. 4:

Transition diagram of the HMM problem.

The system state is assumed to be partially observable through the following measurement process:

$${\mathbf{y}}_k=\left\{\begin{array}{ll}{\mathbf{x}}^{\prime }\in \left\{𝒳-{\mathbf{x}}_k\right\}& \text{with}\text{probability}1-\varphi \\ {\mathbf{x}}_k& \text{with}\text{probability}\varphi \end{array}.\right.$$ (32)

where the observation ${\mathbf{y}}_k$ represents the true underlying state ${\mathbf{x}}_k$ with a probability of $\varphi$ and a non-true state with a probability of $1-\varphi$. We also assume $\varphi$ be an unknown parameter of the system, where it takes its values in the measurement parameter set $\Phi =\left\{\mathrm{0.25,0.55,0.85}\right\}$ with the true parameter ${\varphi }^{*}=0.85$.

Given the unknown parts of the transition probabilities and measurement process, the entire parameter space $\Theta :\Psi \times \Phi$ includes 196 × 3 = 588 distinct elements. Given the state space size of 4, the belief state takes its value in a ${\Delta }_{4\times 588}$. This represents a large and continuous space for the belief state, underscoring the need for a systematic and efficient data collection strategy to accurately infer the unknown parameters.

Fig. 5 represents the average negative entropy achieved under different policies. The proposed Bayesian data collection method is trained using the reward in (16). One can see that the proposed data collection method reaches a higher average negative entropy in almost all the steps compared to the other methods. The next best performance belongs to the BOReL policy, followed by the active learning policy. The worst performance is for the MAP policy, with minimal change observed in the average negative entropy after reaching step 20. The results indicate the significance of accounting for the uncertainty in the future horizons for effective data collection during the inference process.

Fig. 5:

Performance comparison of different data collection policies for entropy reduction of the general HMM problem.

In the next experiment, we investigate the effectiveness of the proposed policy in inferring a subset of unknown parameters rather than the entire unknown parameters. This is referred to as causal inference since the data collection policy should solely focus on these subsets. Given that the entire parameter vector consists of $\theta =\left(\psi =\left[{\psi }_1,{\psi }_2\right],\varphi \right)$, three causal cases for inferring $\psi ,{\psi }_1$ and ${\psi }_2$ are considered. Specifically, the reward function for the causal inference of the parameter ${\psi }_1$, while other parameters ${\psi }_2,\varphi$ are also unknown, can be expressed as:

$$R_{{\psi }_1}=\underset{{\psi }_1\in {\Psi }_1}{\text{max}}p\left({\psi }_1\mid {\mathbf{y}}_{1:k+1},{\mathbf{a}}_{0:k}\right)-\underset{{\psi }_1\in {\Psi }_1}{\text{max}}p\left({\psi }_1\mid {\mathbf{y}}_{1:k},{\mathbf{a}}_{0:k-1}\right),$$ (33)

where $p\left({\psi }_1\mid {\mathbf{y}}_{1:k},{\mathbf{a}}_{0:k-1}\right)=\sum _{{\psi }_2^{\prime }\in {\Psi }_2}\sum _{{\varphi }^{\prime }\in \Phi }p\left({\psi }_1,{\psi }_2^{\prime },{\varphi }^{\prime }\mid \right.\left.{\mathbf{y}}_{1:k},{\mathbf{a}}_{0:k-1}\right)$. $R_{{\psi }_2}$ and $R_{\psi }$ can also be derived similarly. More information about the causal reward functions is provided in the supplementary document.

Fig. 6 represents the performance of different data collection policies in terms of average true model posterior probability for inferring each subset of unknown parameters. Columns 1, 2, and 3 of this figure illustrate the average results over 100 steps for the parameter sets $\Psi ,{\Psi }_1$, and ${\Psi }_2$, respectively. Additionally, Fig. 6(a), (b), and (c) in each of the rows showcase the data collection results for inference of unknown parameter vectors $\psi ,{\psi }_1$, and ${\psi }_2$, correspondingly.

Fig. 6:

Performance of different data collection policies in terms of true model posterior probability for inference of three parameter vectors: a) $\psi =\left[\left(q_{11}^1,q_{12}^1,q_{13}^1,q_{14}^1\right),\left(q_{11}^2,q_{12}^2,q_{13}^2,q_{14}^2\right]\right.$ b) ${\psi }_1=\left[\left(q_{11}^1,q_{12}^1,q_{13}^1,q_{14}^1\right)\right]$ c) ${\psi }_2=\left[\left(q_{11}^2,q_{12}^2,q_{13}^2,q_{14}^2\right)\right]$.

In Fig. 6(a), our method shows an average true model posterior probability of 20% after 100 steps for the inference of $\psi$, which is better than the performance obtained by all the other policies in this case. Moreover, case (a) shows a better performance for inferring $\psi$ compared to cases (b) and (c). This is because in case (a), the proposed method has a superior performance in inferring both ${\psi }_1$ and ${\psi }_2$ simultaneously compared to other policies, whereas cases (b) and (c) aim to infer only one of ${\psi }_1$ or ${\psi }_2$. Further, Fig. 6(b) shows that the proposed method has an average true model posterior probability of near 80% for inferring the unknown parameter vector ${\psi }_1$ after 100 steps, but the other policies show a much lower performance with the BOReL policy showing a performance of about 40%, active learning policy reaching performance of less than 40%, and the MAP policy showing performance of less than 20% after 100 steps. Also, our policy in case (b) performs better in inferring ${\psi }_1$ compared to the other cases, which aligns with its goal; however, it has lower performance for inferring $\psi$ and ${\psi }_2$ compared to the proposed policy in the other cases. Finally, one can see that in Fig. 6(c), our method performs better than the other policies in all the steps for inferring ${\psi }_2$. In this case, for inferring ${\psi }_2$, the proposed method reaches about 60% of performance, the active learning and BOReL policies get to about 20%, and the MAP policy remain less than 20% after 100 steps. Additionally, as expected, case (c) shows the best performance for inference of ${\psi }_2$ compared to the other cases.

For further analysis regarding the last experiments, a snapshot of the average posterior probabilities for different elements of the parameter sets ${\Psi }_1$ and ${\Psi }_2$ at time step 100 is visualized for different inference cases in Fig. 7. As a reminder, the 4th element of ${\Psi }_1$ and the 14th element of ${\Psi }_2$ correspond to ${\psi }_1^{*}$ and ${\psi }_2^{*}$ respectively. It can be seen that in case (a), these elements have the largest average posterior probabilities among the other elements under the proposed policy, which shows that our method has recognized the true underlying parameters with high confidence. This is where the other policies show low and similar to each other posterior probabilities for different elements and fail to distinguish the true parameters. In case (b), the 4th element of ${\Psi }_1$ using our approach shows the highest posterior probability among all the other cases. However, in this case, the average posterior probabilities for different elements of ${\Psi }_2$ show almost the same value under all the policies. This is because in case (b), the focus was on inferring the parameter vector ${\psi }_1$, and considering this objective, the best set of data collection actions was selected to maximize the confidence for inferring ${\psi }_1$ without considering the consequences for inference of ${\psi }_2$. Lastly, in case (c), the 14th element of ${\Psi }_2$ has the largest average posterior probability among the other elements under the proposed policy. This clearly shows the effectiveness of our approach in targeted data collection for inference of the parameter vector ${\psi }_2$.

Fig. 7:

The average posterior probabilities of the parameter sets ${\Psi }_1$ and ${\Psi }_2$ at time step 100 for three data collection objectives: a) Inferring the parameter vector $\psi$ b) Causal Inference of ${\psi }_1$ c) Causal Inference of ${\psi }_2$.

B. p53-MDM2 Regulatory Network with Unknown State Process

Gene regulatory networks are commonly modeled through HMMs with binary state variables [46]–[49]. We consider the p53-MDM2 gene regulatory network [50]–[53] shown in Fig. 8 for our analysis. The suppressive interactions are shown using blunt arrows, and activating regulations are represented by standard arrows. The protein known as $p53$, which is produced from the p53 gene, is crucial for regulating how cells respond to different stress signals that have the potential to induce genomic instability.

Fig. 8:

Activation/repression pathway diagram for the p53-MDM2 network model.

The p53-MDM2 network consists of 4 genes with the state vector at time step $k$ represented by ${\mathbf{x}}_k={\left[{\mathbf{x}}_k\left(1\right),{\mathbf{x}}_k\left(2\right),{\mathbf{x}}_k\left(3\right),{\mathbf{x}}_k\left(4\right)\right]}^T=\left[AT{M}_k,p{53}_k\right.,{Wip1}_k,{\left.MDM{2}_k\right]}^T$. Each state variable takes a value from the binary set ${\mathbf{x}}_k\left(i\right)\in \left\{\mathrm{0,1}\right\}$, where 1 and 0 represent the activation and inactivation of the $i\text{th}$ gene, respectively. The state process for this gene regulatory network can be represented by [54]:

$${\mathbf{x}}_k=\overline{\mathbf{R}{\mathbf{x}}_{k-1}}\oplus {\mathbf{a}}_{k-1}\oplus {\mathbf{n}}_k,\text{for}k=\mathrm{1,2},\dots ,$$ (34)

where $\overline{\mathbf{v}}$ maps the positive elements of a vector to 1 and others to 0, “$\oplus$” indicates component-wise modulo-2 addition (i.e., XOR), ${\mathbf{a}}_{k-1}$ is the action (i.e., perturbation or excitation) at time $k-1$, and ${\mathbf{n}}_k$ is the transition noise at time $k$ which is modeled by an independent Bernoulli random variable with intensity $p=0.001$. Additionally, $\mathbf{R}$ in (34) is the connectivity matrix of size 4 × 4, expressing the types of interactions between the genes, represented by:

$$\mathbf{R}={\left[\begin{array}{cccc}0& +1& 0& -1\\ 0& 0& +1& +1\\ -1& -1& 0& +1\\ 0& -1& 0& 0\end{array}\right]}^T.$$ (35)

The element in the $i\text{th}$ row and $j\text{th}$ column of the connectivity matrix ($\mathbf{R}$) is denoted as $r_{ij}$, which represents the type of regulations from gene $i$ to gene $j$. $r_{ij}$ takes in $\{-\mathrm{1,0},+1\}$, where +1 and −1 represent the positive and negative regulations and 0 indicates no regulation. The following 5 regulations in (35) are assumed to be unknown: $\theta =\left(r_{21},r_{24},r_{42},r_{22},r_{41}\right)$ with true values of ${\theta }^{*}=(+1,-1,+\mathrm{1,0},-1)$. Since each regulation can take three values in $\{-\mathrm{1,0},+1\}$, the parameter set $\Theta$ consists of $3^5=243$ distinct parameter vectors. With $2^4=16$ states within the p53-MDM2 network, the belief space for this problem can be characterized as ${\Delta }_{16\times 243}$.

The following action or perturbation space is considered for our analyses: $𝒜=\left\{{\mathbf{a}}^1=1000,{\mathbf{a}}^2=0100,{\mathbf{a}}^3=0010,{\mathbf{a}}^4=\right.\left.0001,{\mathbf{a}}^5=0000\right\}$. The first four actions alter the state value of a single gene at each time step. For instance, considering (34), ${\mathbf{a}}^1$ flips the first gene’s value from 0 to 1 or 1 to 0. Note that this assumption closely aligns with common practice for data collection of gene regulatory networks in experimental settings, where drugs designed for targeted genes are used for gene excitations and data collection [55]. Note that the fifth action, ${\mathbf{a}}^5$, leaves the state of all genes unchanged.

The specific type of gene-expression data determines the characteristics of the HMM observation model. This paper considers the following Gaussian model, which is often used for representing cDNA microarrays or live-cell imaging-based assays [56], [57]:

$${\mathbf{y}}_k\left(i\right)=m+\delta {\mathbf{x}}_k\left(i\right)+{\mathbf{v}}_k\left(i\right),k=\mathrm{1,2},\dots ,$$ (36)

for $i=1,\dots ,4$. ${\mathbf{v}}_k\left(i\right)~𝒩\left(0,{\sigma }^2\right)$ represents a zero-mean Gaussian noise vector with uncorrelated elements, $m$ denotes the baseline expression which refers to the “zero” states, and $\delta$ is the differential expression value. For our experiments, we assume that the parameters of the observation model in (36) are known and provided as follows: $m=20,\delta =30$, and $\sigma =5$.

The performance of the proposed method using the point-based reward in (15) is assessed in terms of average maximum posterior probability for 50 time steps. Fig. 9 shows that our method achieves a maximum posterior probability of nearly 1 after about 15 steps, whereas other approaches exhibit notably lower performance, with active learning and MAP reaching around 0.3 and BOReL demonstrating a performance of 0.5 by the 15th step. Further, after 50 steps, active learning and MAP policies achieve a performance of less than 0.8, while BOReL reaches about 0.9, which still falls short of the performance demonstrated by the proposed method. The active learning and MAP policies exhibit nearly identical performance, primarily due to their greedy nature. Unlike our approach, these policies prioritize immediate gains without considering the long-term implications of their actions on the inference objective, resulting in slower improvement over time. Meanwhile, BOReL’s lack of joint modeling of state and parameter uncertainty has also led to a lower performance compared to our proposed method.

Fig. 9:

Performance of different data collection policies for inferring all unknown parameters of the p53-MDM2 network.

Four variations of the previous experiment are conducted to evaluate the robustness of our method under different noise conditions. Two process noise levels ($p=\mathrm{0.001,0.05}$) and two measurement noise levels ($\sigma =\mathrm{5,10}$) are considered, resulting in four different combinations. The performance of various data collection methods in these experiments is reported in terms of the average maximum posterior probability per step. As shown in Table I, the results show that performance declines as noise increases. However, our method consistently outperforms all other approaches across all noise scenarios, demonstrating its robustness in noisy environments.

TABLE I:

Average maximum posterior probability per step obtained using different data collection policies for inferring unknown parameters of the p53-MDM2 network under various process noise ($p$) and measurement noise ($\sigma$) values.

$\sigma$╲$p$	0.001				0.05
	Proposed Policy	Active Learning	MAP Policy	BOReL Policy	Proposed Policy	Active Learning	MAP Policy	BOReL Policy
5	0.811	0.457	0.445	0.638	0.701	0.425	0.421	0.632
10	0.654	0.382	0.369	0.499	0.485	0.265	0.263	0.329

The performance of the proposed method for causal inference of the unknown regulation $r_{24}=-1$, while other regulations are also unknown, is analyzed here. Fig. 10(a) shows the average maximum posterior probability over the entire parameter set (i.e., $\Theta$). The low average maximum posterior probability over the entire parameter space under the proposed method comes from the focus of the proposed policy on the inference of the regulation $r_{24}$. The average posterior probability of all the unknown regulations under the proposed method and active learning method is presented in Fig. 10(b). The average posterior probabilities of regulations being −1, 0, and 1 are shown using black, blue, and red, respectively. It can be seen that the posterior probability of being −1 for $r_{24}$ increases quickly to near 1 only after collecting 10 data and gets to 1 after 15 data, whereas the same posterior probability has a very slow increase using active learning and BOReL policies. This posterior probability reaches close to 1 after 50 steps using the active learning policy, and it reaches about 0.75 using the BOReL policy. Moreover, Fig. 10(b) illustrates that the true model posterior probability of the regulations $r_{21}=+1$ and $r_{42}=+1$ have increased considerably under data collected for inference of $r_{24}$ by the proposed policy, whereas these data have not substantially increased the true model posterior probability of $r_{22}=0$ and $r_{41}=-1$. This justifies the low maximum posterior probability of the entire parameter set in Fig. 10(a), as well as the capability of the proposed policy in systematically taking actions for causal inference.

Fig. 10:

(a) The average posterior probability of the entire parameters, and (b) average posterior probabilities of all the unknown interactions, for causal inference of the unknown regulation $r_{24}$ in the p53-MDM2 network.

The distribution-based criterion (entropy) in (16) is considered in this part for the inference of the p53-MDM2 network. Fig. 11 shows the average distribution of the posterior probabilities under different policies at time steps 5 and 10. It can be seen that after only 5 steps/data, the posterior probabilities over the parameter set $\Theta$ are more peaked (roughly between 150 and 243) using the proposed policy (the true parameter belongs to this set). Further, it can be seen that using 10 data, the posterior probability of the true parameter (model 184), shown as red, is peaked at almost 1, which means that the true model has been identified with high confidence after only 10 data. On the other hand, the other policies show a more uncertain (not peaked) distribution of the posterior probabilities in both time steps 5 and 10. More results related to the p53-MDM2 network can be found in the supplementary materials.

Fig. 11:

Average posterior probabilities of the parameter set $\Theta$ at time steps 5 and 10 while reducing the entropy using different data collection policies; the non true models are shown with black while the true model is indicated by red.

C. Littman’s Network with Unknown Measurement Process

Littman’s network [58] is considered to study the effectiveness of our proposed data collection method in communication networks. The Littman’s network comprises 7 states as $𝒳=\left\{{\mathrm{x}}^1=00,{\mathrm{x}}^2=20,{\mathrm{x}}^3=40,{\mathrm{x}}^4=60,{\mathrm{x}}^5=80,{\mathrm{x}}^6=\right.\left.100,{\mathbf{x}}^7=\text{CRASH}\right\}$. Each of $\{\mathrm{00,20},\dots ,100\}$ corresponds to the current load of the communication network, and CRASH refers to a state where the network is overloaded and eventually crashes. Note that when the system is in states ${\mathbf{x}}^1=00$ and ${\mathbf{x}}^7=\text{CRASH}$, the network load cannot be decreased or increased, respectively. Four actions are defined for this network as $𝒜=\left\{{\mathbf{a}}^1=\right.\text{UNRESTRICT},{\mathbf{a}}^2=\text{STEADY},{\mathbf{a}}^3=\text{RESTRICT},{\mathbf{a}}^4=\text{REBOOT}\}$. UNRESTRICT means that the network should increase its load; conversely, RESTRICT shows that the network should decrease its current load. STEADY is an action that enforces the network to stay in its current state, and finally, REBOOT is an action that relieves all the load from the network. It is assumed that the state transition parameters for this system are completely known. Complete information about the transition probabilities is provided in the supplementary material.

The measurement process provides the status of the network, with the measurement space $𝒴=\left\{{\mathbf{y}}^1=\text{UP},{\mathbf{y}}^2=\right.\text{DOWN}\}$. The measurement process can be expressed using the update vectors defined in (19) as:

$$T^{\theta }\left({\mathbf{y}}^1,\mathbf{a}\right)=\left[\begin{array}{c}1\\ 1\\ 1\\ t_1\\ t_2\\ t_3\\ 0\end{array}\right],T^{\theta }\left({\mathbf{y}}^2,\mathbf{a}\right)=\left[\begin{array}{c}0\\ 0\\ 0\\ 1-t_1\\ 1-t_2\\ 1-t_3\\ 0\end{array}\right],\text{for}\mathbf{a}\in 𝒜,$$

for each $\mathbf{a}\in 𝒜$, where $\theta =\left(t_1,t_2,t_3\right)$ is a single vector expressing the unknown parameters with the true parameters being ${\theta }^{*}=\left(t_1^{*}=0.8,t_2^{*}=0.5,t_3^{*}=0.2\right)$. The unknown elements of the update vector can take values as $t_i\in \left\{\mathrm{0.2,0.5,0.8}\right\}$, leading to a parameter set $\Theta$ of size $3^3=27$. Since Littman’s network has 7 states, the belief space for this problem can be described as ${\Delta }_{7\times 27}$.

The results of the proposed method for causal inference of each of the three unknown parameters, $t_1,t_2$, and $t_3$, are presented in Fig. 12. Note that each plot in Fig. 12 represents the average true posterior probability for the causal parameter that the methods aim to infer. For instance, Fig. 12(c) represents that the proposed method has achieved a value near 0.8 for the true posterior probability of $t_3$, meaning that with a confidence of near 80%, $t_3$ is equal to its true value. In contrast, active learning, BOReL, and MAP policies achieve much lower confidence values of about 45%, 40%, and 30%, respectively. Similarly, the proposed policy achieves a high confidence of about 75% in case (a), while the causal inference of parameter $t_2$ in case (b) has been more challenging, and with 50 data a confidence of near 40% is achieved in this case. This can be explained by the fact that $t_1^{*}=0.2$ and $t_3^{*}=0.8$ are associated with less uncertain parameters, whereas $t_2^{*}=0.5$ represents the parameter with the highest uncertainty in the observation model, making its inference more challenging compared to the others. Additional experiments on the Littmans’s network are included in the supplementary materials.

Fig. 12:

Performance comparison of different data collection techniques for three cases of causal inference in Littman’s network: a) causal inference of the parameter $t_1$ b) causal inference of the parameter $t_2$ c) causal inference of the parameter $t_3$.

VI. Conclusion and Future Work

This paper developed a deep reinforcement learning data collection approach for Bayesian inference of general hidden Markov models. At first, a belief state was defined as the joint distribution of the possible system states and models. Using this belief state, the optimal Bayesian lookahead data collection process was formulated to maximize the inference performance given the system and data uncertainty. Since the belief state is continuous, a deep reinforcement learning approach was utilized to approximate the optimal Bayesian data collection policy. The proposed policy can be trained offline, using simulated trajectories from the belief transition. This enables our data collection policy to account for the long-term effects of data collection decisions on the inference performance, as opposed to the short-term impacts. The trained policy can then be used for sequential and real-time data collection as new observations become available from the real system. The effectiveness of our proposed method in terms of point-based, distribution-based, and causal inference objectives was demonstrated using comprehensive numerical experiments over three systems.

Our future work will focus on scaling the proposed data collection method to handle larger systems with continuous state and action spaces. Additionally, we aim to explore more complex data collection objectives, such as fault detection, abnormality or attack detection, and systems with semi-Markov and non-Markovian dynamics. We also plan to investigate variations of our method for different domains, such as image processing, which would require compact representations of the belief state tailored to each domain. Furthermore, we will work on accelerating the training process and improving accuracy by incorporating techniques such as softmax. Finally, we will adapt our data collection method to other learning paradigms, including supervised and unsupervised learning.

Impact Statement—

Recent technological advances leading to data growth have prompted a simultaneous focus on studying increasingly complex systems and phenomena. Modeling complex systems through available data might require thousands or millions of data, which are impossible to collect due to time or cost constraints. This paper provides a systematic approach for collecting targeted data depending on the objective of the analyses. The results demonstrate that the amount of data is not the only factor for accurate modeling of complex systems; rather, targeted and informative data can shed light on the underlying complex mechanisms of these systems. Examples of such systems include ecology, economics, engineering, and medicine. In particular, the biological application of this work will help biologists collect genomic data that best identifies specific interactions between genes causing chronic diseases.

Biographies

Mohammad Alali is a fourth-year Ph.D. candidate in Electrical Engineering at Northeastern University. He received his B.Sc. degree in Electrical Engineering from University of Tehran, Tehran, Iran, in 2018, and M.Sc. degree in Electrical Engineering from Montana State University, Bozeman, USA, in 2021. His research interests lie in the domains of machine learning, Bayesian statistics, reinforcement learning, and experimental design. More specifically, he is currently investigating several projects with a focus on fast and efficient inference of various observable/partially observable Markov models. He is the recipient of the Best Paper Finalist Award at the 20th IFAC Symposium on System Identification in 2024 and the American Control Conference in 2023.

Mahdi Imani received his Ph.D. degree in Electrical and Computer Engineering from Texas A&M University, College Station, TX, in 2019. He is currently an Assistant Professor in the Department of Electrical and Computer Engineering at Northeastern University. His research interests include machine learning, Bayesian statistics, and decision theory, with applications spanning from computational biology to cyber-physical systems. He has received several prestigious awards, including the NIH NIBIB Trailblazer Award in 2022, the Oracle Research Award in 2022, the Outstanding Associate Editor Award from IEEE Transactions on Neural Networks and Learning Systems, the NSF CISE Career Research Initiation Initiative Award in 2020, the Association of Former Students Distinguished Graduate Student Award for Excellence in Research-Doctoral in 2019, and the Best Paper Finalist Award at the American Control Conference in 2023 and the 49th Asilomar Conference on Signals, Systems, and Computers in 2015. He currently serves as an Associate Editor for IEEE Transactions on Neural Networks and Learning Systems (TNNLS) and IEEE Transactions on Vehicular Technology (TVT).