Original: R EINFORCEMENT L EARNING WITH E FFICIENT A CTIVE F EATURE A CQUISITION
Revised: R EINFORCEMENT L EARNING WITH E FFICIENT A CTIVE F EATURE A CQUISITION

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Solving real-life sequential decision making problems under partial observability involves an exploration-exploitation problem.
Revised: Solving real-life sequential decision making problems under partial observability involves an exploration-exploitation problem.

Original: To be successful, an agent needs to efﬁciently gather valuable information about the state of the world for making rewarding decisions.
Revised: To be successful, an agent needs to efﬁciently gather valuable information about the state of the world for making rewarding decisions.

Original: 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.
Revised: 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.

Original: This poses a signiﬁcant challenge for the agent to perform optimally for the task.
Revised: This poses a signiﬁcant challenge for the agent to learn optimal policy for the task.

Original: In this paper, we propose a modelbased reinforcement learning framework that learns a policy which solves this exploration-exploitation problem during its execution.
Revised: In this paper, we propose a model-based reinforcement learning framework that learns a policy which solves this exploration-exploitation problem during its execution.

Original: Key to the success is a novel sequential variational auto-encoder that learns high-quality representations from partially observed states, which are then used by the policy to maximize the task reward in a cost efﬁcient manner.
Revised: 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-efﬁcient manner.

Original: We demonstrate the efﬁcacy of our proposed framework in a control domain as well as using a medical simulator.
Revised: We demonstrate the efﬁcacy of our proposed framework in a control domain as well as using a medical simulator.

Original: In both tasks, our proposed method outperforms conventional baselines and results in policies with greater cost efﬁciency.
Revised: In both tasks, our proposed method outperforms conventional baselines and results in policies with greater cost efﬁciency.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: 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;
Revised: 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;

Original: Zhang et al., 2018), medical diagnosis (Ling et al., 2017;
Revised: Zhang et al., 2018), medical diagnosis (Ling et al., 2017;

Original: Peng et al., 2018) and computer games (Mnih et al., 2015; Silver et al., 2016).
Revised: Peng et al., 2018) and computer games (Mnih et al., 2015; Silver et al., 2016).

Original: 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.
Revised: 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.

Original: Typically this 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).
Revised: 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).

Original: However, the fundamental assumption that the same features are  readily available during deployment does not hold in many real-world applications.
Revised: However, the fundamental assumption that the same features are always readily available during deployment could not hold in many real-world applications.

Original: 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.
Revised: 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.

Original: To provide the best possible treatment, the system might need to perform several measurements of the patient over time, and some of them could be costly or  pose a health risk.
Revised: 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.

Original: That is, at the deployment the system should function with minimal features while during training more features might have been available.
Revised: Therefore, during deployment, it is more ideal that the system could function with minimal features while during training more features might have been available.

Original: Hence, in such cases, we are interested in decision making models that  take the measurement process, i.e.,the feature acquisition, into account and only acquire the information relevant for making a decision.
Revised: 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.

Original: In this paper, we consider the challenging problem of learning effective policies when the cost of information acquisition cannot be neglected.
Revised: In this paper, we consider the challenging problem of learning effective policies when the cost of information acquisition cannot be neglected.

Original: To be successful,we need to learn policies which acquires the information required for solving a task in the cheapest possible way.
Revised: To be successful,we need to learn policies which acquires the information required for solving a task in the cheapest way possible.

Original: For simplicity, we can think of the policy as being constituted of an acquisition policy which  selects the features to be observed and a task policy , which selects actions to change the state of the system towards some goal.
Revised: 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.

Original: 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.
Revised: 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.

Original: for an illustration).
Revised: for an illustration).

Original: (i)
Revised: (i)

Original: 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
Revised: 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

Original: [Figure]
Revised: [Figure]

Original: controlled by the acquisition policy.
Revised: controlled by the acquisition policy.

Original: Thus, the resulting POMDP is different from typically considered
Revised: Thus, the resulting POMDP is different from the conventional

Original: POMDPs in RL literature (Cassandra, 1998) where the partial observability  stems from a ﬁxed and action-independent observation model.
Revised: POMDPs in RL literature (Cassandra, 1998) where the partial observability for later stems from a ﬁxed and action-independent observation model.

Original: Also, the state-transitions in conventional POMDPs areoften 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.
Revised: 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.

Original: (ii)
Revised: (ii)

Original: The learning of the acquisition policy introduces an additional dimension to the explore-exploit problem: each execution of the policy needs to solve an explore-exploit problem, and thus we often need to learn sophisticated policies.
Revised: 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.

Original: Most reinforcement learning research has not taken active feature acquisition into consideration.
Revised: Most reinforcement learning research has not taken active feature acquisition into consideration.

Original: In this work, we propose a uniﬁed approach that jointly learns a policy for optimizing the task reward while performing active feature acquisition.
Revised: In this work, we propose a uniﬁed approach that jointly learns a policy for optimizing the task reward while performing active feature acquisition.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: As such, both the model dynamics of the underlying MDP and the choice of feature acquisition introduce considerable challenges to  learning  the sequential feature acquisition strategy.
Revised: 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.

Original: Due to the challenge of the exploration-exploitation problem, it is a non-trivial task to jointly learn the  policies.
Revised: Due to the challenge of the exploration-exploitation problem, it is a non-trivial task to jointly learn the two policies.

Original: The conventional end-to-end approaches often result in inferior solutions in complex scenarios.
Revised: The conventional end-to-end approaches often result in inferior solutions in complex scenarios.

Original: Ideally, policies based on high-quality representations would be easier for the algorithm to search for better solutions through exploration-exploitation.
Revised: Ideally, policies based on high-quality representations would be easier for the algorithm to search for better solutions through exploration-exploitation.

Original: Therefore, our presented framework also tackles the joint policy training task from a representation learning perspective.
Revised: Therefore, our proposed framework also tackles the joint policy training task from a representation learning perspective.

Original: Speciﬁcally, we introduce anovel representation learning method that not only encodes the sequential partially observed information into its latent features, but also efﬁciently imputes the unobserved features to offer more meaningful information for the policy training.
Revised: Speciﬁcally, we introduce a representation learning model that not only encodes the sequential partially observed information into its latent features, but also efﬁciently imputes the unobserved features to offer more meaningful information for the policy training.

Original: To this end, we present a sequential generative model that can efﬁciently learn model dynamics during representation learning.
Revised: To this end, we formulate a sequential generative model that can efﬁciently learn model dynamics during representation learning.

Original: Overall, our contributions  are three-fold:
Revised: Overall, the contributions of our paper are three-fold:

Original: • We propose an approach for learning  policies with active feature acquisition through  reinforcement learning.
Revised: • We propose an approach for learning sequential decision making policies with active feature acquisition through a uniﬁed reinforcement learning framework.

Original: Our proposed approach simultaneously learns policies for reward optimization and active feature acquisition.
Revised: Our proposed approach simultaneously learns policies for reward optimization and active feature acquisition.

Original: • We present a novel sequential representation learning approach to account for the encoding of the partially observed states.
Revised: • We present a novel sequential representation learning approach to account for the encoding of the partially observed states.

Original: Our proposed approach is based on variational autoencoders (VAE) with amortized inference.
Revised: Our proposed approach is based on variational autoencoders (VAE) with amortized inference.

Original: The imputation of the unobserved features is achieved via learning  model dynamics.
Revised: The imputation of the unobserved features is achieved via learning the model dynamics.

Original: • We demonstratethat our presented framework can be applied to various applications.
Revised: • We demonstrate our proposed framework can be applied to various applications.

Original: We conduct extensive experiments on an image-based control task and a medical simulator ﬁtted from real-life data. Our method shows clear improvements over conventional baselines.
Revised: We conduct extensive experiments on an image-based control task as well as a medical simulator ﬁtted from real-life data where our method shows clear improvements over conventional baselines.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: In this work, we integrate active learning with reinforcement learning to accomplish the policy training task while acquiring fewer observed features as possible.
Revised: 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.

Original: We thus review related methods for active feature acquisition and representation learning for POMDP, respectively.
Revised: We thus review related methods on active feature acquisition and representation learning for POMDP, respectively.

