Reinforcement Learning with Efficient Active Feature Acquisition
Authors: Haiyan Yin
Yingzhen Li
Sinno Pan
Cheng Zhang
Sebastian Tschiatschek
Venue: ICLR 2022 Submitted (2022)
Keywords: Representation Learning
Reinforcement Learning
Active Learning

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-life, 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. Thus it poses a significant challenge for the agent to learn optimal task policy while efficiently reducing the cost for information acquisition. In this paper, we introduce a model-based framework to solve such exploration-exploitation problem during its execution. Key to the success is a sequential variational auto-encoder which could learn high-quality representations over the partially observed/missing features, where such representation learning serves as a prime factor to drive efficient policy training under the cost-sensitive setting. We demonstrate our proposed method could significantly outperform conventional approaches in a control domain as well as using a medical simulator. 

Under review as a conference paper at ICLR 2022
REINFORCEMENT LEARNING WITH EFFICIENT ACTIVE FEATURE ACQUISITION
Anonymous authors Paper under double-blind review
ABSTRACT
Solving real-life sequential decision making problems under partial observability involves an
problem. To be successful, an agent needs efficiently gather valuable information about the state of the world for making rewarding decisions. However, in real world applications, acquiring valuable infor- mation is often highly costly, e-g.. in the medical domain, information acquisition might correspond to performing a medical test on 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 policy which solves problem during its execution. Key to the success is a
this exploration-exploiaaieee 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 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.
INTRODUCTION
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 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- -ideppendedd 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 wellwwrldd applications. For instance, consider 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 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 learn policies which acquires the information required for solving a task in the cheapest way possible. For simplicity, we think of the policy as being constituted of an acquisition policy which actively selects meaningful features to be observed and task policy, which selects actions to change the state of the system towards some goal As such, we consider 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
Clearly, these two policies are not independent in general, e.g., acquiring features can change the state of the system.

---

