Source-linked AI summary
Model Selection for Offline Reinforcement Learning: Practical Considerations for Healthcare Settings
Shengpu Tang, Jenna Wiens
TL;DR
Healthcare offline RL needs model selection, but candidate policies cannot always be run in the real environment and OPE methods bring practical costs. The paper evaluates OPE-based selection for sepsis treatment and finds FQE provides the strongest ranking, then proposes a two-stage approach to reduce unnecessary computation.
Problem
Offline RL model selection lacks a well-accepted validation framework because policies cannot always be run in the real environment, while OPE methods’ usefulness and practical trade-offs remain uncertain.
Source: §1 · p. 2Method
The paper analyzes four OPE methods for ranking candidate policies and proposes a two-stage approach that shortlists candidates with a cheaper estimator before applying a more accurate estimator.
Source: §1 · p. 2Results
On simulated sepsis treatment, FQE consistently gives the most accurate ranking over candidate policies and is most robust to varied sample sizes and behavior policies.
Source: §6 · p. 18Takeaways & Limitations
The two-stage approach outperforms simple averaging and more sophisticated doubly robust methods while avoiding expensive computation on low-quality policies.
Source: §6 · p. 18Experiments use a simulated RL task, whose discrete actions, short horizons, noiseless observations, lack of missingness, and sparse terminal rewards may limit generalizability to real observational data and other RL problems.
Source: §6 · p. 18
Abstract
from arXiv · showhide
Reinforcement learning (RL) can be used to learn treatment policies and aid decision making in healthcare. However, given the need for generalization over complex state/action spaces, the incorporation of function approximators (e.g., deep neural networks) requires model selection to reduce overfitting and improve policy performance at deployment. Yet a standard validation pipeline for model selection requires running a learned policy in the actual environment, which is often infeasible in a healthcare setting. In this work, we investigate a model selection pipeline for offline RL that relies on off-policy evaluation (OPE) as a proxy for validation performance. We present an in-depth analysis of popular OPE methods, highlighting the additional hyperparameters and computational requirements (fitting/inference of auxiliary models) when used to rank a set of candidate policies. We compare the utility of different OPE methods as part of the model selection pipeline in the context of learning to treat patients with sepsis. Among all the OPE methods we considered, fitted Q evaluation (FQE) consistently leads to the best validation ranking, but at a high computational cost. To balance this trade-off between accuracy of ranking and computational efficiency, we propose a simple two-stage approach to accelerate model selection by avoiding potentially unnecessary computation. Our work serves as a practical guide for offline RL model selection and can help RL practitioners select policies using real-world datasets. To facilitate reproducibility and future extensions, the code accompanying this paper is available online at https://github.com/MLD3/OfflineRL_ModelSelection.
1. Introduction
Healthcare RL often relies on historical data, making model selection important but difficult: policies cannot always be tested in the real environment. The paper examines OPE-based selection and proposes a two-stage approach to balance ranking quality and computational cost.
1. Introduction
- Healthcare offline RL uses historical data, while complex state and action spaces can require function approximators whose model selection is crucial to prevent overfitting.The paper notes a lack of a well-accepted training-validation framework for offline RL.Source: §1 · p. 2
- Testing candidate policies in real healthcare settings may be costly or dangerous, so offline RL model selection must rely on OPE, whose usefulness for ranking policies is uncertain.OPE methods add hyperparameters and auxiliary models, which can substantially increase computational cost.Source: §1 · p. 2
- The paper analyzes four OPE methods and their hyperparameters, auxiliary models, and computational requirements when ranking policies, evaluating them on sepsis treatment.The experiments consider early stopping and neural-network architecture selection in both tabular and function-approximation settings.Source: §1 · p. 2, Generalizable Insights about Machine… · p. 3
Generalizable Insights about Machine Learning in the Context of Healthcare
- The proposed two-stage approach first uses a more efficient, less accurate estimator to shortlist policies, then a more accurate, less efficient estimator to choose the final policy.The approach aims to avoid expensive computation on candidates unlikely to be selected.Source: §1 · p. 2, Generalizable Insights about Machine… · p. 3
- The authors frame their practical lessons as relevant beyond healthcare to high-stakes domains without accurate simulators, including tutoring, dialogue, and online advertising.They present the work as informing future research and fairer comparisons in offline RL.Source: Generalizable Insights about Machine… · p. 3
2. Offline Reinforcement Learning - Problem Setup
The setup considers offline, model-free value-based RL: policies are learned from previously collected observational data rather than interaction with the environment. The objective is to learn a policy with low suboptimality, while OPE estimates its performance using that historical data.
2. Offline Reinforcement Learning - Problem Setup
- An MDP is specified by state and action spaces, transition and reward models, an initial-state distribution, and a discount factor; the setup uses discrete actions and discrete or continuous states.A policy maps states to action distributions, and policy value is expected cumulative reward from the initial-state distribution.Source: §2 · p. 3
- Offline RL uses previously collected observational transitions or episodes from behavior policies, rather than interacting with the environment.Because exploration is limited, the goal is a policy with low suboptimality rather than necessarily identifying an optimal policy.Source: §2 · pp. 3–4
- The paper focuses on model-free, value-based fitted Q iteration for policy learning and uses OPE to estimate candidate-policy performance from data collected by other policies.OPE estimates counterfactual outcomes because the historical data reflect actions taken by behavior policies.Source: §2 · p. 4
3. Offline RL Model Selection using OPE
The section adapts validation-based model selection to offline RL by using OPE to rank candidate policies, then analyzes the methods’ practical costs and proposes a two-stage alternative. The OPE methods differ in their auxiliary-model requirements, hyperparameters, and computational demands.
3.1. A Model Selection Framework for Offline RL
- Offline RL model selection splits observational data into training and validation sets, learns candidate policies on training data, and ranks them by OPE estimates on validation data.The selected policy is the one with the highest estimated validation performance.Source: Figure 1 · p. 4, §3.1 · p. 5
3.2. Using OPE for Model Selection
- Validation returns cannot directly score a policy that may choose different actions from the behavior policy, so offline model selection requires OPE.The section compares popular methods and their extra hyperparameters, auxiliary models, and computational requirements.Source: §3.2 · p. 5
3.2.1. Weighted Importance Sampling (WIS)
- WIS estimates policy value by reweighting episodes using relative likelihood under the evaluation and behavior policies, with policy softening used to address zero-weight episodes.WIS is biased but consistent and can provide more stable performance than ordinary importance sampling; its softening parameter ε controls the softened policy.Source: §3.2.1 · pp. 5–6
3.2.2. Approximate Model (AM)
- AM learns an approximate environment model and evaluates a policy using that model, but rollout horizon H and auxiliary model choices affect its estimates and cost.For continuous states, simplifying assumptions can reduce model complexity, while multiple rollouts to address stochasticity require an additional hyperparameter and more computation.Source: §3.2.2 · pp. 6–8
3.2.3. Fitted Q Evaluation (FQE)
- FQE learns a policy’s Q-function from historical data, but fitting separate networks across H iterations for each of K policies makes it substantially more expensive than WIS or AM.Bootstrapping prevents sequential iterations of FQE for a policy from being parallelized; different policies can be evaluated in parallel.Source: §3.2.3 · p. 8
3.2.4. Weighted Doubly Robust (WDR) Estimator
- WDR combines importance sampling with Q- and V-value estimates at every time step, but it can be less accurate and have larger errors than its constituent methods.Its two variants use AM or FQE value estimates; WDR inherits the hyperparameters of its components and requires Q and V estimates across all N state-action pairs.Source: §3.2.4 · p. 9
3.3. Combining Multiple OPE Estimators for Model Selection
- WIS and AM are more efficient but tend to be less accurate, whereas FQE tends to be more accurate but less efficient, motivating a two-stage ranking procedure.An efficient estimator first identifies a subset; a more accurate estimator then ranks that subset. The subset fraction α controls cost and the relative influence of the estimators.Source: §3.3 · p. 10
- Average score, average ranking, and WDR are comparison baselines; the two averaging approaches do not provide the computational savings of two-stage selection and may be vulnerable to high-variance outliers.WDR is a hybrid of WIS with AM or FQE.Source: §3.3 · p. 10
4. Experimental Setup
The experiments evaluate whether OPE validation scores can identify strong policies among candidates learned from simulated sepsis data. They compare OPE rankings with known policy values using rank correlation and regret.
4. Experimental Setup
- The experiments generate training and validation datasets, learn candidate policies on training data, estimate validation performance with OPE, and compare rankings and regret with true policy values.Candidate policies vary through training hyperparameters.Source: §4 · p. 11
4.1. Environment & Dataset Generation
- The sepsis simulator represents patients with discrete or continuous states, eight treatment actions, terminal survival or death rewards, and episodes truncated at 20 steps.Ten training-validation dataset pairs, each with 10,000 episodes generated by a uniformly random policy, provide the data; the authors note this generation process makes assumptions that may not hold in practice.Source: §4.1 · p. 11
4.2. Implementation Details
- FQI uses a neural-network function approximator, and experiments repeat each run 10 times with ε = 0.01 and H = 20 as default OPE hyperparameters.The default networks have one hidden layer with 1,000 neurons and ReLU activation; these settings apply unless otherwise specified.Source: §4.2 · pp. 11–12
4.3. Evaluation
- Evaluation uses Spearman’s rank correlation and regret, with regret@n measuring the smallest regret among the n highest-scoring policies.True policy values come from analytic evaluation using the simulator’s ground-truth MDP parameters; rank correlation and regret@n assess whether OPE can identify promising subsets for two-stage selection.Source: §4.3 · p. 12
5. Experiments & Results
Across experiments on sepsis-treatment policy selection, FQE ranked candidates best, while a two-stage WIS–FQE strategy balanced regret and computation. Sensitivity analyses found the two-stage strategy remained competitive across validation-data conditions, with very small samples as an exception.
5.1. Model Selection of Neural Architectures & Training Hyperparameters
- The experiments compared OPE-based model selection across neural-network architectures and training hyperparameters, evaluating 96 FQI-trained Q-networks on validation data.Ranking quality and regret were measured; the main experiments used continuous states, with discrete-state results in Appendix D.2.Source: §5 · p. 12, §5.1 · p. 13, Table 2 · p. 13
5.1.1. Comparison Across OPE Methods
- ρ = 0.89 for FQE, the best ranking; WIS reached 0.74, AM 0.66, and WDR-FQE 0.64.WDR-FQE sometimes had regret above 0.5, while all methods had near-zero regret@5 and regret@10; FQE and WIS consistently had the lowest regret@1.Source: §5.1.1 · p. 13, Figure 2 · p. 14
5.1.2. Combining Multiple OPE Estimators
- A WIS-first two-stage strategy had comparable or lower regret than FQE across subset sizes, with the lowest regret at 24 and 48, and could reduce computation.With minimal additional computation over WIS, it achieved lower regret than FQE and potentially greater computational savings; the figure normalizes cost to full FQE over 96 candidates for illustration.Source: §5.1.2 · p. 14, Figure 3 · p. 15
5.2. Sensitivity Analyses
- Experiments varied auxiliary hyperparameters, validation sample size, and behavior policy while holding training procedures and candidate policies fixed.They compared WIS, AM, FQE, and WIS–FQE two-stage selection with a subset of 24; additional discrete- and continuous-state analyses are in Appendix D.2.Source: §5 · p. 12, §5.2 · p. 15
5.2.1. Sensitivity to OPE Hyperparameters
- AM was most robust to H but had higher regret; WIS and FQE achieved low regret at reasonable ε or H values.WIS had its lowest regret at ε = 0.1; FQE was more stable at H ≥ 10, while very small or large ε increased WIS regret variance.Source: Figure 4 · p. 15, §5.2.1 · p. 16
5.2.2. Sample Size of Validation Data
- With at least 1,000 validation episodes, two-stage selection matched or slightly outperformed FQE in regret while reducing computation; very small samples were an exception.At 500 episodes, WIS’s higher variance could make first-stage pruning less reliable.Source: Figure 5 · p. 16, §5.2.2 · p. 17
5.2.3. Type of Behavior & Extent of Exploration
- Less exploratory behavior could increase regret, but two-stage selection maintained or improved on FQE’s regret while reducing computation, including with mixed behaviors.Random behavior generally produced slightly lower regret than ε-greedy behavior; WIS still had low regret under mixed behavior policies.Source: Figure 5 · p. 16, §5.2.3 · p. 17
6. Conclusion
The paper uses OPE to estimate validation performance for offline RL model selection and finds FQE ranks sepsis-treatment policies most accurately, despite its high computational cost. A two-stage approach combines OPE estimators while avoiding expensive evaluation of low-quality policies; conclusions remain bounded by the simulated setting and other study limitations.
6. Conclusion
- FQE ranked candidate policies most accurately and most robustly across sample sizes and behavior policies in the simulated sepsis task.The proposed two-stage approach outperformed simple averaging and doubly robust methods while avoiding expensive computation on low-quality policies.Source: §6 · p. 18
- The two-stage approach balances computation and policy performance by combining OPE estimators and avoiding costly evaluation of low-quality policies.The conclusion notes that the approach outperformed simple averaging and more sophisticated doubly robust methods.Source: §6 · p. 18
- The experiments use a simulator whose discrete actions, short horizons, noiseless observations, and sparse terminal rewards may limit generalization to real observational data and other RL problems.The simulator also cannot address problem-design choices such as defining an action space or meaningful initial states.Source: §6 · p. 18
- The study considers four commonly used OPE methods, uses the same validation data to fit auxiliary models and calculate estimates, and trains policies with vanilla FQI on datasets with sufficient state-action coverage.Reusing validation data may violate independence assumptions required for OPE theoretical guarantees; the study also omits some accuracy-improving estimator variants.Source: §6 · pp. 18–19
- The authors frame their practical guidance as a basis for reporting model-selection techniques to improve reproducibility and enable fairer comparisons.They encourage further healthcare RL work using the released code.Source: §6 · p. 19
Appendix A. Pseudocode for FQI & FQE
- Appendix A presents pseudocode for FQI and FQE, with each algorithm taking a dataset, function class, and maximum iteration count as inputs.FQE also takes the policy to be evaluated as input.Source: Appendix A · p. 25
Appendix B. OPE Computational Requirements
The appendix explains how OPE auxiliary models are obtained and summarizes their additional model and inference requirements. The modeling approach differs between discrete and continuous state spaces.
Appendix B. OPE Computational Requirements
- WIS requires a behavior-policy model, whereas AM requires transition and reward models.The appendix assumes a discrete action space and leaves continuous action spaces to future work.Source: Appendix B · p. 25
- For discrete state spaces, auxiliary probability or value tables can be estimated by maximum-likelihood count-based averaging over the dataset.The relevant state-action and state-action-next-state counts can be accumulated in one pass through the data.Source: Appendix B · p. 25
- For continuous states, auxiliary models use parameterized functions and are trained as supervised classification or regression tasks, independently of policy learning.The models may require repeated inference for different inputs.Source: Appendix B · p. 25
- Table 3 details auxiliary supervised models required by each OPE method for continuous-state problems.Table 4 summarizes computational requirements for K candidate policies, given m episodes, N transitions, and horizon H.Source: Table 3 · p. 26, Table 4 · p. 26
Appendix C. Sepsis Simulator - Details
The sepsis experiments use a simulator modeled after patient physiology, with a defined treatment action space, state representation, and terminal outcomes. The simulator's dynamics and reward design establish the scope of the evaluation setting.
Appendix C. Sepsis Simulator - Details
- The simulator offers eight actions formed from combinations of antibiotics, vasopressors, and mechanical ventilation.Treatments can affect vital signs and may raise or lower their values with prespecified probabilities.Source: Appendix C · p. 26
- Patient states encode diabetes status, four ordinal vital signs, and three indicators for previously administered treatments.The experiments use either 1,440 discrete states or a 21-dimensional one-hot feature vector, with separate discharge and death absorbing states.Source: Appendix C · p. 26, Table 5 · p. 27
- The initial-state distribution assigns diabetes status a probability of 0.2 and samples uniformly from 606 eligible nonterminating states.Eligible states include states with treatments already enabled.Source: Appendix C · p. 26
- Treatments are applied sequentially and can be added or withdrawn, with effects on different state variables.The transition-dynamics details are referenced in Table 6.Source: Appendix C · p. 27, Table 6 · p. 27
- Discharge requires normal vital signs and stopped treatments, while death occurs with at least three abnormal vital signs.Episodes assign terminal rewards of +1 for survival and −1 for death, followed by zero reward in the absorbing state.Source: Appendix C · p. 27
D.1. Model Selection using FQI Values & TD Errors
FQI values and RMS-TDE perform poorly as validation scores for offline RL model selection, with weak ranking correlation and high regret. The early-stopping analysis considers policies along an FQI training run and focuses on WIS, AM, and FQE.
D.1. Model Selection using FQI Values & TD Errors
- FQI values use the learned Q-network's predicted values as a validation score, while RMS-TDE is the root-mean-squared TD-learning objective calculated over validation transitions.The corresponding mean-squared TD error is the loss without the square root.Source: §D.1 · p. 28
- FQI values and RMS-TDE achieved Spearman’s rank correlation |ρ| < 0.3 and high top-5/top-10 regret, underperforming the OPE methods studied.Figure 6 shows scatter plots of true policy value against FQI values and RMS-TDE, plus top-n regret for the two scores.Source: §D.1 · p. 28, Figure 6 · p. 28
D.2. Sensitivity Analyses for Early Stopping
- For early stopping, the candidate policies are those derived from Q-functions along an FQI training run; the analysis focuses on WIS, AM, and FQE.WDR is excluded because its large variance heavily skews the plots.Source: §D.2 · p. 29
D.2.1. Sensitivity to OPE Hyperparameters
The analysis tests how a priori OPE hyperparameters affect model-selection regret. Results differ by state setting: WIS and FQE are robust in the tabular case, while WIS and FQE have distinct favorable settings in the continuous case.
D.2.1. Sensitivity to OPE Hyperparameters
- The experiment varies policy softening ε and evaluation horizon H for AM/FQE, while keeping continuous-state auxiliary-model hyperparameters fixed.In the tabular setting, AM is analytic and has no hyperparameters, so it is omitted.Source: §D.2.1 · p. 29
- The results compare discrete-state and continuous-state sepsis settings.Source: §D.2.1 · p. 29
- In the tabular setting, WIS and FQE maintain regret below 0.1 across reasonable ε and H values; in continuous states, WIS is best at low ε and FQE at H = 20.AM is most robust to H in continuous states, but has higher regret that does not decrease as H increases; longer FQE horizons may overfit auxiliary models when approximation errors are present.Source: §D.2.1 · p. 29, Figure 7 · p. 29
D.2.2. Size of Validation Set
The appendix examines how validation data size and behavior affect OPE model selection, then places the work among related research and gives a probability result for random first-stage pruning.
D.2.2. Size of Validation Set
- FQE consistently achieves lower regret with larger validation sets in both state settings; WIS shows a weaker trend and noisier estimates, while AM is affected by model misspecification.Source: §D.2.2 · p. 30, Figure 8 · p. 30
D.2.3. Type of Behavior Policy and Extent of Exploration
- More exploratory random behavior generally yields slightly lower regret; FQE and AM remain stable across behavior settings and usually have the lowest regret with smaller variance.WIS is not heavily affected by mixed-behavior policy misspecification, though its median regret is slightly higher for mixed behavior in the tabular setting.Source: Figure 8 · p. 30, §D.2.3 · p. 31
Appendix E. Additional Related Work
- Prior work studied FQE for offline-RL hyperparameter selection and OPE benchmarks, while this paper compares estimators for practical healthcare model selection.The cited examples include Paine et al. (2020) and Fu et al. (2021).Source: Appendix E · p. 31
Appendix F. Additional Details of the Two-Stage Selection Procedure
- For a random ranking of C values, the theorem gives the probability that the first A positions contain at least one of the B largest values.When A + B > C, that probability is 1; otherwise the proof counts the complementary event.Source: Appendix F · p. 32
F.1. Performance of Two-Stage Selection
The two-stage analysis evaluates whether first-stage pruning retains good policies. FQE retains good policies more often than WIS, while the further analysis assumes a perfect second-stage ranking and leaves overall performance estimation for future work.
F.1. Performance of Two-Stage Selection
- FQE has a higher chance than WIS of retaining a good policy across α and β, and both outperform the random-ranking lower bound.The comparison uses empirical CDFs from 10 replication runs; the random-ranking bound indicates that pruning by half can still retain a good policy with high probability when some regret is allowed.Source: Figure 9 · p. 33, §F.1 · pp. 33–34
F.2. Further Analyses & Discussion
- The analysis assumes a perfect second-stage ranking and does not estimate overall performance because the two stages share validation data and their rankings are not independent.The authors leave these further analyses for future work.Source: §F.2 · p. 34