Original: 2. A CTIVE F EATURE A CQUISITION
Revised: 2. A CTIVE F EATURE A CQUISITION

Original: Our work draws motivation from  existing instance-wise active feature selection approaches.
Revised: Our work draws motivation from the existing instance-wise active feature selection approaches.

Original: 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.
Revised: 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.

Original: One typical example is the conventional linear model that poses sparsity inducing prior distribution to the model (Tibshirani, 1996).
Revised: One typical example is the conventional linear model that poses sparsity inducing prior distribution to the model (Tibshirani, 1996).

Original: Recently, there also emerged approaches that adopt reinforcement learning to actively ﬁnd optimal feature subsets (Yoon et al., 2018; Zannone et al., 2019; Shim et al., 2018)
Revised: Recently, there also emerged approaches that adopt reinforcement learning to actively ﬁnd optimal feature subsets (Yoon et al., 2018; Shim et al., 2018; Zannone et al., 2019).

Original: Though such attempts have demonstrated certain efﬁcacy in handling non time-series instance-wise data, they do not sufﬁce for handling sequential dataset.
Revised: Though such attempts have demonstrated certain efﬁcacy in handling non time-series instance-wise data, they do not sufﬁce for handling sequential dataset.

Original: There is  an alternative category that models feature acquisition as a Bayesian experiment design (Ma et al., 2019;
Revised: There is also an alternative category that models feature acquisition as a Bayesian experimental design (Ma et al., 2019;

Original: Gong et al., 2019).
Revised: Gong et al., 2019).

Original: However, the sequential decision making is for variable-wise feature acquisition and the problems are still non time-series tasks in nature.
Revised: However, the sequential decision making is for variable-wise feature acquisition and the problems are still non time-series tasks in nature.

Original: The key difference between all the existing approaches and ours isin that we tackle active feature acquisition for the time-series data, where an active feature selection decision needs to be formed at each time step along the multi-step reinforcement learning trajectory.
Revised: 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.

Original: Therefore, the feature acquisition for our presented work needs to consider more complex information over model dynamics and control, apart from  static instance-wise features.
Revised: 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.

Original: 2. R EPRESENTATION L EARNING IN POMDP
Revised: 2. R EPRESENTATION L EARNING IN POMDP

Original: In complex tasks, policies trained upon different representations can even converge to different performance levels.
Revised: In complex tasks, policies trained upon different representations can even converge to different performance levels.

Original: Most conventional (deep) reinforcement learning approaches uniﬁes the process of representation learning with policy training and results in policies trained in an end-to-end fashion (Mnih et al., 2013;
Revised: Most conventional deep reinforcement learning approaches uniﬁes the process of representation learning with policy training and results in policies trained in an end-to-end fashion (Mnih et al., 2013;

Original: Lillicrap et al., 2016; Mnih et al., 2016).
Revised: Lillicrap et al., 2016; Mnih et al., 2016).

Original: However, to accomplish the representation learning task, such models often introduce additional training parameters which could come with considerable size and thereby result in signiﬁcant  sample inefﬁciency.
Revised: However, to accomplish the representation learning task, such models often engage trainable parameters which could come with considerable size and thereby result in signiﬁcant degradation in sample efﬁciency.

Original: When considering problems with POMDPs where the state space are partially presented to the agent, representation learning becomes an important and non-trivial research challenge.
Revised: 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.

Original: 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.
Revised: 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.

Original: Higgins et al.
Revised: Higgins et al.

Original: (2016; 2017) adopt pretrained VAE models as  representation module to build agents with strong domain adaptation performance.
Revised: (2016; 2017) adopt pretrained VAE models as a representation module to build agents with strong domain adaptation performance.

Original: The key difference between their works and ours isin that they encode instance-wise image frames from POMDP domains where the image presents  partial view over the task environment,  our work considersthe setting with a distinct partial observability for cost-sensitive reinforcement learning problems, i.e.,  feature-level information is missing at each time step.
Revised: 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.

Original: We thus adopt a sequential representation learning approach to infer a more representative state information.
Revised: We thus adopt a sequential representation learning approach to infer a more representative state information.

Original: Recently, there  emerged a fruitful literature over sequential representation learning for POMDP (Gregor et al., 2019; Vezzani et al., 2019).
Revised: Recently, there also emerged several works on sequential representation learning for POMDP (Gregor et al., 2019; Vezzani et al., 2019).

Original: However, most of the works utilize VAE training as an auxiliary task to jointly update the representation model with the policy.
Revised: However, most of the works utilize VAE training as an auxiliary task to jointly update the representation model with the policy learning loss.

Original: In our work, due to the high acquisition cost to  the entire feature set, we consider off-line representation learning.
Revised: In our work, due to the high acquisition cost to observe the features, we adopt an off-line representation learning setting.

Original: Moreover, we introduce a novel model-based representation learning mechanism, where imputing the missing features brings signiﬁcant beneﬁt to derive high-quality representation for policy training.
Revised: Also, our proposed representation learning is model-based, where the model learns to impute the missing features with such attempt yielding signiﬁcant beneﬁt to derive high-quality representation for policy training.

Original:  M ETHODOLOGY 3. T ASK S ETTING
Revised:  M ETHODOLOGY 3. T ASK S ETTING

Original: In this section, we formally deﬁne the problem settings for the task of jointly learning the task and feature acquisition policy.
Revised: In this section, we formally deﬁne the problem settings for the task of jointly learning the task and feature acquisition policy.

Original: To this end, we deﬁne the active feature acquisition POMDP , a rich class of discrete-time stochastic control processes generalizing standard POMDPs:
Revised: To this end, we deﬁne the active feature acquisition POMDP , a rich class of discrete-time stochastic control processes generalizing standard POMDPs:

Original: Deﬁnition 1 (AFA-POMDP) .
Revised: Deﬁnition 1 (AFA-POMDP) .

Original: The active feature acquisition POMDP is a tuple M = (cid:104)S , A , T , O , R , C , γ (cid:105) , where S is the state space and A
Revised: The active feature acquisition POMDP is a tuple M = (cid:104)S , A , T , O , R , C , γ (cid:105) , where S is the state space and A

Original: =
Revised: =

Original: ( A f , A c ) is a joint action space of feature acquisition actions A f and control actions A c .
Revised: ( A f , A c ) is a joint action space of feature acquisition actions A f and control actions A c .

Original: The transition kernel T : S ×A c ×A f → P S maps any joint action a = ( a f , a c ) in state s ∈ S to a distribution P S over next states.
Revised: The transition kernel T : S ×A c ×A f → P S maps any joint action a = ( a f , a c ) in state s ∈ S to a distribution P S over next states.

Original: In each state s , when taking action a f , the agent observes x p = x
Revised: In each state s , when taking action a f , the agent observes x p = x

Original: ( a f ) , i.e., a subset of the features x = ( x p , x u ) ∼ O ( s ) indicated by a f , where O ( s ) is a distribution over possible feature observation for state s and x u
Revised: ( a f ) , i.e., a subset of the features x = ( x p , x u ) ∼ O ( s ) indicated by a f , where O ( s ) is a distribution over possible feature observation for state s and x u

Original: are features not observed by the agent.
Revised: are features not observed by the agent.

Original: When taking a joint action, the agent obtains rewards according to the reward function R : S
Revised: When taking a joint action, the agent obtains rewards according to the reward function R : S

Original: × A c
Revised: × A c

Original: → R and pays a cost of C : S × A f → R + for feature acquisition.
Revised: → R and pays a cost of C : S × A f → R + for feature acquisition.

Original: Rewards and costs are discounted by the discount factor γ ∈ [0 , 1) .
Revised: Rewards and costs are discounted by the discount factor γ ∈ [0 , 1) .

Original: Simplifying assumptions For simplicity, we assume that x consists of a ﬁxed number of features
Revised: Simplifying assumptions For simplicity, we assume that x consists of a ﬁxed number of features

Original: N f for all states, that A f = 2 [ N f ] is the power set of all the N f features, and that x p ( a f ) consists of all the features in x indicated by the subset a f ∈ A f .
Revised: N f for all states, that A f = 2 [ N f ] is the power set of all the N f features, and that x p ( a f ) consists of all the features in x indicated by the subset a f ∈ A f .

