Reinforcement Learning with Efficient Active Feature Acquisition

1 INTRODUCTION . Recently , machine learning models for automated sequential decision making have shown remarkable success across many application areas , such as visual recognition ( Mathe et al. , 2016 ; Das et al. , 2017 ) , robotics control ( Finn et al. , 2016 ; Zhang et al. , 2018 ) , medical diagnosis ( Ling et al. , 2017 ; Peng et al. , 2018 ) and computer games ( Mnih et al. , 2015 ; Silver et al. , 2016 ) . One fundamental reason that drives the success of such models and enables them to outperform classical algorithms is the availability of large amounts of training data . Typically such training data is either fully observed or the features stem from an action-independent observation model ( which clearly can depend on the state of the system ) . However , the fundamental assumption that the same features are always readily available during deployment could not hold in many real-world applications . For instance , consider a medical support system for monitoring and treating patients during their stay at hospital which was trained on rich historical medical data . To provide the best possible treatment , the system might need to perform several measurements of the patient over time , while some of them could be costly or even pose a health risk . Therefore , during deployment , it is more ideal that the system could function with minimal features while during training more features might have been available . In such cases , we are interested in decision making models that actively take the measurement process , i.e. , feature acquisition , into account and only acquire the information relevant for making a decision . In this paper , we consider the challenging problem of learning effective policies when the cost of information acquisition can not be neglected . To be successful , we need to learn policies which acquires the information required for solving a task in the cheapest way possible . For simplicity , we can think of the policy as being constituted of an acquisition policy which actively selects meaningful features to be observed and a task policy , which selects actions to change the state of the system towards some goal.1 As such , we consider a partially observable learning problem with the following two distinguishing properties compared to the most commonly studied problems ( see also Figure 3.2 for an illustration ) . ( i ) By incorporating active feature acquisition , the training of the task policy is based upon subsets of features only , i.e. , there are missing features , where the missingness is 1Clearly , these two policies are not independent in general , e.g. , acquiring features can change the state of the system . observe action acquire ( e.g . navigation ) ( e.g . medical treatments ) observe action controlled by the acquisition policy . Thus , the resulting POMDP is different from the conventional POMDPs in RL literature ( Cassandra , 1998 ) where the partial observability for later stems from a fixed and action-independent observation model . Also , the state transitions in conventional POMDPs are only determined by the choice of the task action , whereas in our setting the state-transition is affected by both the task action and the feature acquisition choice . ( ii ) The learning of the acquisition policy introduces an additional dimension to the exploration-exploitation problem : each execution of the policy needs to solve an exploration-exploitation problem , and thus we often need to learn sophisticated policies . Most reinforcement learning research has not taken active feature acquisition into consideration . In this work , we propose a unified approach that jointly learns a policy for optimizing the task reward while performing active feature acquisition . Although some of the prior works have exploited the use of reinforcement learning for sequential feature acquisition tasks ( Shim et al. , 2018 ; Zannone et al. , 2019 ) , they considered variable-wise information acquisition in a static setting only , corresponding to feature selection for non-time-dependent prediction tasks . However , our considered setting is truly time-dependent and feature acquisitions need to be made at each time step while the state of the system evolves simultaneously . As such , both the model dynamics of the underlying MDP and the choice of feature acquisition introduce considerable challenges to the learning of the sequential feature acquisition strategy . Due to the challenge of the exploration-exploitation problem , it is a non-trivial task to jointly learn the two policies . The conventional end-to-end approaches often result in inferior solutions in complex scenarios . Ideally , policies based on high-quality representations would be easier for the algorithm to search for better solutions through exploration-exploitation . Therefore , our proposed framework also tackles the joint policy training task from a representation learning perspective . Specifically , we introduce a representation learning model that not only encodes the sequential partially observed information into its latent features , but also efficiently imputes the unobserved features to offer more meaningful information for the policy training . To this end , we formulate a sequential generative model that can efficiently learn model dynamics during representation learning . Overall , the contributions of our paper are three-fold : • We propose an approach for learning sequential decision making policies with active feature acquisition through a unified reinforcement learning framework . Our proposed approach simultaneously learns policies for reward optimization and active feature acquisition . • We present a novel sequential representation learning approach to account for the encoding of the partially observed states . Our proposed approach is based on variational autoencoders ( VAE ) with amortized inference . The imputation of the unobserved features is achieved via learning the model dynamics . • We demonstrate our proposed framework can be applied to various applications . We conduct extensive experiments on an image-based control task as well as a medical simulator fitted from real-life data where our method shows clear improvements over conventional baselines . 2 RELATED WORK . In this work , we integrate active learning with reinforcement learning to accomplish the policy training task while attempting to acquire fewest observed features as possible . We thus review related methods on active feature acquisition and representation learning for POMDP , respectively . 2.1 ACTIVE FEATURE ACQUISITION . Our work draws motivation from the existing instance-wise active feature selection approaches . One category of the instance-wise feature selection methods consider feature acquisition as a one time effort to select a subset of features as a whole . One typical example is the conventional linear model that poses sparsity inducing prior distribution to the model ( Tibshirani , 1996 ) . Recently , there also emerged approaches that adopt reinforcement learning to actively find optimal feature subsets ( Yoon et al. , 2018 ; Shim et al. , 2018 ; Zannone et al. , 2019 ) . Though such attempts have demonstrated certain efficacy in handling non time-series instance-wise data , they do not suffice for handling sequential dataset . There is also an alternative category that models feature acquisition as a Bayesian experimental design ( Ma et al. , 2019 ; Gong et al. , 2019 ) . However , the sequential decision making is for variable-wise feature acquisition and the problems are still non time-series tasks in nature . The key difference between all the aforementioned approaches and ours is that we tackle active feature acquisition problems with time-series data , where an active feature selection decision needs to be formed at each time step along the multi-step reinforcement learning trajectory . Therefore , the feature acquisition for our presented work needs to consider more complex information over model dynamics and control , apart from the static instance-wise features . 2.2 REPRESENTATION LEARNING IN POMDP . In complex tasks , policies trained upon different representations can even converge to different performance levels . Most conventional deep reinforcement learning approaches unifies the process of representation learning with policy training and results in policies trained in an end-to-end fashion ( Mnih et al. , 2013 ; Lillicrap et al. , 2016 ; Mnih et al. , 2016 ) . However , to accomplish the representation learning task , such models often engage trainable parameters which could come with considerable size and thereby result in significant degradation in sample efficiency . When considering problems with POMDPs where the state space is partially accessible to the agent , representation learning becomes an important and non-trivial research challenge . Among the existing literature , one prominent line of research tackles the representation learning for POMDP in an off-line fashion and thus resulting in multi-stage reinforcement learning . Higgins et al . ( 2016 ; 2017 ) adopt pretrained VAE models as a representation module to build agents with strong domain adaptation performance . The key difference between their work and ours is that they encode instance-wise image frames from POMDP domains where each image presents a partial view over the task environment , while our work considers cost-sensitive reinforcement learning with distinct partial observability , i.e. , the feature-level information is missing at each time step for the agent . We thus adopt a sequential representation learning approach to infer a more representative state information . Recently , there also emerged several works on sequential representation learning for POMDP ( Gregor et al. , 2019 ; Vezzani et al. , 2019 ) . However , most of the works utilize VAE training as an auxiliary task to jointly update the representation model with the policy learning loss . In our work , due to the high acquisition cost to observe the features , we adopt an off-line representation learning setting . Also , our proposed representation learning is model-based , where the model learns to impute the missing features with such attempt yielding significant benefit to derive high-quality representation for policy training . 3 METHODOLOGY . 3.1 TASK SETTING . In this section , we formally define the problem settings for the task of jointly learning the task and feature acquisition policy . To this end , we define the active feature acquisition POMDP , a rich class of discrete-time stochastic control processes generalizing standard POMDPs : Definition 1 ( AFA-POMDP ) . The active feature acquisition POMDP is a tuple M = 〈S , A , T , O , R , C , γ〉 , where S is the state space and A = ( Af , Ac ) is a joint action space of feature acquisition actionsAf and control actionsAc . The transition kernel T : S ×Ac×Af → PS maps any joint action a = ( af , ac ) in state s ∈ S to a distribution PS over next states . In each state s , when taking action af , the agent observes xp = x ( af ) , i.e. , a subset of the features x = ( xp , xu ) ∼ O ( s ) indicated by af , whereO ( s ) is a distribution over possible feature observation for state s and xu are features not observed by the agent . When taking a joint action , the agent obtains rewards according to the reward functionR : S ×Ac → R and pays a cost of C : S ×Af → R+ for feature acquisition . Rewards and costs are discounted by the discount factor γ ∈ [ 0 , 1 ) . Simplifying assumptions For simplicity , we assume that x consists of a fixed number of features Nf for all states , that Af = 2 [ Nf ] is the power set of all the Nf features , and that xp ( af ) consists of all the features in x indicated by the subset af ∈ Af . Note that the feature acquisition action for a specific application can take various different forms . For instance , in our experiments in Section 4 , for the Sepsis task , we define feature acquisition as selecting a subset over possible measurement tests , whereas for the Bouncing Ball+ task , we divide an image into four observation regions and let the feature acquisition policy select a subset of observation regions ( rather than raw pixels ) . Please also note that while in a general AFA-POMDP , the transition between two states depends on the joint action , we assume in the following that it depends only on the control action , i.e. , T ( s , ac , af ′ ) = T ( s , ac , af ) for all af ′ , af ∈ Af . While not true for all possible applications , this assumption can be a reasonable approximation for instance for medical settings in which tests are non-invasive . For simplicity we furthermore assume that acquiring each feature has the same cost , denoted as c , i.e. , C ( af , s ) = c |af | , but our approach can be straightforwardly adapted to have different costs for different feature acquisitions . Objective We aim to learn a policy which trades off reward maximization and the cost for feature acquisition by jointly optimizing a task policy πc and a feature acquisition policy πf . That is , we aim to solve the optimization problem max πf , πc E ∞∑ t=0 γt ( R ( xt , act ) − |Af |∑ i c · I ( af ( i ) t ) ) , ( 1 ) where the expectation is over the randomness of the stochastic process and the policies , af ( i ) t denotes the i-th feature acquisition action at timestep t , and I ( · ) is an indicator function whose value equals to 1 if that feature has been acquired . Note that the above optimization problem is very challenging : an optimal solution needs to maintain beliefs bt over the state of the system at time t which is a function of partial observations obtained so far . Both the the feature acquisition policy πf ( aft | bt ) and the task policy i.e. , πc ( act | bt ) depend on this belief . The information in the belief itself can be controlled by the feature acquisition policy through querying subsets from the features xt and hence the task policy and feature acquisition policy itself strongly depend on effectiveness of the feature acquisition policy .

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