Year: 2021

Abstract:
Solving real-life sequential decision making problems under partial observability involves an exploration-exploitation problem. To be successful, an agent needs to efficiently gather valuable information about the state of the world for making rewarding decisions. However, in real world applications, acquiring valuable information is often highly costly, e.g., in the medical domain, information acquisition might correspond to performing a medical test on a patient. This poses a significant challenge for the agent to learn optimal policy for the task. In this paper, we propose a model-based reinforcement learning framework that learns a policy which solves this exploration-exploitation problem during its execution. Key to the success is a novel sequential variational autoencoder that learns high-quality representations from partially observed states, which are then used by the policy to maximize the task reward in a cost-efficient manner. We demonstrate the efficacy of our proposed framework in a control domain as well as using a medical simulator. In both tasks, our proposed method outperforms conventional baselines and results in policies with greater cost efficiency.

heading: 1 INTRODUCTION
text: Recently, machine learning models for automated sequential decision making have shown remarkable success across many application areas, such as visual recognition (Mathe et al., 2016; Das et al., 2017), robotics control (Finn et al., 2016; Zhang et al., 2018), medical diagnosis (Ling et al., 2017; Peng et al., 2018) and computer games (Mnih et al., 2015; Silver et al., 2016). One fundamental reason that drives the success of such models and enables them to outperform classical algorithms is the availability of large amounts of training data. Typically such training data is either fully observed or the features stem from an action-independent observation model (which clearly can depend on the state of the system). However, the fundamental assumption that the same features are always readily available during deployment could not hold in many real-world applications. For instance, consider a medical support system for monitoring and treating patients during their stay at hospital which was trained on rich historical medical data. To provide the best possible treatment, the system might need to perform several measurements of the patient over time, while some of them could be costly or even pose a health risk. Therefore, during deployment, it is more ideal that the system could function with minimal features while during training more features might have been available. In such cases, we are interested in decision making models that actively take the measurement process, i.e., feature acquisition, into account and only acquire the information relevant for making a decision.
In this paper, we consider the challenging problem of learning effective policies when the cost of information acquisition cannot be neglected. To be successful,we need to learn policies which acquires the information required for solving a task in the cheapest way possible. For simplicity, we can think of the policy as being constituted of an acquisition policy which actively selects meaningful features to be observed and a task policy, which selects actions to change the state of the system towards some goal.1 As such, we consider a partially observable learning problem with the following two distinguishing properties compared to the most commonly studied problems (see also Figure 3.2 for an illustration). (i) By incorporating active feature acquisition, the training of the task policy is based upon subsets of features only, i.e., there are missing features, where the missingness is
1Clearly, these two policies are not independent in general, e.g., acquiring features can change the state of the system.
observe
action
acquire
(e.g. navigation) (e.g. medical treatments)
observe action
controlled by the acquisition policy. Thus, the resulting POMDP is different from the conventional POMDPs in RL literature (Cassandra, 1998) where the partial observability for later stems from a fixed and action-independent observation model. Also, the state transitions in conventional POMDPs are only determined by the choice of the task action, whereas in our setting the state-transition is affected by both the task action and the feature acquisition choice. (ii) The learning of the acquisition policy introduces an additional dimension to the exploration-exploitation problem: each execution of the policy needs to solve an exploration-exploitation problem, and thus we often need to learn sophisticated policies.
Most reinforcement learning research has not taken active feature acquisition into consideration. In this work, we propose a unified approach that jointly learns a policy for optimizing the task reward while performing active feature acquisition. Although some of the prior works have exploited the use of reinforcement learning for sequential feature acquisition tasks (Shim et al., 2018; Zannone et al., 2019), they considered variable-wise information acquisition in a static setting only, corresponding to feature selection for non-time-dependent prediction tasks. However, our considered setting is truly time-dependent and feature acquisitions need to be made at each time step while the state of the system evolves simultaneously. As such, both the model dynamics of the underlying MDP and the choice of feature acquisition introduce considerable challenges to the learning of the sequential feature acquisition strategy.
Due to the challenge of the exploration-exploitation problem, it is a non-trivial task to jointly learn the two policies. The conventional end-to-end approaches often result in inferior solutions in complex scenarios. Ideally, policies based on high-quality representations would be easier for the algorithm to search for better solutions through exploration-exploitation. Therefore, our proposed framework also tackles the joint policy training task from a representation learning perspective. Specifically, we introduce a representation learning model that not only encodes the sequential partially observed information into its latent features, but also efficiently imputes the unobserved features to offer more meaningful information for the policy training. To this end, we formulate a sequential generative model that can efficiently learn model dynamics during representation learning. Overall, the contributions of our paper are three-fold:
• We propose an approach for learning sequential decision making policies with active feature acquisition through a unified reinforcement learning framework. Our proposed approach simultaneously learns policies for reward optimization and active feature acquisition.
• We present a novel sequential representation learning approach to account for the encoding of the partially observed states. Our proposed approach is based on variational autoencoders (VAE) with amortized inference. The imputation of the unobserved features is achieved via learning the model dynamics.
• We demonstrate our proposed framework can be applied to various applications. We conduct extensive experiments on an image-based control task as well as a medical simulator fitted from real-life data where our method shows clear improvements over conventional baselines.