Original: Note that the feature acquisition action for a speciﬁc application can take various different forms.
Revised: Note that the feature acquisition action for a speciﬁc application can take various different forms.

Original: For instance, in our experiments in Section 4, for the Sepsis task, we deﬁne 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).
Revised: For instance, in our experiments in Section 4, for the Sepsis task, we deﬁne 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).

Original: 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.,
Revised: 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.,

Original: T ( s, a c , a f (cid:48) )
Revised: T ( s, a c , a f (cid:48) )

Original: =
Revised: =

Original: T ( s, a c , a f ) for all a f (cid:48) , a f ∈ A f .
Revised: T ( s, a c , a f ) for all a f (cid:48) , a f ∈ A f .

Original: 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.
Revised: 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.

Original: For simplicity we furthermore assume that acquiring each feature has the same cost, denoted as c , i.e., C ( a f , s )
Revised: For simplicity we furthermore assume that acquiring each feature has the same cost, denoted as c , i.e., C ( a f , s )

Original: =
Revised: =

Original: c | a f | , but our approach can be straightforwardly adapted to have different costs for different feature acquisitions.
Revised: c | a f | , but our approach can be straightforwardly adapted to have different costs for different feature acquisitions.

Original: 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 .
Revised: 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 .

Original: That is, we aim to solve the optimization problem
Revised: That is, we aim to solve the optimization problem

Original: [Equation]
Revised: [Equation]

Original: where the expectation is over the randomness of the stochastic process and the policies, a f ( 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.
Revised: where the expectation is over the randomness of the stochastic process and the policies, a f ( 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.

Original: Note that the above optimization problem is very challenging: an optimal solution needs to maintain beliefs b t over the state of the system at time t which is a function of partial observations obtained so far.
Revised: Note that the above optimization problem is very challenging: an optimal solution needs to maintain beliefs b t over the state of the system at time t which is a function of partial observations obtained so far.

Original: Both the the feature acquisition policy π f ( a ft | b t ) and the task policy i.e., π c ( a ct | b t ) depend on this belief.
Revised: Both the the feature acquisition policy π f ( a ft | b t ) and the task policy i.e., π c ( a ct | b t ) depend on this belief.

Original: The information in the belief itself can be controlled by the feature acquisition policy through querying subsets from the features x t and hence the task policy and feature acquisition policy itself strongly depend on effectiveness of the feature acquisition policy.
Revised: The information in the belief itself can be controlled by the feature acquisition policy through querying subsets from the features x t and hence the task policy and feature acquisition policy itself strongly depend on effectiveness of the feature acquisition policy.

Original: 3. S EQUENTIAL R EPRESENTATION L EARNING WITH P ARTIAL O BSERVATIONS
Revised: 3. S EQUENTIAL R EPRESENTATION L EARNING WITH P ARTIAL O BSERVATIONS

Original: We introduce a sequential representation learning approach to facilitate the task of policy training with active feature acquisition.
Revised: We introduce a sequential representation learning approach to facilitate the task of policy training with active feature acquisition.

Original: Let x 1: T = ( x 1 , ..., x T ) and a 1: T = ( a 1 , ..., a T ) denote a sequence of observations and actions, respectively.
Revised: Let x 1: T = ( x 1 , ..., x T ) and a 1: T = ( a 1 , ..., a T ) denote a sequence of observations and actions, respectively.

Original: Alternatively, we also denote these sequences as x ≤ T and a ≤ T .
Revised: Alternatively, we also denote these sequences as x ≤ T and a ≤ T .

Original: Overall, our task of interest is to train a sequential representation learning model to learn the distribution of the full sequential observations x 1: T , i.e., for both the observed part x p 1: T and the unobserved part x u 1: T .
Revised: Overall, our task of interest is to train a sequential representation learning model to learn the distribution of the full sequential observations x 1: T , i.e., for both the observed part x p 1: T and the unobserved part x u 1: T .

Original: Given only partial observations, we can perform inference only with the observed features x p 1: T .
Revised: Given only partial observations, we can perform inference only with the observed features x p 1: T .

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Furthermore, unlike many conventional sequential representation learning models for reinforcement learning that only reason over the observation sequence x p 1: T , in our work, we take into account both the observation sequence x p 1: T and the action sequence
Revised: Furthermore, unlike many conventional sequential representation learning models for reinforcement learning that only reason over the observation sequence x p 1: T , in our work, we take into account both the observation sequence x p 1: T and the action sequence

Original: a 1: T for conducting inference.
Revised: a 1: T for conducting inference.

Original: 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.
Revised: 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.

Original: To summarize, our proposed sequential representation model learns to encode x p 1: T and a 1: T into meaningful latent
Revised: To summarize, our proposed sequential representation model learns to encode x p 1: T and a 1: T into meaningful latent

Original: [Equation]
Revised: [Equation]

Original: 
Revised:

Original: [Equation]
Revised: [Equation]

Original: 
Revised:

Original: [Equation]
Revised: [Equation]

Original: 
Revised:

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: 
Revised:

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: features, for predicting x p 1: T and x u 1: T .
Revised: features, for predicting x p 1: T and x u 1: T .

Original: The architecture of our proposed sequential representation learning model is shown in Figure 2.
Revised: The architecture of our proposed sequential representation learning model is shown in Figure 2.

Original: Observation Decoder Let z 1: T = ( z 1 , ..., z T ) denote a sequence of latent states.
Revised: Observation Decoder Let z 1: T = ( z 1 , ..., z T ) denote a sequence of latent states.

Original: We consider the following probabilistic model:
Revised: We consider the following probabilistic model:

Original: [Equation]
Revised: [Equation]

Original: For simplicity of the notations, we assume z 0 = 0 .
Revised: For simplicity of the notations, we assume z 0 = 0 .

Original: 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 ( z t | z t − 1 , x p 1: t , a 0: t − 1 ) .
Revised: 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 ( z t | z t − 1 , x p 1: t , a 0: t − 1 ) .

Original: The reason is that using a static prior distribution results in latent representation z t that is stronger regularized and more normalized than using a learned prior distribution which is stochastically changing over time.
Revised: The reason is that using a static prior distribution results in latent representation z t that is stronger regularized and more normalized than using a learned prior distribution which is stochastically changing over time.

Original: This is crucial for deriving stable policy training performance.
Revised: This is crucial for deriving stable policy training performance.

Original: At time t , the generation of data x pt and x ut depends on the corresponding latent variable z t .
Revised: At time t , the generation of data x pt and x ut depends on the corresponding latent variable z t .

Original: Given z t , the observed variables are conditionally independent of the unobserved ones.
Revised: Given z t , the observed variables are conditionally independent of the unobserved ones.

Original: Therefore,
Revised: Therefore,

Original: [Equation]
Revised: [Equation]

Original: Belief Inference Model
Revised: Belief Inference Model

Original: During policy training we only assume access to partially observed data.
Revised: During policy training we only assume access to partially observed data.

Original: This requires an inference model which takes in the past observation and action sequences to infer the latent states z .
Revised: This requires an inference model which takes in the past observation and action sequences to infer the latent states z .

Original: Speciﬁcally, we present a structured inference network q φ as shown in Figure 2, which has an autoregressive structure:
Revised: Speciﬁcally, we present a structured inference network q φ as shown in Figure 2, which has an autoregressive structure:

Original: [Equation]
Revised: [Equation]

Original: where q φ ( · ) is a function that aggregates the ﬁltering posteriors of the history of observation and action sequences.
Revised: where q φ ( · ) is a function that aggregates the ﬁltering posteriors of the history of observation and action sequences.

Original: Following the common practice in existing sequential VAE literature, we adopt a forward RNN model as the backbone for the ﬁltering function q φ ( · ) (Gregor et al., 2019).
Revised: Following the common practice in existing sequential VAE literature, we adopt a forward RNN model as the backbone for the ﬁltering function q φ ( · ) (Gregor et al., 2019).

Original: Speciﬁcally, at step t , the RNN processes the encoded partial observation x pt , action a t − 1 and its past hidden state h t − 1 to update its hidden state h t .
Revised: Speciﬁcally, at step t , the RNN processes the encoded partial observation x pt , action a t − 1 and its past hidden state h t − 1 to update its hidden state h t .

Original: Then the latent distribution z t is inferred from h t .
Revised: Then the latent distribution z t is inferred from h t .

Original: The belief state b t is deﬁned as the mean of the distribution z t .
Revised: The belief state b t is deﬁned as the mean of the distribution z t .

Original: 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 beneﬁt from it and progress faster towards getting better convergent performance.
Revised: 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 beneﬁt from it and progress faster towards getting better convergent performance.

Original: Learning We proposed to pre-train both the generative and inference models ofﬂine before learning the RL policies.
Revised: Learning We proposed to pre-train both the generative and inference models ofﬂine before learning the RL policies.

Original: 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.
Revised: 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.

Original: Concretely, the pre-training task update the parameters θ, φ by maximizing the following variational lower-bound (Jordan et al., 1999;
Revised: Concretely, the pre-training task update the parameters θ, φ by maximizing the following variational lower-bound (Jordan et al., 1999;

Original: Kingma & Welling, 2013):
Revised: Kingma & Welling, 2013):

Original: [Equation]
Revised: [Equation]

Original: By incorporating the term log p θ ( x ut | z t ) , 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.
Revised: By incorporating the term log p θ ( x ut | z t ) , 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.

Original: We perform multi-stage reinforcement learning to jointly learn the feature acquisition policy and the task policy.
Revised: We perform multi-stage reinforcement learning to jointly learn the feature acquisition policy and the task policy.

Original: The VAE model is pretrained and kept ﬁxed during policy learning.
Revised: The VAE model is pretrained and kept ﬁxed during policy learning.

Original: 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.
Revised: 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.

Original: The pseudocode for our proposed method is in Appendix A.
Revised: The pseudocode for our proposed method is in Appendix A.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: 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 ﬁtted from real-world data (Oberst & Sontag, 2019).
Revised: 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 ﬁtted from real-world data (Oberst & Sontag, 2019).

Original: 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.
Revised: 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.

Original: For representing the missing features, we adopt the zero-imputing method, proposed in (Nazabal et al., 2018) over the unobserved features.
Revised: For representing the missing features, we adopt the zero-imputing method, proposed in (Nazabal et al., 2018) over the unobserved features.

Original: Thus, we denote the VAE baseline as NonSeq-ZI .
Revised: Thus, we denote the VAE baseline as NonSeq-ZI .

Original: 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).
Revised: 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).

