Reinforcement Learning with Efficient Active Feature Acquisition

Abstract: Solving real-life sequential decision making problems under partial observability involves an exploration-exploitation problem. To be successful, an agent needs to efficiently gather valuable information about the state of the world for making rewarding decisions. However, in real-life, acquiring valuable information is often highly costly, e.g., in the medical domain, information acquisition might correspond to performing a medical test on a patient. Thus it poses a significant challenge for the agent to learn optimal task policy while efficiently reducing the cost for information acquisition. In this paper, we introduce a model-based framework to solve such exploration-exploitation problem during its execution. Key to the success is a sequential variational auto-encoder which could learn high-quality representations over the partially observed/missing features, where such representation learning serves as a prime factor to drive efficient policy training under the cost-sensitive setting. We demonstrate our proposed method could significantly outperform conventional approaches in a control domain as well as using a medical simulator. 

Reviews and author responses:

This paper proposes a reinforcement learning + representation learning approach for simultaneously learning a control policy and feature acquisition policy in environments where feature observation is costly. The authors formulate an approach for learning time series latent variable models that incorporate information from both observation and action histories. They demonstrate through a series of experiments that their approach leads to better imputation (i.e., filling in missing values) and better rewards. Strengths:
- The problem is an important one: feature acquisition is indeed a major problem in the healthcare space, for example, and the ability to learn policies that effectively trade off information gathering/feature acquisition costs with environment control.
- The proposed VAE model seems to learn what it's supposed to learn - Figure 4 was a nice way to qualitatively capture exactly how the model is trading off feature acquisition costs with control rewards.
- 

Weaknesses:
- My main concern is that there is already a fairly similar work by Igl et al (2018) [1] that proposed a sequential VAE and had some promising results on a related set of tasks. I would expect to see their DVRL algorithm as well as something like "Deep Recurrent Q-Learning" [2] implemented as baselines before I agree with any claims around significant contributions or novelty.
- There are a few critical experimental details that are left ambiguous. For example, after reading through the manuscript and searching explicitly for information on the "end-to-end" baseline, it is not clear to me exactly what model was used and how it selectively acquired features. 
- The requirement that the VAE be pretrained offline is a fairly restrictive one. It's not clear to me how realistic this is. Also, does that preclude any improvement to the VAE model while the agent is interacting with the environment online? That seems like a missed opportunity.
- There was little, if any, discussion about the potential ethical impacts of the work. This would be a requirement, in my mind, for acceptance.

Minor comments:
- There are a number of minor grammatical errors throughout. The paper would benefit from some detailed proofreading. I tried to keep track of the errors but it quickly got out of hand.
- Citations are formatted awkwardly (use \citep rather than \citet or \cite).
- I believe $\mathcal{A}_f$ in equation (1) should be $\mathcal{A}^f$ in order to be consistent with prior notation.
- It wasn't fully clear to me based on the text what the difference is between the "full loss" and the "partial loss" (I understand that one is over "the entire features" and the other "only applies to the observed features", but how do you calculate the loss over "the entire features"?)

[1] Igl, Maximilian, et al. "Deep variational reinforcement learning for POMDPs." International Conference on Machine Learning. PMLR, 2018.
[2] Hausknecht, Matthew, and Peter Stone. "Deep recurrent q-learning for partially observable mdps." 2015 aaai fall symposium series. 2015. The paper proposes an interesting approach to tackle an important problem, but there are some critical comparisons missing related to prior work that preclude claims of significance, contribution, or novelty.

Post-response: 
A3-1: 
I'm still not convinced about the novelty of the VAE approach. I understand that the sequential VAE in DVRL does not have ground truth labels and thus requires inference of latent states, but that's because the assumptions don't allow for imputation of missing features/supervised learning. It could easily be adapted to such a setting, however, as representation learning with ground truth labels will in general be easier/more sample efficient than learning without. Similarly, I'm more impressed by DVRL's ability to continuously improve its state inference model online, which seems like a more difficult approach than the two-staged approach suggested by the authors here.

A3-2: 
The authors claim that "when dealing with AFA-POMDP, if DVRL wants to learn a high quality representation that could effectively impute the missing features, it would need the access to the fully observed data to update itself during online training time based on some loss similar to Eq(5) [...] such an attempt would consume a very large amount of feature acquisition cost to get the fully observed data, which is way larger than adopting a multi-stage RL as we propose where we only acquire 2K trajectories to train the sequential VAE model." --> It is not clear to me from the paper that, in general, online feature acquisition costs would be any worse than the cost of observing all features offline and collecting those in a batch dataset, as proposed by the paper. By contrast, consider a scenario in which the state includes thousands of noisy and useless features but a few useful and informative features. Then doing the batch representation learning as the authors suggest would lead to unnecessarily high observation costs whereas something like the online DVRL approach could potentially learn much more efficiently. So this claim feels unsubstantiated.