heading: 2 RELATED WORK
text: In this work, we integrate active learning with reinforcement learning to accomplish the policy training task while attempting to acquire fewest observed features as possible. We thus review related methods on active feature acquisition and representation learning for POMDP, respectively.

heading: 2.1 ACTIVE FEATURE ACQUISITION
text: Our work draws motivation from the existing instance-wise active feature selection approaches. One category of the instance-wise feature selection methods consider feature acquisition as a one time effort to select a subset of features as a whole. One typical example is the conventional linear model that poses sparsity inducing prior distribution to the model (Tibshirani, 1996). Recently, there also emerged approaches that adopt reinforcement learning to actively find optimal feature subsets (Yoon et al., 2018; Shim et al., 2018; Zannone et al., 2019). Though such attempts have demonstrated certain efficacy in handling non time-series instance-wise data, they do not suffice for handling sequential dataset. There is also an alternative category that models feature acquisition as a Bayesian experimental design (Ma et al., 2019; Gong et al., 2019). However, the sequential decision making is for variable-wise feature acquisition and the problems are still non time-series tasks in nature.
The key difference between all the aforementioned approaches and ours is that we tackle active feature acquisition problems with time-series data, where an active feature selection decision needs to be formed at each time step along the multi-step reinforcement learning trajectory. Therefore, the feature acquisition for our presented work needs to consider more complex information over model dynamics and control, apart from the static instance-wise features.

heading: 2.2 REPRESENTATION LEARNING IN POMDP
text: In complex tasks, policies trained upon different representations can even converge to different performance levels. Most conventional deep reinforcement learning approaches unifies the process of representation learning with policy training and results in policies trained in an end-to-end fashion (Mnih et al., 2013; Lillicrap et al., 2016; Mnih et al., 2016). However, to accomplish the representation learning task, such models often engage trainable parameters which could come with considerable size and thereby result in significant degradation in sample efficiency.
When considering problems with POMDPs where the state space is partially accessible to the agent, representation learning becomes an important and non-trivial research challenge. Among the existing literature, one prominent line of research tackles the representation learning for POMDP in an off-line fashion and thus resulting in multi-stage reinforcement learning. Higgins et al. (2016; 2017) adopt pretrained VAE models as a representation module to build agents with strong domain adaptation performance. The key difference between their work and ours is that they encode instance-wise image frames from POMDP domains where each image presents a partial view over the task environment, while our work considers cost-sensitive reinforcement learning with distinct partial observability, i.e., the feature-level information is missing at each time step for the agent. We thus adopt a sequential representation learning approach to infer a more representative state information. Recently, there also emerged several works on sequential representation learning for POMDP (Gregor et al., 2019; Vezzani et al., 2019). However, most of the works utilize VAE training as an auxiliary task to jointly update the representation model with the policy learning loss. In our work, due to the high acquisition cost to observe the features, we adopt an off-line representation learning setting. Also, our proposed representation learning is model-based, where the model learns to impute the missing features with such attempt yielding significant benefit to derive high-quality representation for policy training.

heading: 3 METHODOLOGY

heading: 3.1 TASK SETTING
text: In this section, we formally define the problem settings for the task of jointly learning the task and feature acquisition policy. To this end, we define the active feature acquisition POMDP, a rich class of discrete-time stochastic control processes generalizing standard POMDPs:
Definition 1 (AFA-POMDP). The active feature acquisition POMDP is a tuple M = 〈S,A, T ,O,R, C, γ〉, where S is the state space and A = (Af ,Ac) is a joint action space of feature acquisition actionsAf and control actionsAc. The transition kernel T : S ×Ac×Af → PS maps any joint action a = (af ,ac) in state s ∈ S to a distribution PS over next states. In each state s, when taking action af , the agent observes xp = x(af ), i.e., a subset of the features x = (xp,xu) ∼ O(s) indicated by af , whereO(s) is a distribution over possible feature observation for state s and xu are features not observed by the agent. When taking a joint action, the agent obtains
rewards according to the reward functionR : S ×Ac → R and pays a cost of C : S ×Af → R+ for feature acquisition. Rewards and costs are discounted by the discount factor γ ∈ [0, 1).
Simplifying assumptions For simplicity, we assume that x consists of a fixed number of features Nf for all states, that Af = 2[Nf ] is the power set of all the Nf features, and that xp(af ) consists of all the features in x indicated by the subset af ∈ Af . Note that the feature acquisition action for a specific application can take various different forms. For instance, in our experiments in Section 4, for the Sepsis task, we define feature acquisition as selecting a subset over possible measurement tests, whereas for the Bouncing Ball+ task, we divide an image into four observation regions and let the feature acquisition policy select a subset of observation regions (rather than raw pixels). Please also note that while in a general AFA-POMDP, the transition between two states depends on the joint action, we assume in the following that it depends only on the control action, i.e., T (s,ac,af ′) = T (s,ac,af ) for all af ′ ,af ∈ Af . While not true for all possible applications, this assumption can be a reasonable approximation for instance for medical settings in which tests are non-invasive. For simplicity we furthermore assume that acquiring each feature has the same cost, denoted as c, i.e., C(af , s) = c |af |, but our approach can be straightforwardly adapted to have different costs for different feature acquisitions.
Objective We aim to learn a policy which trades off reward maximization and the cost for feature acquisition by jointly optimizing a task policy πc and a feature acquisition policy πf . That is, we aim to solve the optimization problem
max πf ,πc
E  ∞∑ t=0 γt ( R(xt,act)− |Af |∑ i c · I (af(i)t ) ), (1)
where the expectation is over the randomness of the stochastic process and the policies, af(i)t denotes the i-th feature acquisition action at timestep t, and I (·) is an indicator function whose value equals to 1 if that feature has been acquired.
Note that the above optimization problem is very challenging: an optimal solution needs to maintain beliefs bt over the state of the system at time t which is a function of partial observations obtained so far. Both the the feature acquisition policy πf (aft | bt) and the task policy i.e., πc(act | bt) depend on this belief. The information in the belief itself can be controlled by the feature acquisition policy through querying subsets from the features xt and hence the task policy and feature acquisition policy itself strongly depend on effectiveness of the feature acquisition policy.