Original: We denoted our proposed sequential VAE model for POMDPs as Seq-PO-VAE .
Revised: We denoted our proposed sequential VAE model for POMDPs as Seq-PO-VAE .

Original: All the
Revised: All the

Original: VAE-based approaches adopt an identical policy architecture.
Revised: VAE-based approaches adopt an identical policy architecture.

Original: Detailed information on the model architecture is presented in appendix.
Revised: Detailed information on the model architecture is presented in appendix.

Original: Data Collection To pre-train the VAE models, data generated by a non-random policy is unavoidably needed to incorporate abundant dynamics information.
Revised: Data Collection To pre-train the VAE models, data generated by a non-random policy is unavoidably needed to incorporate abundant dynamics information.

Original: 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).
Revised: 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).

Original: The simple mixture of dataset works very well on both tasks without the need of further ﬁne-tuning the
Revised: The simple mixture of dataset works very well on both tasks without the need of further ﬁne-tuning the

Original: VAEs.
Revised: VAEs.

Original: We also create a testing set that consists of 2000 trajectories to evaluate the models.
Revised: We also create a testing set that consists of 2000 trajectories to evaluate the models.

Original: 4. B OUNCING B ALL +
Revised: 4. B OUNCING B ALL +

Original: Task Settings We adapted the original bouncing ball experiment presented in (Fraccaro et al., 2017) by adding a navigation objective and introducing control actions.
Revised: Task Settings We adapted the original bouncing ball experiment presented in (Fraccaro et al., 2017) by adding a navigation objective and introducing control actions.

Original: Speciﬁcally, a ball moves in a 2D box and at each step, a binary image of size 32 × 32 showing the box and the ball.
Revised: Speciﬁcally, a ball moves in a 2D box and at each step, a binary image of size 32 × 32 showing the box and the ball.

Original: Initially, the ball appears at a random position in the upper left quadrant
Revised: Initially, the ball appears at a random position in the upper left quadrant

Original: [Table]
Revised: [Table]

Original: and has a random velocity.
Revised: and has a random velocity.

Original: The objective is to control the ball to reach a ﬁxed target location set at (5 , 25) .
Revised: The objective is to control the ball to reach a ﬁxed target location set at (5 , 25) .

Original: We incorporate ﬁve RL actions: a null action and four actions for changing the velocity of the ball in either the x or y direction with a ﬁxed scale: { ∆ V x : ± 0 . 5 , ∆ V y : ± 0 .
Revised: We incorporate ﬁve RL actions: a null action and four actions for changing the velocity of the ball in either the x or y direction with a ﬁxed scale: { ∆ V x : ± 0 . 5 , ∆ V y : ± 0 .

Original: 5 , null } .
Revised: 5 , null } .

Original: A reward of 1.0 is issued if the ball reaches its target location.
Revised: A reward of 1.0 is issued if the ball reaches its target location.

Original: Each episode runs up to 50 time steps.
Revised: Each episode runs up to 50 time steps.

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: impute the missing features (i.e., ball can be reconstructed even when they are unobserved from consequent frames).
Revised: impute the missing features (i.e., ball can be reconstructed even when they are unobserved from consequent frames).

Original: Representation Learning Results We evaluate the missing feature imputing performance of each
Revised: Representation Learning Results We evaluate the missing feature imputing performance of each

Original: VAE model in terms of negative log likelihood (nll) loss and present results in Table 1.
Revised: VAE model in terms of negative log likelihood (nll) loss and present results in Table 1.

Original: We notice that our proposed model yields to signiﬁcantly better imputing result than all the other baselines.
Revised: We notice that our proposed model yields to signiﬁcantly better imputing result than all the other baselines.

Original: This reveals the fact that our proposed sequential VAE model can efﬁciently capture the environment dynamics and learn meaningful information over the missing features.
Revised: This reveals the fact that our proposed sequential VAE model can efﬁciently capture the environment dynamics and learn meaningful information over the missing features.

Original: Such effect is vital in determining the policy training performance in AFA-POMDP, since the policy is conditioned on the
Revised: Such effect is vital in determining the policy training performance in AFA-POMDP, since the policy is conditioned on the

Original: VAE latent features.
Revised: VAE latent features.

Original: We also demonstrate sample trajectories reconstructed by different VAE models in the Appendix.
Revised: We also demonstrate sample trajectories reconstructed by different VAE models in the Appendix.

Original: The result shows that our model learns to impute signiﬁcant amount of missing information given the partially observed sequence.
Revised: The result shows that our model learns to impute signiﬁcant amount of missing information given the partially observed sequence.

Original: Policy Training Results We evaluate the policy training performance in terms of episodic number of acquired observations and the task rewards (w/o cost).
Revised: Policy Training Results We evaluate the policy training performance in terms of episodic number of acquired observations and the task rewards (w/o cost).

Original: The results are presented in Figure 3 (a) and (b), respectively.
Revised: The results are presented in Figure 3 (a) and (b), respectively.

Original: First, we notice that the end-to-end method is vital and fails to learn task skills under the given feature acquisition cost.
Revised: First, we notice that the end-to-end method is vital and fails to learn task skills under the given feature acquisition cost.

Original: However, the VAE-based representation learning methods manage to learn the navigation skill under the same cost setting.
Revised: However, the VAE-based representation learning methods manage to learn the navigation skill under the same cost setting.

Original: This veriﬁes our assumption that representation learning could bring signiﬁcant beneﬁt to the policy training under the AFA-POMDP scenario.
Revised: This veriﬁes our assumption that representation learning could bring signiﬁcant beneﬁt to the policy training under the AFA-POMDP scenario.

Original: 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.
Revised: 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.

Original: Our method also converges to a standard where much less feature acquisition is required to perform the task.
Revised: Our method also converges to a standard where much less feature acquisition is required to perform the task.

Original: [Figure]
Revised: [Figure]

