[Paper Review] Leveraging Factored Action Spaces for Efficient Offline Reinforcement Learning in Healthcare

Shengpu Tang, Maggie Makar+3 more|arXiv (Cornell University)|May 2, 2023

Health Systems, Economic Evaluations, Quality of Life4 citations

TL;DR

This paper proposes a linear Q-function decomposition for offline reinforcement learning in healthcare by exploiting factored action spaces—where actions are composed of independent sub-actions—enabling better sample efficiency and improved generalization in low-data regimes. The method achieves superior policy performance on both synthetic and real-world sepsis treatment datasets by reducing bias and variance, especially in underexplored action combinations.

ABSTRACT

Many reinforcement learning (RL) applications have combinatorial action spaces, where each action is a composition of sub-actions. A standard RL approach ignores this inherent factorization structure, resulting in a potential failure to make meaningful inferences about rarely observed sub-action combinations; this is particularly problematic for offline settings, where data may be limited. In this work, we propose a form of linear Q-function decomposition induced by factored action spaces. We study the theoretical properties of our approach, identifying scenarios where it is guaranteed to lead to zero bias when used to approximate the Q-function. Outside the regimes with theoretical guarantees, we show that our approach can still be useful because it leads to better sample efficiency without necessarily sacrificing policy optimality, allowing us to achieve a better bias-variance trade-off. Across several offline RL problems using simulators and real-world datasets motivated by healthcare, we demonstrate that incorporating factored action spaces into value-based RL can result in better-performing policies. Our approach can help an agent make more accurate inferences within underexplored regions of the state-action space when applying RL to observational datasets.

Motivation & Objective

• To address the challenge of combinatorial action spaces in offline reinforcement learning, where standard methods fail to generalize across rare sub-action combinations.
• To develop a linear Q-function decomposition that leverages the inherent factorization of action spaces to improve value function approximation.
• To demonstrate that this approach enhances sample efficiency and policy performance even when theoretical assumptions for unbiased estimation are only partially satisfied.
• To validate the method on both simulated sepsis treatment environments and real-world MIMIC-III clinical data, showing improved robustness in low-data settings.

Proposed method

• The method decomposes the Q-function into a linear combination of sub-action Q-values, parameterizing Qπ(s,a) as the sum of individual Q functions over each sub-action space.
• It uses a factored action space A = A₁ × ⋯ × A_D, where each action a is a vector of sub-actions ad ∈ Ad, enabling modular value function learning.
• The Q-function is approximated using a neural network with D × |Ad| output heads, one per sub-action, instead of 2^D heads for all combinations, reducing model complexity.
• The approach incorporates domain knowledge about action independence and synergy, allowing better generalization to rare or unobserved sub-action combinations.
• It is compatible with existing offline RL algorithms such as BCQ, enabling integration into established pipelines without architectural overhaul.
• The method includes value clipping and early stopping during training to stabilize learning and ensure bounded function approximation.

Experimental results

Research questions

• RQ1 Under what conditions does linear Q-function decomposition lead to an unbiased estimate of the true Q-function in factored action spaces?
• RQ2 How does the proposed method improve sample efficiency in offline RL when data are limited and action combinations are underexplored?
• RQ3 Can linear decomposition still yield optimal policies even when the decomposition is biased, particularly when sub-actions exhibit reinforcing effects?
• RQ4 How does the method compare to non-factored baselines in terms of policy performance on real-world clinical datasets with sparse treatment combinations?
• RQ5 To what extent does incorporating domain knowledge about action structure improve policy generalization in offline healthcare RL?

Key findings

• The proposed method outperforms non-factored baselines in both simulated sepsis and real MIMIC-III datasets, especially under low data regimes, with a 95th percentile WIS score of 98.5 on the validation Pareto frontier.
• On the MIMIC-III sepsis dataset, the proposed approach achieved a test WIS value of 96.2, significantly outperforming the baseline (92.1) while maintaining comparable ESS (effective sample size).
• The method demonstrated improved generalization to rare treatment combinations, learning policies that better capture the effects of underrepresented sub-action pairs.
• Even when theoretical conditions for zero bias were violated (e.g., non-additive rewards), the method still reduced variance and improved sample efficiency.
• Model selection with ESS cutoffs ≥200 showed that the proposed method consistently achieved higher test WIS values than the baseline across all thresholds.
• The approach achieved better bias-variance trade-offs, with reduced variance in value estimates leading to more stable and reliable policy learning.