heading: 3.2 SEQUENTIAL REPRESENTATION LEARNING WITH PARTIAL OBSERVATIONS
text: We introduce a sequential representation learning approach to facilitate the task of policy training with active feature acquisition. Let x1:T = (x1, ...,xT ) and a1:T = (a1, ...,aT ) denote a sequence of observations and actions, respectively. Alternatively, we also denote these sequences as x≤T and a≤T . Overall, our task of interest is to train a sequential representation learning model to learn the distribution of the full sequential observations x1:T , i.e., for both the observed part x p 1:T and the unobserved part xu1:T . Given only partial observations, we can perform inference only with the observed features xp1:T . Therefore, our proposed approach extends the conventional unsupervised representation learning task to a supervised learning task, which learns to impute the unobserved features by synthesizing the acquired information and learning the model dynamics.
As such, the key underlying assumption is that learning to impute the unobserved features would result in better representations which can be leveraged by the task policy. And performing sequential representation learning, as we propose, is a more adequate choice than non-sequential modeling, for our task of interest with partial observability. Furthermore, unlike many conventional sequential representation learning models for reinforcement learning that only reason over the observation sequence xp1:T , in our work, we take into account both the observation sequence x p 1:T and the action sequence a1:T for conducting inference. The intuition is that since x p 1:T by itself consists of very limited information over the agent’s underlying MDP state, incorporating the action sequence would be an informative add-on to the agent’s acquired information to infer the belief state. To summarize, our proposed sequential representation model learns to encode xp1:T and a1:T into meaningful latent
1
2
3
0 1 2
ℎ1
1 2 3
ℎ2 ℎ3
1
2
3 Decoder Inference
Figure 2: Observation decoder and belief inference model for the partially observable sequential VAE. Shaded nodes represent the observed variables. The inference model filters information over the partial observations and actions, to predict both the observed and unobserved features.
features, for predicting xp1:T and x u 1:T . The architecture of our proposed sequential representation learning model is shown in Figure 2.
Observation Decoder Let z1:T = (z1, ..., zT ) denote a sequence of latent states. We consider the following probabilistic model:
pθ(x p,xu, z) = T∏ t=1 p(xpt ,x u t |zt) p(zt), (2)
For simplicity of the notations, we assume z0 = 0. We impose a simple prior distribution over z, i.e., a standard Gaussian prior, instead of incorporating some learned prior distribution over the latent space of z, such as an autoregressive prior distribution like p(zt|zt−1,xp1:t,a0:t−1). The reason is that using a static prior distribution results in latent representation zt that is stronger regularized and more normalized than using a learned prior distribution which is stochastically changing over time. This is crucial for deriving stable policy training performance. At time t, the generation of data xpt and x u t depends on the corresponding latent variable zt. Given zt, the observed variables are conditionally independent of the unobserved ones. Therefore, p(xpt , x u t |zt) = p(x p t |zt) p(xut |zt). (3)
Belief Inference Model During policy training we only assume access to partially observed data. This requires an inference model which takes in the past observation and action sequences to infer the latent states z. Specifically, we present a structured inference network qφ as shown in Figure 2, which has an autoregressive structure:
qφ(z|x,a) = ∏ t qφ(zt|xp≤t,a<t), (4)
where qφ(·) is a function that aggregates the filtering posteriors of the history of observation and action sequences. Following the common practice in existing sequential VAE literature, we adopt a forward RNN model as the backbone for the filtering function qφ(·) (Gregor et al., 2019). Specifically, at step t, the RNN processes the encoded partial observation xpt , action at−1 and its past hidden state ht−1 to update its hidden state ht. Then the latent distribution zt is inferred from ht. The belief state bt is defined as the mean of the distribution zt. By accomplishing the supervised learning task, the belief state could provide abundant information for not only the observed sequential features, but also for the missing features, so that the policy trained over it could benefit from it and progress faster towards getting better convergent performance.
Learning We proposed to pre-train both the generative and inference models offline before learning the RL policies. In this case, we assume the access to the unobserved features, so that we can construct a supervised learning task to learn to impute unobserved features. Concretely, the pre-training task update the parameters θ, φ by maximizing the following variational lower-bound (Jordan et al., 1999; Kingma & Welling, 2013):
log p(xp,xu) ≥ Eqφ [∑
t
log pθ(x p t |zt) + log pθ(xut |zt)− KL ( qφ(zt|xp≤t,a<t) || p(zt) )] = ELBO(xp,xu). (5)
By incorporating the term log pθ(xut |zt), the training of sequential VAE generalizes from an unsupervised task to a supervised task that learns the model dynamics from past observed transitions and imputes the missing features.
We perform multi-stage reinforcement learning to jointly learn the feature acquisition policy and the task policy. The VAE model is pretrained and kept fixed during policy learning. The reason for not updating VAE online is that computing the loss in Eq (5) would require the access to unobserved features and therefore, is cost intensive. The pseudocode for our proposed method is in Appendix A.