Original: We also show that our proposed method can learn meaningful feature acquisition policies.
Revised: We also show that our proposed method can learn meaningful feature acquisition policies.

Original: To this end, we show three sampled trajectories upon convergence of training in Figure 4.
Revised: To this end, we show three sampled trajectories upon convergence of training in Figure 4.

Original: From the examples, we notice that our feature acquisition policy acquires meaningful features with a majority grasping the exact ball location.
Revised: From the examples, we notice that our feature acquisition policy acquires meaningful features with a majority grasping the exact ball location.

Original: Thus, it demonstrates that the feature acquisition policy adapts to the dynamics of the problem and learns to acquire meaningful features.
Revised: Thus, it demonstrates that the feature acquisition policy adapts to the dynamics of the problem and learns to acquire meaningful features.

Original: We also show the actively learned feature acquisition policy works better than random acquisition.
Revised: We also show the actively learned feature acquisition policy works better than random acquisition.

Original: From the results in Figure 4 (c), our method converges to better standard than random policies with considerably high selection probabilities.
Revised: From the results in Figure 4 (c), our method converges to better standard than random policies with considerably high selection probabilities.

Original: 4. S EPSIS M EDICAL S IMULATOR
Revised: 4. S EPSIS M EDICAL S IMULATOR

Original: Task Setting Our second evaluation domain adopts a medical simulator for treating sepsis among
Revised: Task Setting Our second evaluation domain adopts a medical simulator for treating sepsis among

Original: ICU patients, proposed in (Oberst & Sontag, 2019).
Revised: ICU patients, proposed in (Oberst & Sontag, 2019).

