Table of Contents
Fetching ...

Q-function Decomposition with Intervention Semantics with Factored Action Spaces

Junkyu Lee, Tian Gao, Elliot Nelson, Miao Liu, Debarun Bhattacharjya, Songtao Lu

TL;DR

This work tackles the challenge of sample-efficient reinforcement learning in environments with large factored action spaces by introducing Q-function decomposition under intervention semantics and a practical action-decomposed RL framework. It formalizes projected action space MDPs and MB-FPI to analyze the theoretical properties and sample complexity of decomposed Q-functions, then implements these ideas in model-free and offline settings through AD-DQN and AD-BCQ. Empirically, the approach yields improved sample efficiency on online 2D control tasks and gains in offline sepsis treatment evaluation, outperforming baseline decompositions and demonstrating robust performance across discretized action spaces. The results highlight the potential of incorporating causal intervention semantics and action-projection structure to scale RL to complex, real-world decision problems.

Abstract

Many practical reinforcement learning environments have a discrete factored action space that induces a large combinatorial set of actions, thereby posing significant challenges. Existing approaches leverage the regular structure of the action space and resort to a linear decomposition of Q-functions, which avoids enumerating all combinations of factored actions. In this paper, we consider Q-functions defined over a lower dimensional projected subspace of the original action space, and study the condition for the unbiasedness of decomposed Q-functions using causal effect estimation from the no unobserved confounder setting in causal statistics. This leads to a general scheme which we call action decomposed reinforcement learning that uses the projected Q-functions to approximate the Q-function in standard model-free reinforcement learning algorithms. The proposed approach is shown to improve sample complexity in a model-based reinforcement learning setting. We demonstrate improvements in sample efficiency compared to state-of-the-art baselines in online continuous control environments and a real-world offline sepsis treatment environment.

Q-function Decomposition with Intervention Semantics with Factored Action Spaces

TL;DR

This work tackles the challenge of sample-efficient reinforcement learning in environments with large factored action spaces by introducing Q-function decomposition under intervention semantics and a practical action-decomposed RL framework. It formalizes projected action space MDPs and MB-FPI to analyze the theoretical properties and sample complexity of decomposed Q-functions, then implements these ideas in model-free and offline settings through AD-DQN and AD-BCQ. Empirically, the approach yields improved sample efficiency on online 2D control tasks and gains in offline sepsis treatment evaluation, outperforming baseline decompositions and demonstrating robust performance across discretized action spaces. The results highlight the potential of incorporating causal intervention semantics and action-projection structure to scale RL to complex, real-world decision problems.

Abstract

Many practical reinforcement learning environments have a discrete factored action space that induces a large combinatorial set of actions, thereby posing significant challenges. Existing approaches leverage the regular structure of the action space and resort to a linear decomposition of Q-functions, which avoids enumerating all combinations of factored actions. In this paper, we consider Q-functions defined over a lower dimensional projected subspace of the original action space, and study the condition for the unbiasedness of decomposed Q-functions using causal effect estimation from the no unobserved confounder setting in causal statistics. This leads to a general scheme which we call action decomposed reinforcement learning that uses the projected Q-functions to approximate the Q-function in standard model-free reinforcement learning algorithms. The proposed approach is shown to improve sample complexity in a model-based reinforcement learning setting. We demonstrate improvements in sample efficiency compared to state-of-the-art baselines in online continuous control environments and a real-world offline sepsis treatment environment.
Paper Structure (39 sections, 6 theorems, 27 equations, 15 figures, 1 table, 3 algorithms)

This paper contains 39 sections, 6 theorems, 27 equations, 15 figures, 1 table, 3 algorithms.

Key Result

proposition 1

Given a projected MDP $\mathcal{M}^{k}$ over action variables $\mathbf{A}_k$, the Q-function $Q_{\pi_{k}}(\mathbf{s}, \mathbf{a}_k)$ can be recursively written as, where $\pi_{k}:\mathcal{S} \rightarrow \mathcal{A}_k$ is a factored policy for $\mathbf{A}_k$.

Figures (15)

  • Figure 1: Decomposable Structures in Factored MDPs. The diagrams show factored MDPs, where the circles, squares, and diamonds represent state, action, and reward variables. Fully Separable Structure: Fig \ref{['f11']} shows a factored MDP that can be fully separable into two independent MDPs, considered in the previous work tang2022leveragingseyde2022solving. Separable Effects: Fig \ref{['f12']} shows a factored MDP that has non-separable dynamics and rewards. However, the effects of factored actions are non-interacting, and we study this structure with intervention semantics. Non-separable Structure: Fig \ref{['f13']} shows a non-separable factored MDP.
  • Figure 2: Comparing the values in the test environment on three action spaces. The X-axis shows 2 million steps of training (2000 episodes) and the Y-axis is the value evaluated from a test environment. The plot aggregates 10 trials with the random seeds from 1 to 10 for training and 1001 to 1010 for testing.
  • Figure 3: Figure \ref{['f3a']} and \ref{['f3b']} visualize the number of samples for two discrete action spaces in the training set. The offline RL dataset is split into 70% training, 15% for validation, and 15% test sets. Figure \ref{['f3c']} shows the test effective sample size (ESS) score of the policy functions selected by the validation ESS score. Figure \ref{['f3d']}--\ref{['f3f']} show the performance score evaluated from the test set. The AD-BCQ clearly improves upon two baselines, BCQ and factored-BCQ.
  • Figure 4: Discretized action spaces ranging from 5x5 to 14x14
  • Figure 5: The train and validation loss (Gaussian NLL) for 5 different action spaces ranging from 5x5 to 14x14. For each action space, we selected the state encoder model with the minimum validation loss and encoded the trajectory of all patients.
  • ...and 10 more figures

Theorems & Definitions (14)

  • definition 1: Projected Action Space MDP
  • proposition 1: Projected Q-function
  • proof
  • definition 2: Weighted Projected Q-functions
  • theorem 1: Convergence of MB-FPI
  • proof
  • theorem 2: Sample Complexity of MB-FPI
  • proof
  • proposition 1: Projected Q-function
  • proof
  • ...and 4 more