heading: 4 EXPERIMENTS
text: We examine the characteristics of our proposed model in the following two experimental domains: a bouncing ball control task with high-dimensional image pixels as input, adapted from (Fraccaro et al., 2017); a sepsis medical simulator fitted from real-world data (Oberst & Sontag, 2019).
Baselines For comparison, we mainly consider variants of the strong VAE baseline beta-VAE (Higgins et al., 2016), which works on non-time-dependent data instances. For representing the missing features, we adopt the zero-imputing method, proposed in (Nazabal et al., 2018) over the unobserved features. Thus, we denote the VAE baseline as NonSeq-ZI. We train the VAE with either the full loss over the entire features, or the partial loss which only applies to the observed features (Ma et al., 2019). We denoted our proposed sequential VAE model for POMDPs as Seq-PO-VAE. All the VAE-based approaches adopt an identical policy architecture. Detailed information on the model architecture is presented in appendix.
Data Collection To pre-train the VAE models, data generated by a non-random policy is unavoidably needed to incorporate abundant dynamics information. For both tasks, we collect a small scale dataset of 2000 trajectories, where half of the data is collected from a random policy and the the other half from a policy which better captures the state space that would be encountered by a learned model (e.g., by training a data collection policy end-to-end or using human generated trajectories). The simple mixture of dataset works very well on both tasks without the need of further fine-tuning the VAEs. We also create a testing set that consists of 2000 trajectories to evaluate the models.
4.1 BOUNCING BALL+
Task Settings We adapted the original bouncing ball experiment presented in (Fraccaro et al., 2017) by adding a navigation objective and introducing control actions. Specifically, a ball moves in a 2D box and at each step, a binary image of size 32× 32 showing the box and the ball. Initially, the ball appears at a random position in the upper left quadrant, and has a random velocity. The objective is to control the ball to reach a fixed target location set at (5, 25). We incorporate five RL actions: a null action and four actions for changing the velocity of the ball in either the x or y direction with a fixed scale: {∆Vx : ±0.5, ∆Vy : ±0.5, null}. A reward of 1.0 is issued if the ball reaches its target location. Each episode runs up to 50 time steps.
Representation Learning Results We evaluate the missing feature imputing performance of each VAE model in terms of negative log likelihood (nll) loss and present results in Table 1. We notice that our proposed model yields to significantly better imputing result than all the other baselines. This reveals the fact that our proposed sequential VAE model can efficiently capture the environment dynamics and learn meaningful information over the missing features. Such effect is vital in determining the policy training performance in AFA-POMDP, since the policy is conditioned on the VAE latent features. We also demonstrate sample trajectories reconstructed by different VAE models in the Appendix. The result shows that our model learns to impute significant amount of missing information given the partially observed sequence.
Policy Training Results We evaluate the policy training performance in terms of episodic number of acquired observations and the task rewards (w/o cost). The results are presented in Figure 3 (a) and (b), respectively. First, we notice that the end-to-end method is vital and fails to learn task skills under the given feature acquisition cost. However, the VAE-based representation learning methods manage to learn the navigation skill under the same cost setting. This verifies our assumption that representation learning could bring significant benefit to the policy training under the AFA-POMDP scenario. Furthermore, we also notice that the joint policies trained by Seq-PO-VAE can develop the target navigation skill at a much faster pace than the non-sequential baselines. Our method also converges to a standard where much less feature acquisition is required to perform the task.
We also show that our proposed method can learn meaningful feature acquisition policies. To this end, we show three sampled trajectories upon convergence of training in Figure 4. From the examples, we notice that our feature acquisition policy acquires meaningful features with a majority grasping the exact ball location. Thus, it demonstrates that the feature acquisition policy adapts to the dynamics of the problem and learns to acquire meaningful features. We also show the actively learned feature acquisition policy works better than random acquisition. From the results in Figure 4 (c), our method converges to better standard than random policies with considerably high selection probabilities.