Under review as a conference paper at ICLR 2022
conventional POMDP (e q. nivigation)
feature acquisition (e.g. medical troimentss
An illustrative figure for com- paring POMDPs considered in conventional RL literature and those in our proposed set- ting. By performing active feature acquisi- tion, the partial observability is determined by the feature acquisition policy besides the evolvement of the system.
Figure 1:
action
action
observe
obseeve
acquire
controlled by the acquisition policy. Thus, the resulting POMDP is different from the conventional POMDPs in RL literature (Cassandra I998) where the partial observability for later stems from a fixed and aetion-independe are only determined by the choice of the task action, whereas in our setting the state-trand sate-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 of the policy needs to solve an exploration-exploittion problem, and thus we often need to learn sophisticated policies.
observation model. Also, the state transitions in conventional POMDPs
problem: each execution
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 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 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 the challenge of the exploration-exploitttoo problem, it non-trivial task 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-exxloiittoor Therefore, framework also tackles the joint policy training task from a representation learning perspective. Specifically, 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, formulate sequential generative model that can efficiently learn model dynamics during representation learning. Overall, the contributions of our paper are three-fold:
proposed
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.
RELATED WORK
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 active feature acquisition and representation learning for POMDP, respectively.
2

---

Under review as a conference paper at ICLR 2022
2.1 ACTIVE FEATURE ACQUISITION
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 one time effort to select subset of features as a whole. One typical example is the conventional linear model that poses sparsity inducing prior distribution to the model (Tibshirani] (996).. Recently, there also emerged approaches that adopt reinforcement learning to actively find optimal feature subsets (Yoon) al. 2018 Shim 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-wist features.
2.2 REPRESENTATION LEARNING IN POMDP
In complex tasks, policies trained upon different representations can 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 end-to- tnn fashion (Mnih et al. 2013; Lillicrap et al. 2016 Mnih et 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 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 an off-line fashion and thus resulting in multi-stage reinforcement learning. Higgins et al. (0166 2017) adopt pretrained VAE models as representation module 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 more representative state information. Recently, there also emerged several works on sequential representation learning for POMDP (Gregor et al. 2019) Vezzani al. 2019.. However, most of the works utilize VAE training as auxiliary task to jointly update the representation model with the policy learning loss. In our work, due to the high acquisition cost 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.
METHODOLOGY
3.1
TASK SETTING
In this section, 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-tin time stochastic control processes generalizing standard POMDPs:
Definition 1 (AFA-POMDP). The active feature acquisition POMDP is a tuple M (S, A, O, C,,,, 7) where S is the state space and A = (AF, is a joint action space of feature acquisition actions Al and control actions A The transition kernel T: S X A x Al + Ps maps any joint action a = (a a²) in state s E S to a distribution Ps over next states. In each state s. when taking action a!, the agent observes x" x(a²), ). i.e,, subset of the features )²") ~ O(s) indicated by al, where (ss) is a distribution over possible feature observation for state and x are features not observed by the agent. When taking a joint action, the agent obtains
x = (x", xx"
feature
3

---

Under review as a conference paper at ICLR 2022
C: rewards according to the reward function R: S A - R and pays cost C: S S A + B- feature acquisition. Rewards and costs are discounted by the discount factor 7 E (0,1).
Simplifying assumptions For simplicity, we assume that x consists of fixed number of features N, for all states, that A' = - is the power set of all the Ny features, and that x²(a() consists of all the features in x indicated by the subset a E A Note that the feature acquisition action for 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 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 the following that it depends only on the control action, i.e., T(s, a a²) = (s,, a", a a for all a a E A'. While not true for all possible applications, this assumption can be 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(a', s) e cla but 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 task policy ne and feature acquisition policy xf. That is, we aim to solve the optimization problem
LA4 R(xx, - Ee. Ec.
max of _
Σ 0=
(1)
[()
where the expectation is over the randomness of the stochastic process and the policies, a; denotes the i-th feature acquisition action at timestep t, and I1 (-) is an indicator function whose value equals to if that feature has been acquired.
Note that the above optimization problem very challenging: an optimal solution needs to maintain beliefs b, 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 x² (a/ | b.) and the task policy i.e., <a by) depend on this belief. The information in the belief itself can be controlled by the feature acquisition policy through querying subsets from the features x, and hence the task policy and feature acquisition policy itself strongly depend on effectiveness of the feature acquisition policy.
3.2 SEQUENTIAL REPRESENTATION LEARNING WITH PARTIAL OBSERVATIONS
We introduce a sequential representation learning approach to facilitate the task of policy training with active feature acquisition. Let X1.7 (x, Xy) and ai.r (a....., ar) denote a sequence of observations and actions, respectively. Alternatively, we also denote these sequences as xęr and agr. Overall, our task of interest is to train sequential representation learning model to learn the distribution of the full sequential observations X1:7, i.e., for both the observed part xf= and the unobserved part x".r- Given only partial observations, we can perform inference only with the observed features xf.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, 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 x xi.7. in our work, we take into account both the observation sequence x";7 and the action sequence a1.7' for conducting inference. The intuition is that since x".r by itself consists very limited information over the agent's 's underlying MDP state, incorporating the action sequence would be an informative add-on to the agent's 's acquired information to infer the belief state. To summarize, our proposed sequential representation model leamss encode xf.r and a into meaningful latent
4

---

Under review as a conference paper at ICLR 2022
21
22
23
Figure 2: Observation decoder and be- lief inference model for the partially ob- servable 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 unob- served features.
hl
X
Inference Decoder
X
features, for predicting x and x".. The architecture of our proposed sequential representation learning model is shown in Figure
Observation Decoder Z1.T (Z),..., Zr) denote sequence of latent states. We consider the following probabilistic model:
Px x", 2)
II pl x x (1) p(z)).
(2)
For simplicity the notations, we assume Zg We impose simple prior distribution over z, i.e., standard Gaussian prior, instead of incorporating some learned prior distribution over the latent space of z, such as an autoregressive prior distribution like p(z, 1 1,,1, a0.c-1). The reason is that using static prior distribution results in latent representation z, 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 x and x depends on the corresponding latent variable 2,. Given Zr, the observed variables are conditionally independent of the unobserved ones. Therefore,
p(x,, x (2)) = p(x²' 22) (x))
During policy training only assume access partially observed data. This requires an inference model which takes in the past observation and action sequences to infer the latent states Specifically, present a structured inference network qø as shown in Figure 2 which has an autoregressive structure:
Belief Inference Model
qs (z a) = II 4 (z act),
(4)
where (-) is 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 forward RNN model as the backbone for the filtering function (-) (Gregor et al. 2019). Specifically, step f, the RNN processes the encoded partial observation x action a and its past hidden state h to update its hidden state h, Then the latent distribution 2, is inferred from h, The belief state b, defined as the mean of the distribution z₁- By accomplishing the supervised learning task, the belicf 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 proposed to pre-train both the generative and inference models offline before learning the RL policies. In this case, assume the access the unobserved features, so that construct a supervised learning task to learn to impute unobserved features. Coneretely, the pre-training task update the parameters 0, o by maximizing the following variational lower-hound (Jordan et al. 1999 Kingma & Welling 2013):
log x(xxxx")) 2 Eyg
Σ
log po (x{(z.)) + log (o( (LogoZz |x" aci) p(()))
Ex ,x))
(5)
5

---

Under review as a conference paper at ICLR 2022
By incorporating the term log (x (z),, the training of sequential VAE generalizes from unsuper- vised task to 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.
4 EXPERIMENTS
We examine the characteristics of our proposed model in the following two experimental domains: a bouncing ball control task with high-dimensiioal image pixels as input, adapted from (Fraccaro)
al. 2017; sepsis medical simulator fitted from real-world data (Oberst & Sontag 2019).
Baselines For comparison, we mainly consider variants the strong VAE baseline beta-VAL (Hig gins et al. 2016), which works cen-time-deeedeee data instances. For representing the missing features, we adopt the pero-imputing method, proposed in (Nazabal et al. 2018) over the unobserved features. Thus, denote the VAE baseline as NonSeq-ZI. We train the VAÊ with either the full loss over the entire features, or the partial loss which only applies to the observed features (Ma) al. 2019). We denoted proposed sequential VAE model for POMDPs Seq-PO-Ver All the approaches adopt an identical policy architecture, Detailed information on the model architecture presented in appendix.
Data Collection To pre-train the VAE models, data generated by a non-random policy is unavoid- ably needed to incorporate abundant dynamics information. For both tasks, we collect small scale dataset of 2000 trajectories, where half the data is collected from a random policy and the the other half from 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 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 al. 2017) by adding a navigation objective and introducing control actions. Specifically, a ball moves in a 2D box and each step, binary image of size 32 x 32 showing the box and the ball. Initially, the random position in the upper left quadrant, and has random velocity. The objective to control the ball to reach a fixed target location set at (5, 25). We incorporate five RL actions: a direction with
ball appears
null action and four actions for changing the velocity of the ball in either the fixed scale: (AV; : +05, AV,- +0.5, null). reward of 1.0 is issued if the ball reaches its target location. Each episode runs up to 50 time steps.
Table 1: Imputing loss evaluated on missing features for Bouncing Ball- and Sepsis. Each setting is evaluated for 3 random seeds.
BOUNCING BALL (NLL)
SEPSIS (MSE)
VAE MODEL
0.5846 (#: 0.0544)
1.1954
NONSEQ ZI (PARTIAL)
(=0.2521)
0.0728
0.5217
NONSEQ-ZI (FULL)
(+:: 0.0004)
(+ 0.0016)
0.0284
0.1724
SEQ-PO- VAE (OURS)
(+ 0.0002)
(=0.0064)

---

Under review as a conference paper at ICLR 2022
di End to End
Episedic Features
rewards (weo CDS3
1,0
rewards cost
60
0.8
0.6
40
0.4
30
Nundeq di (sarbat)
Nondee
0.2
0.0
10
20
30 Steps (1e4)
40
50
10
20
30
40
50
steps
Steps (1ef)
Figure 3: Performance curves on the bouncing ball+ domain: episodic number of observations acquired by the x b: task rewards w/o cost. Our proposed method outperforms the non-sequential baselines in learning the task as well as acquiring less observations; c: Ablation study on bouncing ball' to illustrate the effect of learning the feature acquisition policy. Each method is run with 10 random seeds. Our proposed approach outperforms the random baseline significantly in terms of task performance.
0
Figure 4: Seq-PQ-Vel VAF enlarged). Each block of three rows corresponds to the results for one trajectory. In each block, the three rows (top-down) correspond to: (1) the partially observable input selected by acquisition policy; (2) the ground-truth full observation; (3) reconstruction from Seq-PO- VAE. The green boxes remark the frames where ball is not observed but model could impute its location. Key takeaways: (1) our learned acquisition policy captures model dynamics (2) Seg-PO-VAL VAE effectively impute the missing features (i.e., ball can be reconstructed even when they are unobserved from consequent frames).
PO-VAE reconstruction for the online trajectories upon convergence (better to view
Representation Learning Results We evaluate the missing feature imputing performance each VAE model in terms of negative log likelihood (nll) loss and present results in Table 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 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, the target navigation skill at a much faster pace than the non-sequential baselines. Our method also converges standard where much less feature acquisition is required to perform the task.
also notice that the joint policies trained by Seq-PO-VAL can develop
7

---

Under review as a conference paper at ICLR 2022
0.49
Manlieg (llonett)) Tei-it-Cht
0.3
0.4
0.35
cuts
0.2
discharge
0.3
roward (wio
0.2
0.29
0.1-
0.15
NonSee (pariia
End and
0.10
-0.3
End co End
-04
200
400
600
800
1000
200
4000
850
1000
600 (1c4)
800
1000
steps (1e4)
steps (1e4)
Figure 5: Performance curves for the compared approaches on Sepsis. The curves are derived under cost value of 0.01. Overall, our method converges treatment policy with substantially better reward compared to the baselines.
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 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 acquire meaningful features. We also show the actively learned feature acquisition policy works better than random acquisition. From the results in Figure (c), our method converges to better standard than random policies with considerably high selection probabilities.
4.2 SEPSIS MEDICAL SIMULATOR
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.c. Lantibiotic, ventilation, vasopressors). The state space consists of 8 features: 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 [ 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 Overall, our proposed method results in substantially better task reward compared all baselines. Noticeably, the learning of discharge for our method progresses significantly faster than baseline approaches and converges 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 metries, 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 sequen- tial feature acquisition strategy with RL. To this end, compare method with baseline that randomly acquires features. We evaluate our method under different cost values, and the results are shown in F From the results, notice that there is clear port-performnnce trade-off, i.e.,
8

---

Under review as a conference paper at ICLR 2022
NonSee Zi(ful) (i(full NonSea Zi(partial)
2.4-
Reward (w/o cost.)
Total Acquisition
End-to-End
4.40
Ours
1.0
0.8
c=00
0.6
50%
c=0.01
8.20
c=0.02
c=0,025
0.2
C-0.05
Random
0.0
-0.6
-0.4
-0.2
0.0
0.2
0.4
Episodic No. Features
Reward (w/o cost)
Figure 6: The task reward and episodic no. fea- ture acquisitions derived upon convergence for our method, when training our model under dif- ferent cost values. We also present the results for baseline that trains the task policy under a random feature acquísition strategy. The actively learned feature acquisition leads to significant advantage compared random acquisition.
Figure This figure shows the total number of features acquired (y-value) to obtain certain re- ward standard (x-value), when training the task policy using different representation learning ap- proaches. The presented methods adopt a cost 0.01. Our method results in the best achiev- able reward (max x-value) with lowest slope.
higher feature acquisition cost results in feature acquisition policies that obtain fewer observations, with 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, acquire only about half of the total number of features (refer to the value derived by Random-10 100%) to obtain comparable task performance. Also, we notice that the specified cost has very clear impact 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.
Impact on Total Acquisition Cost For different representation learning methods, also investigate the total number of features acquired at different stage of training. The results are shown in Figure 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 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 (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.
Ablation:
CONCLUSION
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, 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 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.

---

Under review as a conference paper at ICLR 2022
ETHICS STATEMENT
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 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 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.
trained on biased
REFERENCES
Anthony R Cassandra. A survey of pomdp applications. In Working notes of AAAI 1998 fall symposium on planning with partially observable Markov decision processes, volume 1724, 1998.
Abhishek Das, Satwik Kottur, José MF Moura, Stefan Lee, and Dhruv Batra. Learning cooperative visual dialog agents with deep reinforcement learning. In Proceedings of the IEEE international conference on computer vision, pp. 2951- 29600, 2017.
Chelsea Finn, Sergey Levine, and Pieter Abbeel. Guided cost learning: Deep inverse optimal control via policy optimization. In International conference on machine learning, pp. 49-58, 2016.
Marco Fraccaro, Simon Kamronn, UIrich Paquet, and Ole Winther. A disentangled recognition and nonlinear dynamics model for unsupervised learning. In Advances in Neural Information Processing Systems, pp. 3601-3610,, 2017.
Wenbo Gong, Sebastian Tschiatschek, Sebastian Nowozin, Richard E. Turner, José Miguel Hernández- Lobato, and Cheng Zhang. Icebreaker: Element-wise efficient information acquisition with bayesian deep latent gaussian model. In Annual Conference on Neural Information Processing Systems, pp. 14791-14800 2019.
Karol Gregor, George Papamakarios, Frederic Besse, Lars Buesing, and Theophane Weber. Temporal difference variational auto-encoder. ICLR, 2019.
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. In ICLR, 2016.
Irina Higgins, Arka Pal, Andrei Rusu, Loie Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerchner. Darla: Improving zero-shot transfer in reinforcement learning. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pp. 1480-1490, 2017.
Michael Jordan, Zoubin Ghahramani, Tommi Jaakkola, and Lawrence K Saul. An introduction variational methods for graphical models. Machine learning, 37(2):183-233, 1999.
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint
urXiv:1312.0 14, 2013.
Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. ICLR, 2016.
10

---

Under review as a conference paper at ICLR 2022
Yuan Ling, Sadid A Hasan, Vivek Datla, Ashequl Qadir, Kathy Lee, Joey Liu, and Oladimeji Farri. Learning to diagnose: assimilating clinical narratives using deep reinforcement learning. In Proceedings of the Eighth International Joint Conference Natural Language Processing (Volume I: Long Papers), pp. 895-905, 2017.
Chao Ma, Sebastian Tschiatschek, Konstantina Palla, José Miguel Hernández-Lobato, Sebastian Nowozin, and Cheng Zhang. EDDI: efficient dynamic discovery of high-value information with partial VAE. In Proceedings of the 36th International Conference on Machine Learning, pp.
4234 4243, 2019.
Stefan Mathe, Aleksis Pirinen, and Cristian Sminchisescu. Reinforcement learning for visual object detection. In Proceedings the the IEEE Conference on Computer Vision and Pattern Recognition, pp. 2894-2902, 2016.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:13/2.5602, 2013.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Belle- mare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, loannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518(7540):529- 533, 2015.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. International conference on machine learning, pp. 1928-1937, 2016.
Alfredo Nazabal, Pablo M Olmos, Zoubin Ghahramani, and Isabel Valera. Handling incomplete heterogeneous data using vaes. arXiv preprint
2018
Michael Oberst and David Sontag. Counterfactual off-policy evaluation with gumbel-max structural causal models. In ICML, 2019.
Yu-Shao Peng, Kai-Fu Tang, Hsuan-Tiees Lin, and Edward Chang. Refuel: Exploring sparse features in deep reinforcement learning for fast disease diagnosis. In Advances in Neural Information Processing Systems, pp. 7322-7331, 2018.
Hajin Shim, Sung Ju Hwang, and Funho Yang. Joint active feature acquisition and classification with variable-size set encoding. Advances Neural Information Processing Systems, pp. 1368-1378, 2018.
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schritwieser, loannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484, 2016.
Robert Tibshirani. Regression shrinkage and selection via the lasso. Journal of the Royal Staristical Society: Series B (Methodological), 58(1):267-22 1996.
Giulia Vezzani, Abhishek Gupta, Lorenzo Natale, and Pieter Abbeel. Learning latent state representa- tion for speeding up exploration. arXiv preprint arXiv: 1905. /2621, 2019.
Jinsung Yoon, James Jordon, and Mihaela van der Schaar. using neural networks. In ICLR, 2018
Invase: Instance- variable selection
Sara Zannone, José Miguel Hernández-Lobate Cheng Zhang, and Konstantina Palla. Odin: Optimal discovery of high-value information using model-based deep reinforcement learning. In ICML Real-world Sequential Decision Making Workshop. 2019.
Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew Johnson, and Sergey Levine. Solar: Deep structured latent representations for model-based reinforcement learning. In Proceed- ings of the 35th International Conference on Machine Learning, 2018
11

---

Under review as a conference paper at ICLR 2022
APPENDIX
This appendix organized as follows:
Sec A the detailed algorithm.
Sec B
experimental settings and additional results the Bouncing Ball domain.
See
experimental settings and additional results on the Sepsis domain.
A
RL WITH ACTIVE FEATURE ACQUISITION ALGORITHM
Algorithm 1 RL with Active Feature Acquisition
1: Input: learning rate 0, dataset D
2: Initialize RL policy f, Re. VAE parameters 8, Ø.
3: Train VAE on dataset D using Eq (5)
4: while Not Converge do
5: Reset the environment.
6: Initialize null observation x = Ø, feature acquisition action a and control action ag.
7: for to do
8: Compute representation with VAE: b, g0 9o ace).
9: Sample feature acquisition action a- ~ y (b:) and a control action a. ~ .T((b)).
Step the environment and receive partial features, reward and terminal: x X Tty term cnv(a af)
Compute cost c E, cI(((( (()) 12: Save the transitions {b,, af af, r2 C2 term].
if term then 14: break end if
16: end for
Update y. ne using the saved transitions with an RI. algorithm under learning rate a. end while
18:
B BOUNCING BALL+
B. TASK SPECIFICATION
The task consists of ball moving in 2D box of size 32 x 32 pixels. The radius of the ball equals to 2 pixels. At each step, 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.. VyJ = 4 0/ where the and y-component of 7 are sampled uniformly from the interval |-0.5, 0.5). There is 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 threshold of pixel along both x/y-axis.
The action spaces is defined follows. There are five task actions A
Increase velocity leftwards, i.e., change Vz by -0.5
Increase velocity rightwards, i.e., change Vg by +0.5
Increase velocity downwards, i.c., change V,, by +0.5
Increase velocity upwards, i.e., change V,, by -0.5
Keep velocities unchanged
The maximum velocity along the x/y-axis is 5.0. The velocity will stay unchanged it exceeds this threshold. The feature acquisition action a A is specified acquiring the observation of subset
12

---

Under review as a conference paper at ICLR 2022
of the quadrants (this also includes acquiring the observation of all 4 quadrants). Thus, the agent can acquire - 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 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 ReLU activation function. Then the output 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 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 generate the mean and variance of a Gaussian distribution. To decode image, the sampled code first passes through fully connected layer with size 256, followed by 3 convolutional layers with filters of 32, 32, and ne and strides of 2, 2 and 1, respectively, where nc is the channel size that equals to 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 employed as the target to train the model to impute the missing features. The hyperparameters for training NonSeq-ZI are summarized in Table
all the variables, i.e., the ground-truth -ruth image with full observation is
Hyperparameter
3 (KI. weight)
KL. reduction
Loss reduction
learning rate
epochs
NonSeq-Z4 ZI (partial)
1.0
sum sum
lc-4
1k
NonSeq-ZI (full)
1.0
sum sum
le-4
1k
Seq-PO-VAE PAE (ours)
1.0
sum
5e-4
2k
Table 2: Hyperparameter settings for training VAE models the Bouncing Ball* dataset.
Seq-PO-VAE (ours) At each step, the Sog-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 generate latent representation for the image f, 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 f. [fz, fa, fg * fa]. where [:] denotes concatenation of features. Then f. 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 h, 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 image, the model adopts deconvolutional layers that identical to those for NonSeq-ZI. Zi. The hyperparameters for training Seq-PO-VAL are shown in Table 2
LSTM-A3C We adopt LSTM-A3C (Mnih et al. (0166) 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 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.
13

---

Under review as a conference paper at ICLR 2022
B.3 DATA COLLECTION
To train the VAEs, we prepare a training set that consists of 2000 trajectories. Half of the trajectories derived from a random policy and the other half is derived from 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 0.02 for the following 0.5m steps. All the VAE models are evaluated on 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-Enut (o-End) and our representation learning approach in Ta [ 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).
dataset
Model
End-to-End
Ours
Average # of observations per episode
17.94
8.24
Task reward
1.0
1.0
Table 3: Task performance for the data collection policy and our method on Bouncing Ball*
14

---

Under review as a conference paper at ICLR 2022
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-io-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 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 compared to only using information from single time step and thereby offers significant benefit for the policy model learn to acquire meaningful features/gain task reward.
Figure 8: Imputation results for different VAE models. We select 9 trajectories obtained from the trained End-to-End policy. Each block corresponds the results for one trajectory (better view enlarged). The five rows in one block are (top-down): (1) partial observations acquired by the agent; (2) ground-truth image with full observation; (3) Imputation by NoSeg ZI (partial); (4) Imputation by NoSeq-ZI ZI (full): (5) Imputation by Seq-PO- VAE (ours). Our model can often successfully predict the balls location even if it is not present in the acquired observation. Hence successfully employs its learned knowledge of the dynamics. In contrast, the non-sequential model (obviously) fails to predict the balls location when the ball not present the observation.
15

---

Under review as a conference paper at ICLR 2022
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 increase the cost, the exploration- exploitation task becomes more challenging and each compared method has its own upper bound on the cost above which it fails learn an effective task policy while acquiring minimum observation. First, we notice that the End-to-Ene End model takes 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(c) can achieve as low as observations whereas the baselines NonSeq- - (Full) (Figure 9()) and NonSeq-Zi ZI (partial) (Figure achieve a standard of ~00 (lowest point among the solid lines). Thus, we could conclude that 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.
1.0
Exsib-Erdicn IOVE ED
apiiodic num. obserooaaions
100
Rnd EST
Fusta (ne
60
rewards
0.0
50
Stops (164)
100 125 150
175
25
30
75 100 125
staps (1o4)
150
175 200
(a) End-to-Enu
Nor Sea (0H0000
SpI num. doservatooss
Nardeazh
50
Mendios 25
rewordds (WN
0.66
10
20
30
40
60 70
stops (1o4)
52005
(104)
(b) NonSeq-7I (full)
Farldd EED
Nor See Ferid tovll ESS
21
E35
rewards Dele
0.6.
[40
0.4
20
0-
40
(1e6)
(1e4)
(c) NonSeq Zl (partial)
apiedic abservation
ruwards (wil 0094)
50
oo DAI
0.6
UNE G-- DU!
30
0.0
po une und (61)
30
20
30
40
50
60
10
20
30
40
50
60
70
Staps (1e4)
Slaps
(104)
(d) Scq-PO-VAE (ours)
Figure 9: Cost- cost-performance trade-off investigation. Each row corresponds to the performance in terms of task reward and episodic number of acquisitions obtained for specific method (see legends). Each curve is derived from independent runs. We use dotted lines to indicate those instances when the task learning does not always succeed. Thus, the best achievable number of observations should be referred to as the lowest curve among the solid lines.
16

---

Under review as a conference paper at ICLR 2022
C SEPSIS MEDICAL SIMULATOR
C.I TASK SPECIFICATIONS
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., A = (antibioti ventilation, vasopressors). At each time step, the agent selects a subset of the trealment actions to apply. The state space consists of 8 features: 3 of them specify the current trealment 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.c., A {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 significant scale.
C.2 IMPLEMENTATION DETAILS
For all the compared methods, we adopt Zero-Imputinn (Nazabal et al.) 2018) to fill in missing features. In particular, fixed value of -10 which 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 Rel.U activation.
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 partial loss. The details of the hyperparameter settings used for training are presented in Table 4
Hyperparameter
3 (KL weight)
KL reduction
Loss reduction
learning rate
epochs
NonSeq-ZI (partial)
0.01
sum sum sum
sum
le-4
lk
NonSeq-ZI ZI (full)
0.01
sum
le-4
1k
Seq-PO- VAE (ours)
0.01
sum
le-3
1k
Table 4: Hyperparameter settings for training VAE models on the Sepsis dataset.
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 f. and action feature f, are combined in the following manner: f [fz, fa, f. f. f. passed fully connected layers of size 64 and 32 to form the input the LSTM module. The output h, of the LSTM is fed to two independent fully connected layers of size 10 generate the mean and variance for the Gaussian distribution. The decoder for Seq-PO- Pip--O-VAA V-E has identical architecture as NonSeq-ZI. ZI. The details for training Seq-PO-VAL VVE are presented in Table 4
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 h, 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 applied to all fully connected layers. The biases for the LSTM and fully connected layers are initialized as zero
LSTM-A3C
17

---

Under review as a conference paper at ICLR 2022
C.3 DATA COLLECTION
To train the VAEs, we prepare a training set that consists of 2000 trajectories. Half of the trajectories derived from random policy and the other half is derived from a policy learned from the End-to-Enu 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 Pad-to-Erme method and that obtained by our proposed method in Table Noticeably, by performing representation learning, we obtained much better treatment reward compared to the data collection policy, which demonstrates the necessity of performing representation learning.
Model
End-to-End
Ours
Treatment Reward
0.35
0.45
Table 5: Task performance for the data collection policy and our proposed method Sepsis.
C.4 MORE COMPARISON RESULT UNDER DIFFERENT VALUES FOR COST
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.
0.25
RetSee Ziperti))
reward (wio cost)
0.4
mortalty
0.3
0.15
0.1
0.2
0.0
0.2
Nanden D110 Out
0.13
-01
NanSoou-Atill Nossing
0.05
-02
Gurt
0.0
1002
630
(164)
(a) Cost=0
0.65
SIME Randu
0.53
reward (eel cost)
0.30
dischange
0.25
nortolly
0.20
0.25
0.15
0.23
-2.2
Honding Zarel Norday Zaparto
Ninieen 7154
200
400
500
S00
1000
209
400
630
800
1000
200
400
800
1000
steps (1e4)
steps
(1e4)
steps
(b) Cost=0.025
Figure 10: Comparison result between our proposed method and the non-sequential VAE baseline models under different values for cost. Each curve is derived from 3 independent runs.
18

---

Under review as a conference paper at ICLR 2022
Also, demonstrate the restooerforaece 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.
50
Gura (contr10) (0
Ours (cost=B.0)) 0.0
45
40
35
100
25
20
15
10
200
600
1000
steps (1o4)
Figure 11: Average number of observations acquired in each episode when training proposed model under different cost values.
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 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 [2] 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.
heart rate
2.22
Sys0P
percogya
percoxyg
glurase
Blccose
IutS inputed
28
steps
steps
Figure 12: Two example trajectories for illustrating how our method works on the Sepsis medical domain. The acquisition policy trained with cost of Each block corresponds to one trajectory and the four rows correspond to the four measurement features being considered for active feature acquisition. Each dot indicates the employment of feature acquisition on the corresponding measure- ment feature the presented time point. In each trajectory, demonstrate the ground-trun -tuth signal over time as well as the imputed signal over time predicted by our proposed Seq-PO- PO-VAE model. By imputing the missing features via learning model dynamics, proposed method could offer much more informative representation for the policy training compared to the non-sequential VAE baselines, and thus significantly benefit the policy training with partial observability.
19


Reviews:
[{"review_id": "NvjtzsOvBlp", "reviewer": "Reviewer_82uH", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- The authors approach outperforms the other VAE baselines, both in terms of reward, and in terms of MSE in inferring the unobserved state.\n- The problem of joint control and feature acquisition is interesting and important.\n\nWeaknesses:\n- The authors claim that AFA-POMDPs generalize POMDPs is false—an AFA-POMDP is a special case of a POMDP. An example of past work using POMDPs for feature acquisition can be found here: Shi and Cain. Cost-sensitive feature acquisition and classification. 2007.\n- The authors approach of using a neural network-based model to infer the belief state in a POMDP was done here: Karkus et al. QMDP-Net: Deep Learning for Planning under Partial Observability, 2017. \n- The domains seem very simple, like they probably could be solved with planning approaches like QMDP-net. They don't really show off the benefit of using RL.\n- Despite having much higher MSE, the baselines are very competitive on the tasks. NonSeq-ZI (partial) is almost as good as the trained method on bouncing ball and NonSeq-ZI (full) performs better in terms of mortality on sepsis, but this is not mentioned.\n- The authors approach makes a quite strong assumption about having access to fully observed data for training. The paper doesn't really discuss this as a limitation.\n\nPost-response: The authors have addressed some of my concerns, and I have raised my score to a 5. I still believe the paper is not ready for publication.", "summary_of_the_paper": "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.\n\nThe 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.", "summary_of_the_review": "The paper makes some incorrect claims, does not cite relevant past work, and does experiments on very simple problems.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "MRQWltD2jCK", "reviewer": "Reviewer_APEd", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This is an interesting paper which tackles a variation on the perennial challenge of exploration vs. exploitation in POMDPs. The overall approach seems reasonable to me and I liked what I identify to be two key ideas in the paper, which are learning the feature acquisition policy *and* the target policy simultaneously, and introducing the notion of sequence in the VAE inference framework. However, I do have some questions, \n\n- My first concern regarding the proposed approach is that it seems wildly expensive to train even for modestly sized problems, and even in the fraught landscape of POMDP learning strategies. I have not seen concrete estimations of sample complexities required in training in the supplementary material, how do they compare to other approaches for solving POMDPs (more on that in a bit). On that note, I am somewhat concerned by the pretraining of the VAEs and the fact that it is then essentially fixed during learning, implying that while there's a sequential component to it, it isn't really an RL approach for hidden state imputation and therefore might be very brittle in practice, especially when a lot of distributions are unknown prior to training. \n\n- Speaking of other approaches for solving POMDPs, a problem which can be spotted already in the related work section is the fact that the author restrict themselves essentially to a small pool of works which pursue a POMDP solving strategy that is largely similar in spirit to that proposed in this paper. While this focus is understandable, it not only does a disservice to a great deal of prior work on efficient exploration and representation learning in POMDPs, it also substantially limits the set of useful baselines the authors compare against. What about other strategies for learning behavior policies and work on balancing exploration vs. exploitation in POMDPs in general? Or hierarchical abstractions? The current set of experiments seems to mostly show that introducing the notion of sequence in the VAE formulation is beneficial, which is neither negligible nor surprising, but it's not obvious to me that the proposed (again, monstrously complex) approach outperforms simpler approaches under most settings.\n\n- The term \"cost\" here mostly refers to number of features acquired. However, it seems to me that such an approach is overly simplistic, given that certain features are vastly more expensive to acquire than others (consider the difference between a simple blood test and a 24 hour EEG monitoring experiment). Is the proposed approach resilient to varying acquisition costs?\n", "summary_of_the_paper": "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).", "summary_of_the_review": "I find this paper very interesting but am worried about the complexity of the proposed approach and how difficult it would be to apply in most cases. I am also concerned that the baselines compared against are insufficient to get an actual sense of how well this approach compares to other approaches, and what's the sample complexity and computational cost (rather than just feature acquisition cost) associated with it. I am also concerned that essentially learning the VAE component offline makes the approach brittle in practice. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "GIhsuyqNjGA", "reviewer": "Reviewer_jW1b", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n- The problem is an important one: feature acquisition is indeed a major problem in the healthcare space, for example, and the ability to learn policies that effectively trade off information gathering/feature acquisition costs with environment control.\n- The proposed VAE model seems to learn what it's supposed to learn - Figure 4 was a nice way to qualitatively capture exactly how the model is trading off feature acquisition costs with control rewards.\n- \n\nWeaknesses:\n- My main concern is that there is already a fairly similar work by Igl et al (2018) [1] that proposed a sequential VAE and had some promising results on a related set of tasks. I would expect to see their DVRL algorithm as well as something like \"Deep Recurrent Q-Learning\" [2] implemented as baselines before I agree with any claims around significant contributions or novelty.\n- There are a few critical experimental details that are left ambiguous. For example, after reading through the manuscript and searching explicitly for information on the \"end-to-end\" baseline, it is not clear to me exactly what model was used and how it selectively acquired features. \n- The requirement that the VAE be pretrained offline is a fairly restrictive one. It's not clear to me how realistic this is. Also, does that preclude any improvement to the VAE model while the agent is interacting with the environment online? That seems like a missed opportunity.\n- There was little, if any, discussion about the potential ethical impacts of the work. This would be a requirement, in my mind, for acceptance.\n\nMinor comments:\n- There are a number of minor grammatical errors throughout. The paper would benefit from some detailed proofreading. I tried to keep track of the errors but it quickly got out of hand.\n- Citations are formatted awkwardly (use \\citep rather than \\citet or \\cite).\n- I believe $\\mathcal{A}_f$ in equation (1) should be $\\mathcal{A}^f$ in order to be consistent with prior notation.\n- It wasn't fully clear to me based on the text what the difference is between the \"full loss\" and the \"partial loss\" (I understand that one is over \"the entire features\" and the other \"only applies to the observed features\", but how do you calculate the loss over \"the entire features\"?)\n\n[1] Igl, Maximilian, et al. \"Deep variational reinforcement learning for POMDPs.\" International Conference on Machine Learning. PMLR, 2018.\n[2] Hausknecht, Matthew, and Peter Stone. \"Deep recurrent q-learning for partially observable mdps.\" 2015 aaai fall symposium series. 2015.", "summary_of_the_paper": "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.", "summary_of_the_review": "The paper proposes an interesting approach to tackle an important problem, but there are some critical comparisons missing related to prior work that preclude claims of significance, contribution, or novelty.\n\nPost-response: \nA3-1: \nI'm still not convinced about the novelty of the VAE approach. I understand that the sequential VAE in DVRL does not have ground truth labels and thus requires inference of latent states, but that's because the assumptions don't allow for imputation of missing features/supervised learning. It could easily be adapted to such a setting, however, as representation learning with ground truth labels will in general be easier/more sample efficient than learning without. Similarly, I'm more impressed by DVRL's ability to continuously improve its state inference model online, which seems like a more difficult approach than the two-staged approach suggested by the authors here.\n\nA3-2: \nThe authors claim that \"when dealing with AFA-POMDP, if DVRL wants to learn a high quality representation that could effectively impute the missing features, it would need the access to the fully observed data to update itself during online training time based on some loss similar to Eq(5) [...] such an attempt would consume a very large amount of feature acquisition cost to get the fully observed data, which is way larger than adopting a multi-stage RL as we propose where we only acquire 2K trajectories to train the sequential VAE model.\" --> It is not clear to me from the paper that, in general, online feature acquisition costs would be any worse than the cost of observing all features offline and collecting those in a batch dataset, as proposed by the paper. By contrast, consider a scenario in which the state includes thousands of noisy and useless features but a few useful and informative features. Then doing the batch representation learning as the authors suggest would lead to unnecessarily high observation costs whereas something like the online DVRL approach could potentially learn much more efficiently. So this claim feels unsubstantiated.\n\nA3-4: \nI appreciate the authors' statement that \"the representation learning model could be updated in a continual learning fashion by the new data as well as the old data and then we could leverage policy transfer [6] techniques to effectively update the policy parameters without inquiring much fully observed data for training the pipeline.\" But this is all still in the hypothetical. At the end of the day, the proposed method as written is fixed after some number of observations.\n\nA3-5: \nI appreciate the authors' diligence on the Ethics Statement.\n\nSeparately, there is a recent paper at NeurIPS 2021 that seems closely aligned with the proposed work. It shouldn't impact the novelty of the present work, as the two papers are essentially contemporaneous, but may be of interest to the authors: https://neurips.cc/Conferences/2021/Schedule?showEvent=26563\n\nIn light of the comments and clarifications, I'm willing to increase my score to a 6.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: Due to the delayed rebuttal made it very hard for reviewers to react.

The paper proposes a new sub-type of POMDPs dubbed AFA-POMDP. The proposed approach first learns a sequential VAE, then an RL approach learns control and feature acquisition policies jointly. The approach is evaluated on two tasks and shows very promising results compared to baselines. Overall the setting and the approach are very interesting.

The replies and revised paper managed to address some of the concerns of the reviewers. However, there remain a few open questions and doubts (see updated reviews), in particular as some of the arguments of the authors remain in the hypothetical, and the reviewers are still not entirely convinced by the choice of the experimental tasks.
Decision: Reject