Leveraging Factored Action Spaces for Efficient Offline Reinforcement Learning in Healthcare

Shengpu Tang, Maggie Makar, Michael Sjoding, Finale Doshi-Velez, Jenna Wiens

Keywords: reinforcement learning, offline rl, action space factorization, bias-variance trade-off, domain knowledge, healthcare, sepsis

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.

TL;DR: We present a simple way to incorporate action space factorization into offline RL by learning linearly decomposable Q-functions, supported by both theoretical and empirical results.

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2305.01738/code)