heading: 4.2 SEPSIS MEDICAL SIMULATOR
text: Task Setting Our second evaluation domain adopts a medical simulator for treating sepsis among ICU patients, proposed in (Oberst & Sontag, 2019). Overall, the task is to learn to apply three treatment actions to the patient, i.e, {antibiotic, ventilation, vasopressors}. The state space consists of 8 features: 3 of them indicate the current treatment state for the patient; 4 of them are the measurement states over heart rate, sysBP rate, percoxyg state and glucose level; the rest is an index specifying the patent’s diabetes condition. The feature acquisition policy learns to actively select the measurement features. Each episode runs for up to 30 steps. The patient will be discharged if his/her measurement states all return to normal values. An episode terminates upon mortality or discharge, with a reward −1.0 or 1.0.
Representation Learning Result We evaluate the imputation performance for each VAE model on the testing dataset. The loss is evaluated in terms of MSE, presented in Table 1. Our proposed method leads to the lowest MSE loss compared to the baselines. The result reveals that our proposed sequential VAE could promisingly learn model dynamics for tasks with stochastic transitions.
Policy Training Result We show the policy training results for Sepsis in Figure 5. Overall, our proposed method results in substantially better task reward compared to all baselines. Noticeably, the learning of discharge for our method progresses significantly faster than baseline approaches and converges to substantially better values. The result shows that our method can be trained in a much more sample efficient way. Moreover, upon convergence, our model outperforms the best non-sequential VAE baseline with a gap of > 5% for discharge ratio. For all the evaluation metrics, we notice that VAE-based representation learning models outperform the end-to-end baseline by significant margins. This indicates that efficient representation learning may be crucial for deriving satisfying task performance in AFA-POMDP setting. The result also reveals that learning to impute missing features contributes greatly to improve the policy training performance for AFA-POMDP.
Ablation: Efficacy of Active Feature Acquisition We study the effect of actively learning sequential feature acquisition strategy with RL. To this end, we compare our method with a baseline that randomly acquires features. We evaluate our method under different cost values, and the results are shown in Figure 6. From the results, we notice that there is a clear cost-performance trade-off, i.e., a
higher feature acquisition cost results in feature acquisition policies that obtain fewer observations, with a sacrifice of task performance. Overall, our acquisition method results in significantly better task performance than the random acquisition baselines. Noticeably, with the learned active feature acquisition strategy, we acquire only about half of the total number of features (refer to the value derived by Random-100%) to obtain comparable task performance. Also, we notice that the specified cost has a very clear impact on the final task performance, i.e., the number of acquired features per episode decreases significantly as the cost increases. Thereby, our proposed solution can promisingly compute feature acquisition policies that meet different budgets.
Ablation: Impact on Total Acquisition Cost For different representation learning methods, we also investigate the total number of features acquired at different stage of training. The results are shown in Figure 7. As expected, to obtain better task policies, the models need to take longer training steps and thus the total feature acquisition cost would increases accordingly. We notice that policies trained by our method result in the highest convergent task performance (max x-value). Given a certain performance level (same x-value), our method consumes substantially less total feature acquisition cost (y-value) than the others. We also notice that the overall feature acquisition cost increases with a near exponential trend. Therefore, it is essential to train the policy for AFA-POMDP with advanced representation learning method, so that the feature acquisition cost could be reduced.

heading: 5 CONCLUSION
text: We present a novel AFA-POMDP framework that jointly learns the task policy and the active feature acquisition strategy with a unified reinforcement learning formalism. We introduce a model-based sequential VAE model to facilitate policy training under partial observability. We demonstrate that imputing missing features via learning model dynamics could significantly benefit policy training with partial observability. Our proposed model, by efficiently synthesizing the sequential information to impute the missing features, can significantly outperform conventional representation learning baselines and leads to policy training with significantly better sample efficiency as well as obtained solutions. Future work may investigate whether our proposed model could be applied to more diverse and complex application domains. Another promising direction is to integrate our framework with model-based planning for further reducing the feature acquisition cost.