Original: Overall, the task is to learn to apply three treatment actions to the patient, i.e, { antibiotic , ventilation , vasopressors
Revised: Overall, the task is to learn to apply three treatment actions to the patient, i.e, { antibiotic , ventilation , vasopressors

Original: Representation Learning Result We evaluate the imputation performance for each VAE model on the testing dataset.
Revised: Representation Learning Result We evaluate the imputation performance for each VAE model on the testing dataset.

Original: The loss is evaluated in terms of MSE, presented in Table 1.
Revised: The loss is evaluated in terms of MSE, presented in Table 1.

Original: Our proposed method leads to the lowest MSE loss compared to the baselines.
Revised: Our proposed method leads to the lowest MSE loss compared to the baselines.

Original: The result reveals that our proposed sequential VAE could promisingly learn model dynamics for tasks with stochastic transitions.
Revised: The result reveals that our proposed sequential VAE could promisingly learn model dynamics for tasks with stochastic transitions.

Original: Policy Training Result We show the policy training results for Sepsis in Figure 5.
Revised: Policy Training Result We show the policy training results for Sepsis in Figure 5.

Original: Overall, our proposed method results in substantially better task reward compared to all baselines.
Revised: Overall, our proposed method results in substantially better task reward compared to all baselines.

Original: Noticeably, the learning of discharge for our method progresses signiﬁcantly faster than baseline approaches and converges to substantially better values.
Revised: Noticeably, the learning of discharge for our method progresses signiﬁcantly faster than baseline approaches and converges to substantially better values.

Original: The result shows that our method can be trained in a much more sample efﬁcient way.
Revised: The result shows that our method can be trained in a much more sample efﬁcient way.

Original: Moreover, upon convergence, our model outperforms the best non-sequential VAE baseline with a gap of > 5% for discharge ratio.
Revised: Moreover, upon convergence, our model outperforms the best non-sequential VAE baseline with a gap of > 5% for discharge ratio.

Original: For all the evaluation metrics, we notice that VAE-based representation learning models outperform the end-to-end baseline by signiﬁcant margins.
Revised: For all the evaluation metrics, we notice that VAE-based representation learning models outperform the end-to-end baseline by signiﬁcant margins.

Original: This indicates that efﬁcient representation learning may be crucial for deriving satisfying task performance in AFA-POMDP setting.
Revised: This indicates that efﬁcient representation learning may be crucial for deriving satisfying task performance in AFA-POMDP setting.

Original: The result also reveals that learning to impute missing features contributes greatly to improve the policy training performance for AFA-POMDP.
Revised: The result also reveals that learning to impute missing features contributes greatly to improve the policy training performance for AFA-POMDP.

Original: Ablation: Efﬁcacy of Active Feature Acquisition We study the effect of actively learning sequential feature acquisition strategy with RL.
Revised: Ablation: Efﬁcacy of Active Feature Acquisition We study the effect of actively learning sequential feature acquisition strategy with RL.

Original: To this end, we compare our method with a baseline that randomly acquires features.
Revised: To this end, we compare our method with a baseline that randomly acquires features.

Original: We evaluate our method under different cost values, and the results are shown in Figure 6.
Revised: We evaluate our method under different cost values, and the results are shown in Figure 6.

Original: From the results, we notice that there is a clear cost-performance trade-off, i.e., a
Revised: From the results, we notice that there is a clear cost-performance trade-off, i.e., a

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: higher feature acquisition cost results in feature acquisition policies that obtain fewer observations, with a sacriﬁce of task performance.
Revised: higher feature acquisition cost results in feature acquisition policies that obtain fewer observations, with a sacriﬁce of task performance.

Original: Overall, our acquisition method results in signiﬁcantly better task performance than the random acquisition baselines.
Revised: Overall, our acquisition method results in signiﬁcantly better task performance than the random acquisition baselines.

Original: 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.
Revised: 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.

Original: Also, we notice that the speciﬁed cost has a very clear impact on the ﬁnal task performance, i.e., the number of acquired features per episode decreases signiﬁcantly as the cost increases.
Revised: Also, we notice that the speciﬁed cost has a very clear impact on the ﬁnal task performance, i.e., the number of acquired features per episode decreases signiﬁcantly as the cost increases.

Original: Thereby, our proposed solution can promisingly compute feature acquisition policies that meet different budgets.
Revised: Thereby, our proposed solution can promisingly compute feature acquisition policies that meet different budgets.

Original: 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.
Revised: 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.

Original: The results are shown in Figure 7.
Revised: The results are shown in Figure 7.

Original: 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.
Revised: 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.

Original: We notice that policies trained by our method result in the highest convergent task performance (max x-value).
Revised: We notice that policies trained by our method result in the highest convergent task performance (max x-value).

Original: Given a certain performance level (same x-value), our method consumes substantially less total feature acquisition cost (y-value) than the others.
Revised: Given a certain performance level (same x-value), our method consumes substantially less total feature acquisition cost (y-value) than the others.

Original: We also notice that the overall feature acquisition cost increases with a near exponential trend.
Revised: We also notice that the overall feature acquisition cost increases with a near exponential trend.

Original: 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.
Revised: 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.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We present a novel AFA-POMDP framework that jointly learns the task policy and the active feature acquisition strategy with a uniﬁed reinforcement learning formalism.
Revised: We present a novel AFA-POMDP framework that jointly learns the task policy and the active feature acquisition strategy with a uniﬁed reinforcement learning formalism.

Original: We introduce a novel sequential VAE model to facilitate policy training under partial observability.
Revised: We introduce a model-based sequential VAE model to facilitate policy training under partial observability.

Original: We demonstrate that imputing missing features via learning model dynamics could signiﬁcantly beneﬁt policy training with partial observability.
Revised: We demonstrate that imputing missing features via learning model dynamics could signiﬁcantly beneﬁt policy training with partial observability.

Original: Our proposed model, by efﬁciently synthesizing the sequential information and imputing missing features, can signiﬁcantly outperform conventional representation learning baselines and leads to policy training with signiﬁcantly better sample efﬁciency and obtained solutions.
Revised: Our proposed model, by efﬁciently synthesizing the sequential information to impute the missing features, can signiﬁcantly outperform conventional representation learning baselines and leads to policy training with signiﬁcantly better sample efﬁciency as well as obtained solutions.

Original: Future work may investigate whether our proposed model could be applied to more diverse and complex application domains.
Revised: Future work may investigate whether our proposed model could be applied to more diverse and complex application domains.

Original: Another promising direction is to integrate our framework with model-based planning for further reducing the feature acquisition cost.
Revised: Another promising direction is to integrate our framework with model-based planning for further reducing the feature acquisition cost.

Original: E THICS S TATEMENT
Revised: E THICS S TATEMENT

Original: 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.
Revised: 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.

Original: Our work addresses the aforementioned problem via formulating a novel
Revised: Our work addresses the aforementioned problem via formulating a novel

Original: 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.
Revised: 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.

Original: The sequential active feature acquisition module enables the decision making to be performed in a more cost-efﬁcient way when partial features are accessible only.
Revised: The sequential active feature acquisition module enables the decision making to be performed in a more cost-efﬁcient way when partial features are accessed only during model deployment.

Original: Considering that  learning and applying the models is  problem speciﬁc it is unlikely that our method can equally beneﬁt all possible application scenarios.
Revised: Considering that the task of learning and applying machine learning models is rather problem speciﬁc it is unlikely that our method can equally beneﬁt all possible application scenarios.

Original: We also fully acknowledge the existence of risk in applying our model in sensitive and high risk domains, e g., healthcare, and  bias if the model itself or the used representations are trained on biased data.
Revised: 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.

Original: In high risk settings, human supervision of the proposed model might be desired and the model could mainly be  used for decision support.
Revised: 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.

Original: However, there are still many practical scenarios that could satisfy our model assumption and are less sensitive. Also, we believe that our paper has the potential to triggermany follow-up works studying novel models, theoretical aspects, and the application of AFA-POMDP in different domains.
Revised: To alleviate the reliance on fully observed data during representation learning, it is very promising to trigger follow-up works studying data efﬁcient sequential autoencoder training paradigms.

Original: APPENDIX
Revised: APPENDIX

Original: This appendix is organized as follows:
Revised: This appendix is organized as follows:

Original: • Sec A : the detailed algorithm.
Revised: • Sec A : the detailed algorithm.

Original: • Sec B : experimental settings and additional results on the Bouncing Ball domain.
Revised: • Sec B : experimental settings and additional results on the Bouncing Ball domain.

Original: • Sec C : experimental settings and additional results on the Sepsis domain.
Revised: • Sec C : experimental settings and additional results on the Sepsis domain.

Original: A RL WITH A CTIVE F EATURE A CQUISITION A LGORITHM
Revised: A RL WITH A CTIVE F EATURE A CQUISITION A LGORITHM

Original: action
Revised: action

Original: 7: for to do 8:
Revised: 7: for to do 8:

Original: Compute representation with VAE: b t = q φ ( x p ≤
Revised: Compute representation with VAE: b t = q φ ( x p ≤

Original: , a
Revised: , a

Original: Sample a feature acquisition action action a
Revised: Sample a feature acquisition action action a

Original: Step environment receive features env a a
Revised: Step environment receive features env a a

Original: Compute cost transitions then 14: break 15: end while
Revised: Compute cost transitions then 14: break 15: end while

Original: B B OUNCING B ALL + B. T ASK S PECIFICATION
Revised: B B OUNCING B ALL + B. T ASK S PECIFICATION

Original: The task consists of a ball moving in a 2D box of size 32 × 32 pixels.
Revised: The task consists of a ball moving in a 2D box of size 32 × 32 pixels.

Original: The radius of the ball equals topixels.
Revised: The radius of the ball equals topixels.

Original: At each step, a binary image is returned as an observation of the MDP state.
Revised: At each step, a binary image is returned as an observation of the MDP state.

Original: At the beginning of every episode, the ball starts at a random position in the upper left quadrant (sampled uniformly).
Revised: At the beginning of every episode, the ball starts at a random position in the upper left quadrant (sampled uniformly).

Original: The initial velocity of the ball is randomly deﬁned as follows: (cid:126)v = [ V x , V y ] = 4 · ˜ (cid:126)v/ (cid:107) ˜ (cid:126)v (cid:107) , where the x- and y-component of ˜ (cid:126)v are sampled uniformly from the interval [ − 0 .
Revised: The initial velocity of the ball is randomly deﬁned as follows: (cid:126)v = [ V x , V y ] = 4 · ˜ (cid:126)v/ (cid:107) ˜ (cid:126)v (cid:107) , where the x- and y-component of ˜ (cid:126)v are sampled uniformly from the interval [ − 0 .

Original: There is a navigation target set at (5 , 25) pixels, which is in the lower left quadrant.
Revised: There is a navigation target set at (5 , 25) pixels, which is in the lower left quadrant.

Original: The navigation is considered to be successful if the ball reaches the speciﬁed target location within a threshold of 1 pixel along both x/y-axis.
Revised: The navigation is considered to be successful if the ball reaches the speciﬁed target location within a threshold of 1 pixel along both x/y-axis.

Original: The action spaces is deﬁned as follows.
Revised: The action spaces is deﬁned as follows.

Original: There are ﬁve task actions
Revised: There are ﬁve task actions

Original: A c :
Revised: A c :

Original: • Increase velocity leftwards, i.e., change V x by − 0 .•
Revised: • Increase velocity leftwards, i.e., change V x by − 0 .•

Original: Increase velocity rightwards, i.e., change V x by +0 .•
Revised: Increase velocity rightwards, i.e., change V x by +0 .•

Original: Increase velocity downwards, i.e., change V y by +0 .•
Revised: Increase velocity downwards, i.e., change V y by +0 .•

Original: Increase velocity upwards, i
Revised: Increase velocity upwards, i

Original: e., change V y by − 0 .•
Revised: e., change V y by − 0 .•

Original: Keep velocities unchanged
Revised: Keep velocities unchanged

Original: The maximum velocity along the x/y-axis is 5.0.
Revised: The maximum velocity along the x/y-axis is 5.0.

Original: The velocity will stay unchanged if it exceeds this threshold.
Revised: The velocity will stay unchanged if it exceeds this threshold.

Original: The feature acquisition action a f ∈
Revised: The feature acquisition action a f ∈

Original: A f is speciﬁed as acquiring the observation of a subset of the quadrants (this also includes acquiring the observation of all 4 quadrants).
Revised: A f is speciﬁed as acquiring the observation of a subset of the quadrants (this also includes acquiring the observation of all 4 quadrants).

Original: Thus, the agent can acquire 0 − 4 quadrants to observe.
Revised: Thus, the agent can acquire 0 − 4 quadrants to observe.

Original: Each episode runs up to 50 steps.
Revised: Each episode runs up to 50 steps.

Original: The episode terminates if agent reaches the target location.
Revised: The episode terminates if agent reaches the target location.

Original: B. I MPLEMENTATION D ETAILS
Revised: B. I MPLEMENTATION D ETAILS

Original: For all the compared methods, Zero-Imputing (Nazabal et al., 2018) is adopted to ﬁll in missing features with a ﬁxed value of 0.5.
Revised: For all the compared methods, Zero-Imputing (Nazabal et al., 2018) is adopted to ﬁll in missing features with a ﬁxed value of 0.5.

Original: End-to-End The end-to-end model ﬁrst processes the imputed image by 2 convolutional layers with ﬁlter sizes of 16 and 32, respectively.
Revised: End-to-End The end-to-end model ﬁrst processes the imputed image by 2 convolutional layers with ﬁlter sizes of 16 and 32, respectively.

Original: Each convolutional layer is followed by a ReLU activation function.
Revised: Each convolutional layer is followed by a ReLU activation function.

Original: Then the output is passed to a fully connected layer of size 1024.
Revised: Then the output is passed to a fully connected layer of size 1024.

Original: The weights for the fully connected layer are initialized by orthogonal weights initialization and the biases are initialized as zeros.
Revised: The weights for the fully connected layer are initialized by orthogonal weights initialization and the biases are initialized as zeros.

Original: NonSeq-ZI
Revised: NonSeq-ZI

Original: The non-sequential VAE models ﬁrst process the imputed image by 2 convolutional layers with ﬁlter sizes of 32 and 64, respectively.
Revised: The non-sequential VAE models ﬁrst process the imputed image by 2 convolutional layers with ﬁlter sizes of 32 and 64, respectively.

Original: Each convolutional layer is followed by a ReLU activation function.
Revised: Each convolutional layer is followed by a ReLU activation function.

Original: 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.
Revised: 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.

Original: To decode an image, the sampled code ﬁrst passes through a fully connected layer with size 256 , followed by 3 convolutional layers with ﬁlters 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.
Revised: To decode an image, the sampled code ﬁrst passes through a fully connected layer with size 256 , followed by 3 convolutional layers with ﬁlters 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.

Original: 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.
Revised: 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.

Original: The hyperparameters for training NonSeq-ZI are summarized in Table 2.
Revised: The hyperparameters for training NonSeq-ZI are summarized in Table 2.

Original: [Table]
Revised: [Table]

Original: Seq-PO-VAE (ours)
Revised: Seq-PO-VAE (ours)

Original: At each step, the Seq-PO-VAE takes an imputed image and an action vector of size 9 as input.
Revised: At each step, the Seq-PO-VAE takes an imputed image and an action vector of size 9 as input.

Original: The imputed image is processed by 3 convolutional layers with ﬁlter sizeand stride 2.
Revised: The imputed image is processed by 3 convolutional layers with ﬁlter sizeand stride 2.

Original: Each convolutional layer employs ReLU as its activation function.
Revised: Each convolutional layer employs ReLU as its activation function.

Original: Then the output passes through a fully connected layer of size 32 to generate a latent representation for the image f x .
Revised: Then the output passes through a fully connected layer of size 32 to generate a latent representation for the image f x .

Original: The action vector passes through a fully connected layer of 32 to generate latent representation for the action f a .
Revised: The action vector passes through a fully connected layer of 32 to generate latent representation for the action f a .

Original: Then the image and action features are concatenated and augmented to form a feature vector f c =
Revised: Then the image and action features are concatenated and augmented to form a feature vector f c =

Original: [ f x , f a , f x ∗ f a ] , where [ · ] denotes concatenation of features.
Revised: [ f x , f a , f x ∗ f a ] , where [ · ] denotes concatenation of features.

Original: Then f c is fed to fully connected projection layers of size 64 and 32, respectively.
Revised: Then f c is fed to fully connected projection layers of size 64 and 32, respectively.

Original: The output is then fed to an LSTM module, with latent size of 32.
Revised: The output is then fed to an LSTM module, with latent size of 32.

Original: The output h t 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 ﬁltered from the sequential inputs.
Revised: The output h t 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 ﬁltered from the sequential inputs.

Original: To decode an image, the model adopts deconvolutional layers that are identical to those for NonSeq-ZI .
Revised: To decode an image, the model adopts deconvolutional layers that are identical to those for NonSeq-ZI .

Original: The hyperparameters for training Seq-PO-VAE are shown in Table 2.
Revised: The hyperparameters for training Seq-PO-VAE are shown in Table 2.

Original: LSTM-A3C
Revised: LSTM-A3C

Original: We adopt LSTM-A3C
Revised: We adopt LSTM-A3C

Original: (Mnih et al., 2016) to train the RL policy.
Revised: (Mnih et al., 2016) to train the RL policy.

Original: The policy takes the features derived from the representation learning module as input.
Revised: The policy takes the features derived from the representation learning module as input.

Original: For the VAE-based methods, the input features are passed through a fully connected layer of size 1024.
Revised: For the VAE-based methods, the input features are passed through a fully connected layer of size 1024.

Original: Then the features are fed to an
Revised: Then the features are fed to an

Original: LSTM with 1024 units.
Revised: LSTM with 1024 units.

Original: 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.
Revised: 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.

Original: We adopt normalized column initialization for all the fully connected layers and the biases for the LSTM module are set to be zero.
Revised: We adopt normalized column initialization for all the fully connected layers and the biases for the LSTM module are set to be zero.

Original: B. D ATA C OLLECTION
Revised: B. D ATA C OLLECTION

Original: To train the VAEs, we prepare a training set that consists of 2000 trajectories.
Revised: To train the VAEs, we prepare a training set that consists of 2000 trajectories.

Original: 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.
Revised: 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.

Original: To train the end-to-end method, we employ a cost of 0.01 over ﬁrst 2m steps and then increase it to 0.02 for the following 0.5m steps.
Revised: To train the end-to-end method, we employ a cost of 0.01 over ﬁrst 2m steps and then increase it to 0.02 for the following 0.5m steps.

Original: All the VAE models are evaluated on a test dataset that has identical size and data distribution as the training dataset.
Revised: All the VAE models are evaluated on a test dataset that has identical size and data distribution as the training dataset.

Original: We present the best achieved task performance of the data collection policy ( End-to-End ) and our representation learning approach in
Revised: We present the best achieved task performance of the data collection policy ( End-to-End ) and our representation learning approach in

Original: Table 5.
Revised: Table 5.

Original: We notice that our proposed method, by employing an advanced representation model, leads to signiﬁcantly better feature acquisition policy than End-to-End (smaller number of observations while achieving similar or better reward).
Revised: We notice that our proposed method, by employing an advanced representation model, leads to signiﬁcantly better feature acquisition policy than End-to-End (smaller number of observations while achieving similar or better reward).

Original: [Table]
Revised: [Table]

Original: B. I MPUTING M ISSING F EATURES VIA L EARNING M ODEL D YNAMICS
Revised: B. I MPUTING M ISSING F EATURES VIA L EARNING M ODEL D YNAMICS

Original: We present an illustrative example to demonstrate the process of imputing missing features and the role of learning model dynamics.
Revised: We present an illustrative example to demonstrate the process of imputing missing features and the role of learning model dynamics.

Original: 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.
Revised: 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.

Original: From the results presented in Figure 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 signiﬁcant beneﬁt for the policy model to learn to acquire meaningful features/gain task reward.
Revised: From the results presented in Figure 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 signiﬁcant beneﬁt for the policy model to learn to acquire meaningful features/gain task reward.

Original: B. I NVESTIGATION ON C OST -P ERFORMANCE T RADE - OFF
Revised: B. I NVESTIGATION ON C OST -P ERFORMANCE T RADE - OFF

Original: We perform a case study on investigating the cost-performance trade-off for each representation learning method, presented in Figure 9.
Revised: We perform a case study on investigating the cost-performance trade-off for each representation learning method, presented in Figure 9.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: 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).
Revised: 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).

Original: Thus, we could conclude that our proposed approach can signiﬁcantly beneﬁt the cost-sensitive policy training and lead to a policy which acquires much fewer observations while still succeeding in terms of task performance.
Revised: Thus, we could conclude that our proposed approach can signiﬁcantly beneﬁt the cost-sensitive policy training and lead to a policy which acquires much fewer observations while still succeeding in terms of task performance.

Original: [Figure]
Revised: [Figure]

Original: C S EPSIS M EDICAL S IMULATOR C. T ASK S PECIFICATIONS
Revised: C S EPSIS M EDICAL S IMULATOR C. T ASK S PECIFICATIONS

Original: For this task we employ a Sepsis simulator proposed in previous work (Oberst & Sontag, 2019).
Revised: For this task we employ a Sepsis simulator proposed in previous work (Oberst & Sontag, 2019).

Original: The task is to learn to apply three treatment actions for Sepsis patients in intensive care units, i.e.,
Revised: The task is to learn to apply three treatment actions for Sepsis patients in intensive care units, i.e.,

Original: A c = { antibiotic , ventilation , vasopressors } .
Revised: A c = { antibiotic , ventilation , vasopressors } .

Original: At each time step, the agent selects a subset of the treatment actions to apply.
Revised: At each time step, the agent selects a subset of the treatment actions to apply.

Original: 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.
Revised: 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.

Original: The feature acquisition actively selects a subset among the measurement features for observation, i.e., A f = { heart rate , sysBP rate , percoxyg state , glucose level } .
Revised: The feature acquisition actively selects a subset among the measurement features for observation, i.e., A f = { heart rate , sysBP rate , percoxyg state , glucose level } .

Original: The objective for learning a active feature acquisition strategy is to help the decision making system to reduce measurement cost at a signiﬁcant scale.
Revised: The objective for learning a active feature acquisition strategy is to help the decision making system to reduce measurement cost at a signiﬁcant scale.

Original: C. I MPLEMENTATION D ETAILS
Revised: C. I MPLEMENTATION D ETAILS

Original: For all the compared methods, we adopt Zero-Imputing (Nazabal et al., 2018) to ﬁll in missing features.
Revised: For all the compared methods, we adopt Zero-Imputing (Nazabal et al., 2018) to ﬁll in missing features.

Original: In particular, a ﬁxed value of -10 which is outside the range of feature values is used to impute missing values.
Revised: In particular, a ﬁxed value of -10 which is outside the range of feature values is used to impute missing values.

Original: End-to-End The end-to-end model ﬁrst processes the imputed state by 3 fully connected layers of size 32, 64 and 32, respectively.
Revised: End-to-End The end-to-end model ﬁrst processes the imputed state by 3 fully connected layers of size 32, 64 and 32, respectively.

Original: Each fully connected layer is followed by a ReLU activation.
Revised: Each fully connected layer is followed by a ReLU activation.

Original: NonSeq-ZI The VAE model
Revised: NonSeq-ZI The VAE model

Original: ﬁrst processes the imputed state by 2 fully connected layers with size 32 and 64, with the ﬁrst fully connected layer being followed by ReLU activation functions.
Revised: ﬁrst processes the imputed state by 2 fully connected layers with size 32 and 64, with the ﬁrst fully connected layer being followed by ReLU activation functions.

Original: 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.
Revised: 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.

Original: To decode the state, the latent code is ﬁrst processed by a fully connected layer of size 64, then fed into three fully connected layers of size 64, 32, and 8.
Revised: To decode the state, the latent code is ﬁrst processed by a fully connected layer of size 64, then fed into three fully connected layers of size 64, 32, and 8.

Original: The intermediate fully connected layers employ ReLU activation functions.
Revised: The intermediate fully connected layers employ ReLU activation functions.

Original: Also, we adopt two variants for NonSeq-ZI , trained under either full loss or partial loss.
Revised: Also, we adopt two variants for NonSeq-ZI , trained under either full loss or partial loss.

Original: The details of the hyperparameter settings used for training are presented in Table 4.
Revised: The details of the hyperparameter settings used for training are presented in Table 4.

Original: [Table]
Revised: [Table]

Original: Seq-PO-VAE (ours)
Revised: Seq-PO-VAE (ours)

Original: [Table]
Revised: [Table]

Original: Seq-PO-VAE (ours)
Revised: Seq-PO-VAE (ours)

Original: At each time step, the inputs for state and action are ﬁrst processed by their corresponding projection layers.
Revised: At each time step, the inputs for state and action are ﬁrst processed by their corresponding projection layers.

Original: 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.
Revised: 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.

Original: The projection layer for the action input is a fully connected layer of size 10.
Revised: The projection layer for the action input is a fully connected layer of size 10.

Original: Then the projected state feature f c and action feature f a are combined in the following manner: f c =
Revised: Then the projected state feature f c and action feature f a are combined in the following manner: f c =

Original: [ f x , f a , f x ∗ f a ] .
Revised: [ f x , f a , f x ∗ f a ] .

Original: f c is passed to 2 fully connected layers of size 64 and 32 to form the input to the
Revised: f c is passed to 2 fully connected layers of size 64 and 32 to form the input to the

Original: LSTM module.
Revised: LSTM module.

Original: The output h t of the LSTM is fed to two independent fully connected layers of size 10 to generate the mean and variance for the Gaussian distribution.
Revised: The output h t of the LSTM is fed to two independent fully connected layers of size 10 to generate the mean and variance for the Gaussian distribution.

Original: The decoder for Seq-PO-VAE has identical architecture as NonSeq-ZI .
Revised: The decoder for Seq-PO-VAE has identical architecture as NonSeq-ZI .

Original: The details for training Seq-PO-VAE are presented in Table 4.
Revised: The details for training Seq-PO-VAE are presented in Table 4.

Original: LSTM-A3C
Revised: LSTM-A3C

Original: The LSTM-A3C (Mnih et al., 2016) takes encoded state features derived from the corresponding representation model as its input.
Revised: The LSTM-A3C (Mnih et al., 2016) takes encoded state features derived from the corresponding representation model as its input.

Original: The encoded featuresare fed into an LSTM with size 256.
Revised: The encoded featuresare fed into an LSTM with size 256.

Original: Then the h t for the LSTM is fed to three independent fully connected layers, to predict the state value, feature acquisition policy and task policy.
Revised: Then the h t for the LSTM is fed to three independent fully connected layers, to predict the state value, feature acquisition policy and task policy.

Original: Normalized column initialization is applied to all fully connected layers.
Revised: Normalized column initialization is applied to all fully connected layers.

Original: The biases for the LSTM and fully connected layers are initialized as zero.
Revised: The biases for the LSTM and fully connected layers are initialized as zero.

Original: C. D ATA C OLLECTION
Revised: C. D ATA C OLLECTION

Original: To train the VAEs, we prepare a training set that consists of 2000 trajectories.
Revised: To train the VAEs, we prepare a training set that consists of 2000 trajectories.

Original: Half of the trajectories are derived from a random policy and the other half is derived from a policy learned from the
Revised: Half of the trajectories are derived from a random policy and the other half is derived from a policy learned from the

Original: End-to-End method with cost 0 . 0 .
Revised: End-to-End method with cost 0 . 0 .

Original: All the VAE models are evaluated on a test dataset that consists of identical size and data distribution as the the training dataset.
Revised: All the VAE models are evaluated on a test dataset that consists of identical size and data distribution as the the training dataset.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: [Table]
Revised: [Table]

Original: C. M ORE C OMPARISON R ESULT UNDER D IFFERENT V ALUES FOR C OST
Revised: C. M ORE C OMPARISON R ESULT UNDER D IFFERENT V ALUES FOR C OST

Original: We present additional experiment results that compare our proposed method and the non-sequential baselines under the cost values { 0, 0.025 } .
Revised: We present additional experiment results that compare our proposed method and the non-sequential baselines under the cost values { 0, 0.025 } .

Original: The results for cost value of 0.01 are shown in the main paper.
Revised: The results for cost value of 0.01 are shown in the main paper.

Original: Overall, under all the cost settings, our method leads to signiﬁcantly better discharge ratio and task reward compared to the baselines.
Revised: Overall, under all the cost settings, our method leads to signiﬁcantly better discharge ratio and task reward compared to the baselines.

Original: [Figure]
Revised: [Figure]

Original: Also, we demonstrate the cost-performance trade-off on Sepsis domain.
Revised: Also, we demonstrate the cost-performance trade-off on Sepsis domain.

Original: By increasing the value of cost, we could obtain feature acquisition policy that acquires substantially decreased amount of features within each episode.
Revised: By increasing the value of cost, we could obtain feature acquisition policy that acquires substantially decreased amount of features within each episode.

Original: [Figure]
Revised: [Figure]

Original: C. I LLUSTRATIVE E XAMPLES FOR M ISSING F EATURE I MPUTATION IN Sepsis
Revised: C. I LLUSTRATIVE E XAMPLES FOR M ISSING F EATURE I MPUTATION IN Sepsis

Original: 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
Revised: 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

Original: [Figure]
Revised: [Figure]

Original: The policy training process with partial observability could only access very limited information, due to the employment of active feature acquisition.
Revised: The policy training process with partial observability could only access very limited information, due to the employment of active feature acquisition.

Original: Under such circumstances, imputing the missing features would offer much more abundant information to the decision making process.
Revised: Under such circumstances, imputing the missing features would offer much more abundant information to the decision making process.

Original: From the results shown in Figure
Revised: From the results shown in Figure

Original: [Figure]
Revised: [Figure]

Original: 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.
Revised: 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.

Original: The imputed missing information would be greatly beneﬁcial for training the task policy and feature acquisition policy
Revised: The imputed missing information would be greatly beneﬁcial for training the task policy and feature acquisition policy

Original: [Figure]
Revised: [Figure]

Original: By imputing the missing features via learning model dynamics, our proposed method could offer much more informative representation for the training compared to the
Revised: By imputing the missing features via learning model dynamics, our proposed method could offer much more informative representation for the training compared to the

Original: -sequential baselines, and thus signiﬁcantly beneﬁt the training with partial observability.
Revised: -sequential baselines, and thus signiﬁcantly beneﬁt the training with partial observability.