A3-4: 
I appreciate the authors' statement that "the representation learning model could be updated in a continual learning fashion by the new data as well as the old data and then we could leverage policy transfer [6] techniques to effectively update the policy parameters without inquiring much fully observed data for training the pipeline." But this is all still in the hypothetical. At the end of the day, the proposed method as written is fixed after some number of observations.

A3-5: 
I appreciate the authors' diligence on the Ethics Statement.

Separately, there is a recent paper at NeurIPS 2021 that seems closely aligned with the proposed work. It shouldn't impact the novelty of the present work, as the two papers are essentially contemporaneous, but may be of interest to the authors: https://neurips.cc/Conferences/2021/Schedule?showEvent=26563

In light of the comments and clarifications, I'm willing to increase my score to a 6.

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. Strengths:
- The authors approach outperforms the other VAE baselines, both in terms of reward, and in terms of MSE in inferring the unobserved state.
- The problem of joint control and feature acquisition is interesting and important.

Weaknesses:
- The authors claim that AFA-POMDPs generalize POMDPs is false—an AFA-POMDP is a special case of a POMDP. An example of past work using POMDPs for feature acquisition can be found here: Shi and Cain. Cost-sensitive feature acquisition and classification. 2007.
- The authors approach of using a neural network-based model to infer the belief state in a POMDP was done here: Karkus et al. QMDP-Net: Deep Learning for Planning under Partial Observability, 2017. 
- The domains seem very simple, like they probably could be solved with planning approaches like QMDP-net. They don't really show off the benefit of using RL.
- Despite having much higher MSE, the baselines are very competitive on the tasks. NonSeq-ZI (partial) is almost as good as the trained method on bouncing ball and NonSeq-ZI (full) performs better in terms of mortality on sepsis, but this is not mentioned.
- The authors approach makes a quite strong assumption about having access to fully observed data for training. The paper doesn't really discuss this as a limitation.

Post-response: The authors have addressed some of my concerns, and I have raised my score to a 5. I still believe the paper is not ready for publication. The paper makes some incorrect claims, does not cite relevant past work, and does experiments on very simple problems.

In this paper the authors propose an approach for simultaneously learning how to explore more efficiently in POMDPs via targeted feature acquisition, and learning a reward-maximizing control policy, balancing the cost of feature acquisition with the expected reward. Learning is done via a VAE framework which combines a belief inference model and an observation decoder, with a key innovation being that inference is done as a sequential process. Results comparing this approach to other variational inference approaches show the proposed framework reaches better performance with lower cost (particularly, number of acquired features). This is an interesting paper which tackles a variation on the perennial challenge of exploration vs. exploitation in POMDPs. The overall approach seems reasonable to me and I liked what I identify to be two key ideas in the paper, which are learning the feature acquisition policy *and* the target policy simultaneously, and introducing the notion of sequence in the VAE inference framework. However, I do have some questions, 

- My first concern regarding the proposed approach is that it seems wildly expensive to train even for modestly sized problems, and even in the fraught landscape of POMDP learning strategies. I have not seen concrete estimations of sample complexities required in training in the supplementary material, how do they compare to other approaches for solving POMDPs (more on that in a bit). On that note, I am somewhat concerned by the pretraining of the VAEs and the fact that it is then essentially fixed during learning, implying that while there's a sequential component to it, it isn't really an RL approach for hidden state imputation and therefore might be very brittle in practice, especially when a lot of distributions are unknown prior to training. 

- Speaking of other approaches for solving POMDPs, a problem which can be spotted already in the related work section is the fact that the author restrict themselves essentially to a small pool of works which pursue a POMDP solving strategy that is largely similar in spirit to that proposed in this paper. While this focus is understandable, it not only does a disservice to a great deal of prior work on efficient exploration and representation learning in POMDPs, it also substantially limits the set of useful baselines the authors compare against. What about other strategies for learning behavior policies and work on balancing exploration vs. exploitation in POMDPs in general? Or hierarchical abstractions? The current set of experiments seems to mostly show that introducing the notion of sequence in the VAE formulation is beneficial, which is neither negligible nor surprising, but it's not obvious to me that the proposed (again, monstrously complex) approach outperforms simpler approaches under most settings.

- The term "cost" here mostly refers to number of features acquired. However, it seems to me that such an approach is overly simplistic, given that certain features are vastly more expensive to acquire than others (consider the difference between a simple blood test and a 24 hour EEG monitoring experiment). Is the proposed approach resilient to varying acquisition costs?
 I find this paper very interesting but am worried about the complexity of the proposed approach and how difficult it would be to apply in most cases. I am also concerned that the baselines compared against are insufficient to get an actual sense of how well this approach compares to other approaches, and what's the sample complexity and computational cost (rather than just feature acquisition cost) associated with it. I am also concerned that essentially learning the VAE component offline makes the approach brittle in practice. 

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

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

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

Decision: Reject