heading: ETHICS STATEMENT
text: When deploying machine learning models in real-world applications, the fundamental assumption that the features used during training are always readily available during the deployment phase does not necessarily hold. Our work addresses the aforementioned problem via formulating a novel AFA-POMDP framework that extends the conventional instance-wise non-time-dependent active feature acquisition task to a more challenging time-dependent sequential decision making task. The sequential active feature acquisition module enables the decision making to be performed in a more cost-efficient way when partial features are accessed only during model deployment.
Considering that the task of learning and applying machine learning models is rather problem specific, it is unlikely that our method can equally benefit all possible application scenarios. We also fully acknowledge the existence of risk in applying our model in sensitive and high risk domains, e.g., healthcare, and its potential bias if the model itself or the used representations are trained on biased data. In high risk settings, human supervision of the proposed model might be desired and the model is suggested to be mainly used for decision support systems. To alleviate the reliance on fully observed data during representation learning, it is very promising to trigger follow-up works studying data efficient sequential autoencoder training paradigms.

heading: APPENDIX
text: This appendix is organized as follows:
• Sec A: the detailed algorithm. • Sec B: experimental settings and additional results on the Bouncing Ball domain. • Sec C: experimental settings and additional results on the Sepsis domain.

heading: A RL WITH ACTIVE FEATURE ACQUISITION ALGORITHM
text: Algorithm 1 RL with Active Feature Acquisition 1: Input: learning rate α > 0, dataset D 2: Initialize RL policy πf , πc, VAE parameters θ, φ. 3: Train VAE on dataset D using Eq (5). 4: while Not Converge do 5: Reset the environment. 6: Initialize null observation xp1 = Ø, feature acquisition action a f 0 and control action a c 0.
7: for i = 1 to T do 8: Compute representation with VAE: bt = qφ(x p ≤t,a<t).
9: Sample a feature acquisition action aft ∼ πf (bt) and a control action act ∼ πc(bt). 10: Step the environment and receive partial features, reward and terminal: xpt+1, rt, term ∼
env(aft ,a c t) 11: Compute cost ct = ∑ i c · I(a f(i) t ). 12: Save the transitions {bt,aft ,act , rt, ct, term}. 13: if term then 14: break 15: end if 16: end for 17: Update πf , πc using the saved transitions with an RL algorithm under learning rate α. 18: end while
B BOUNCING BALL+

heading: B.1 TASK SPECIFICATION
text: The task consists of a ball moving in a 2D box of size 32×32 pixels. The radius of the ball equals to 2 pixels. At each step, a binary image is returned as an observation of the MDP state. At the beginning of every episode, the ball starts at a random position in the upper left quadrant (sampled uniformly). The initial velocity of the ball is randomly defined as follows: ~v = [Vx, Vy] = 4 · ~̃v/‖~̃v‖, where the x- and y-component of ~̃v are sampled uniformly from the interval [−0.5, 0.5]. There is a navigation target set at (5, 25) pixels, which is in the lower left quadrant. The navigation is considered to be successful if the ball reaches the specified target location within a threshold of 1 pixel along both x/y-axis.
The action spaces is defined as follows. There are five task actions Ac:
• Increase velocity leftwards, i.e., change Vx by −0.5 • Increase velocity rightwards, i.e., change Vx by +0.5 • Increase velocity downwards, i.e., change Vy by +0.5 • Increase velocity upwards, i.e., change Vy by −0.5 • Keep velocities unchanged
The maximum velocity along the x/y-axis is 5.0. The velocity will stay unchanged if it exceeds this threshold. The feature acquisition action af ∈ Af is specified as acquiring the observation of a subset
of the quadrants (this also includes acquiring the observation of all 4 quadrants). Thus, the agent can acquire 0− 4 quadrants to observe. Each episode runs up to 50 steps. The episode terminates if agent reaches the target location.
B.2 IMPLEMENTATION DETAILS
For all the compared methods, Zero-Imputing (Nazabal et al., 2018) is adopted to fill in missing features with a fixed value of 0.5.
End-to-End The end-to-end model first processes the imputed image by 2 convolutional layers with filter sizes of 16 and 32, respectively. Each convolutional layer is followed by a ReLU activation function. Then the output is passed to a fully connected layer of size 1024. The weights for the fully connected layer are initialized by orthogonal weights initialization and the biases are initialized as zeros.
NonSeq-ZI The non-sequential VAE models first process the imputed image by 2 convolutional layers with filter sizes of 32 and 64, respectively. Each convolutional layer is followed by a ReLU activation function. Then the output passes through a fully connected layer of size 256, followed by two additional fully connected layers of size 32 to generate the mean and variance of a Gaussian distribution. To decode an image, the sampled code first passes through a fully connected layer with size 256, followed by 3 convolutional layers with filters of 32, 32, and nc and strides of 2, 2 and 1, respectively, where nc is the channel size that equals to 2 for the binary image. There are two variants for NonSeq-ZI: one employs the partial loss that is only for the observed variables; the other employs the full loss that is computed on all the variables, i.e., the ground-truth image with full observation is employed as the target to train the model to impute the missing features. The hyperparameters for training NonSeq-ZI are summarized in Table 2.
Seq-PO-VAE (ours) At each step, the Seq-PO-VAE takes an imputed image and an action vector of size 9 as input. The imputed image is processed by 3 convolutional layers with filter size 32 and stride 2. Each convolutional layer employs ReLU as its activation function. Then the output passes through a fully connected layer of size 32 to generate a latent representation for the image fx. The action vector passes through a fully connected layer of 32 to generate latent representation for the action fa. Then the image and action features are concatenated and augmented to form a feature vector fc = [fx, fa, fx ∗ fa], where [·] denotes concatenation of features. Then fc is fed to fully connected projection layers of size 64 and 32, respectively. The output is then fed to an LSTM module, with latent size of 32. The output ht of LSTM is passed to two independent fully connected layers of size 32 for each to generate the mean and variance for the Gaussian distribution filtered from the sequential inputs. To decode an image, the model adopts deconvolutional layers that are identical to those for NonSeq-ZI. The hyperparameters for training Seq-PO-VAE are shown in Table 2.
LSTM-A3C We adopt LSTM-A3C (Mnih et al., 2016) to train the RL policy. The policy takes the features derived from the representation learning module as input. For the VAE-based methods, the input features are passed through a fully connected layer of size 1024. Then the features are fed to an LSTM with 1024 units. The output of the LSTM is fed to three independent fully connected layers to generate the estimations for value, task policy and feature acquisition policy. We adopt normalized column initialization for all the fully connected layers and the biases for the LSTM module are set to be zero.

heading: B.3 DATA COLLECTION
text: To train the VAEs, we prepare a training set that consists of 2000 trajectories. Half of the trajectories are derived from a random policy and the other half is derived from a policy learned from end-to-end method. To train the end-to-end method, we employ a cost of 0.01 over first 2m steps and then increase it to 0.02 for the following 0.5m steps. All the VAE models are evaluated on a test dataset that has identical size and data distribution as the training dataset. We present the best achieved task performance of the data collection policy (End-to-End) and our representation learning approach in Table 5. We notice that our proposed method, by employing an advanced representation model, leads to significantly better feature acquisition policy than End-to-End (smaller number of observations while achieving similar or better reward).
B.4 IMPUTING MISSING FEATURES VIA LEARNING MODEL DYNAMICS
We present an illustrative example to demonstrate the process of imputing missing features and the role of learning model dynamics. To this end, we collect trajectories under an End-to-End policy (the choice of the underlying RL policy is not that important since we just want to derive some trajectory samples for the VAE models to reconstruct) and use different VAE models to impute the observations.
From the results presented in Figure 9, we observe that under the partially observable setting with missing features, the latent representation derived from our proposed method provides abundant information as compared to only using information from a single time step and thereby offers significant benefit for the policy model to learn to acquire meaningful features/gain task reward.
B.5 INVESTIGATION ON COST-PERFORMANCE TRADE-OFF
We perform a case study on investigating the cost-performance trade-off for each representation learning method, presented in Figure 9. Apparently, as we increase the cost, the explorationexploitation task becomes more challenging and each compared method has its own upper bound on the cost above which it fails to learn an effective task policy while acquiring minimum observation. First, we notice that the End-to-End model takes a long time to progress in learning task skills, while the VAE-based models can progress much faster. Among the VAE-based methods, we notice that our proposed method (Figure 9(d)) can achieve as low as 8 observations whereas the baselines NonSeq-ZI (Full) (Figure 9(b)) and NonSeq-ZI (partial) (Figure 9(c)) achieve a standard of ∼20 (lowest point among the solid lines). Thus, we could conclude that our proposed approach can significantly benefit the cost-sensitive policy training and lead to a policy which acquires much fewer observations while still succeeding in terms of task performance.

heading: C SEPSIS MEDICAL SIMULATOR

heading: C.1 TASK SPECIFICATIONS
text: For this task we employ a Sepsis simulator proposed in previous work (Oberst & Sontag, 2019). The task is to learn to apply three treatment actions for Sepsis patients in intensive care units, i.e., Ac = {antibiotic, ventilation, vasopressors}. At each time step, the agent selects a subset of the treatment actions to apply. The state space consists of 8 features: 3 of them specify the current treatment status; 4 of them specify the measurement status in terms of heart rate, sysBP rate, percoxyg stage and glucose level; the remaining one is a categorical feature indicating the patent’s antibiotic status. The feature acquisition actively selects a subset among the measurement features for observation, i.e., Af = {heart rate, sysBP rate, percoxyg state, glucose level}. The objective for learning a active feature acquisition strategy is to help the decision making system to reduce measurement cost at a significant scale.
C.2 IMPLEMENTATION DETAILS
For all the compared methods, we adopt Zero-Imputing (Nazabal et al., 2018) to fill in missing features. In particular, a fixed value of -10 which is outside the range of feature values is used to impute missing values.
End-to-End The end-to-end model first processes the imputed state by 3 fully connected layers of size 32, 64 and 32, respectively. Each fully connected layer is followed by a ReLU activation.
NonSeq-ZI The VAE model first processes the imputed state by 2 fully connected layers with size 32 and 64, with the first fully connected layer being followed by ReLU activation functions. Then the output is fed into two independent fully connected layers of size 10 for each, to generate the mean and variance for the Gaussian distribution. To decode the state, the latent code is first processed by a fully connected layer of size 64, then fed into three fully connected layers of size 64, 32, and 8. The intermediate fully connected layers employ ReLU activation functions. Also, we adopt two variants for NonSeq-ZI, trained under either full loss or partial loss. The details of the hyperparameter settings used for training are presented in Table 4.
Seq-PO-VAE (ours) At each time step, the inputs for state and action are first processed by their corresponding projection layers. The projection layers for the state consists of 3 fully connected layers of size 32, 16 and 10, where the intermediate fully connected layers are followed by a ReLU activation function. The projection layer for the action input is a fully connected layer of size 10. Then the projected state feature fc and action feature fa are combined in the following manner: fc = [fx, fa, fx ∗ fa]. fc is passed to 2 fully connected layers of size 64 and 32 to form the input to the LSTM module. The output ht of the LSTM is fed to two independent fully connected layers of size 10 to generate the mean and variance for the Gaussian distribution. The decoder for Seq-PO-VAE has identical architecture as NonSeq-ZI. The details for training Seq-PO-VAE are presented in Table 4.
LSTM-A3C The LSTM-A3C (Mnih et al., 2016) takes encoded state features derived from the corresponding representation model as its input. The encoded featuresare fed into an LSTM with size 256. Then the ht for the LSTM is fed to three independent fully connected layers, to predict the state value, feature acquisition policy and task policy. Normalized column initialization is applied to all fully connected layers. The biases for the LSTM and fully connected layers are initialized as zero.

heading: C.3 DATA COLLECTION
text: To train the VAEs, we prepare a training set that consists of 2000 trajectories. Half of the trajectories are derived from a random policy and the other half is derived from a policy learned from the End-to-End method with cost 0.0. All the VAE models are evaluated on a test dataset that consists of identical size and data distribution as the the training dataset. We present the task treatment reward obtained by our data collection policy derived from the End-to-End method and that obtained by our proposed method in Table 5. Noticeably, by performing representation learning, we obtained much better treatment reward as compared to the data collection policy, which demonstrates the necessity of performing representation learning.

heading: C.4 MORE COMPARISON RESULT UNDER DIFFERENT VALUES FOR COST
text: We present additional experiment results that compare our proposed method and the non-sequential baselines under the cost values {0, 0.025}. The results for cost value of 0.01 are shown in the main paper. Overall, under all the cost settings, our method leads to significantly better discharge ratio and task reward compared to the baselines.
Also, we demonstrate the cost-performance trade-off on Sepsis domain. By increasing the value of cost, we could obtain feature acquisition policy that acquires substantially decreased amount of features within each episode.
C.5 ILLUSTRATIVE EXAMPLES FOR MISSING FEATURE IMPUTATION IN Sepsis
We present two illustrative examples in Figure 12 to demonstrate how imputing missing features via learning model dynamics would help the decision making with partial observability in Sepsis domain. The policy training process with partial observability could only access very limited information, due to the employment of active feature acquisition. Under such circumstances, imputing the missing features would offer much more abundant information to the decision making process. From the results shown in Figure 12, our model demonstrates considerable accuracy in imputing the missing features, even though it is extremely challenging to perform the missing feature imputation task given the distribution shift from the data collection policy and the online policy. The imputed missing information would be greatly beneficial for training the task policy and feature acquisition policy.

Review summaries:
This paper proposes a reinforcement learning + representation learning approach for simultaneously learning a control policy and feature acquisition policy in environments where feature observation is costly. The authors formulate an approach for learning time series latent variable models that incorporate information from both observation and action histories. They demonstrate through a series of experiments that their approach leads to better imputation (i.e., filling in missing values) and better rewards.

The authors study the problem of reinforcement learning in environments where the agent can spend some reward in order to gain access to observations. The authors introduce a generalization of a POMDP, which they call an AFA-POMDP for Active Feature Acquisition, that divides the action into two pieces, an action for control and action for feature acquisition. The authors' solution approach begins with fully observed trajectories that are used to train a sequential VAE as an inference model. Then, using the pre-trained VAE, an RL algorithm jointly learns the control and feature acquisition policies.  The experiments are on a synthetic "bouncing ball" task where there are five discrete control actions that hit the ball from different directions and the feature acquisition chooses which quadrants of the space to acquire. There is also a sepsis task with three discrete actions and 8 features that correspond to measurements of the patient.

In this paper the authors propose an approach for simultaneously learning how to explore more efficiently in POMDPs via targeted feature acquisition, and learning a reward-maximizing control policy, balancing the cost of feature acquisition with the expected reward. Learning is done via a VAE framework which combines a belief inference model and an observation decoder, with a key innovation being that inference is done as a sequential process. Results comparing this approach to other variational inference approaches show the proposed framework reaches better performance with lower cost (particularly, number of acquired features).