Technical Summary: Minimax PAC Bounds for Learning in Exogenous Contextual MDPs

1. Problem Setting

This paper investigates the sample complexity of learning in tabular discounted Markov Decision Processes (MDPs) with exogenous i.i.d. contexts. The model is defined by a tuple $\mathcal{M} = (\mathcal{X}, \mathcal{Z}, \mathcal{A}, P, r, \mu, \gamma)$, where:

• $\mathcal{X}$ is a finite controlled state space.
• $\mathcal{A}$ is a finite action space.
• $\mathcal{Z}$ is a finite context space.
• $\gamma \in [0, 1)$ is the discount factor.
• $P: \mathcal{X} \times \mathcal{A} \times \mathcal{Z} \to \Delta(\mathcal{X})$ is the transition kernel.
• $r: \mathcal{X} \times \mathcal{A} \times \mathcal{Z} \to [0, 1]$ is the reward function.
• $\mu$ is the unknown distribution over contexts.

At each time step $t$, a context $Z_t$ is drawn independently from $\mu$ and revealed to the agent before action selection. The agent chooses $A_t$, receives reward $r(X_t, A_t, Z_t)$, and transitions to $X_{t+1} \sim P(\cdot | X_t, A_t, Z_t)$. Crucially, the context $Z_{t+1}$ is independent of the agent's actions and the current state, depending only on $\mu$.

The learning objectives are Policy Evaluation (PE), Best-Value Estimation (BVE), and Best-Policy Extraction (BPE). The sample complexity is measured by a pair $(n_{\text{learn}}, m_{\text{query}})$, representing the number of oracle calls made during an offline learning phase and the number of calls made during the decision-time (query) phase, respectively. The paper considers two regimes:

1. Known Dynamics: $P$ and $r$ are known; only $\mu$ is unknown.
2. Fully Unknown: Both $\mu$ and $P$ are unknown.

2. Methodology

2.1 Core Insight: Averaged Value Functions

The central methodological contribution is the reduction of the problem from the augmented space $\mathcal{X} \times \mathcal{Z}$ to the controlled state space $\mathcal{X}$. The authors define the averaged value function $\bar{V}(x) := \mathbb{E}_{Z \sim \mu}[V(x, Z)]$.
Because the next context is drawn independently from $\mu$, the Bellman optimality equation for the averaged value $\bar{V}^\star$ depends only on the transition kernel averaged over $\mu$:
$$\bar{V}^\star(x) = \mathbb{E}_{Z \sim \mu} \left[ \max_{a \in \mathcal{A}} \left( r(x, a, Z) + \gamma \sum_{x'} P(x' | x, a, Z) \bar{V}^\star(x') \right) \right]$$
This formulation allows the learner to estimate a vector in $\mathbb{R}^{|\mathcal{X}|}$ rather than $\mathbb{R}^{|\mathcal{X}||\mathcal{Z}|}$, effectively decoupling the complexity from the size of the context space $|\mathcal{Z}|$.

2.2 Algorithmic Framework: Variance-Reduced Halving

The paper employs a variance-reduced value iteration scheme (adapted from Sidford et al., 2019) to estimate the averaged value functions. The algorithm operates in two phases:

1. Offline Phase (Anchor Estimation): The algorithm computes a high-probability lower bound (anchor) for the Bellman operator applied to an initial value estimate. This requires a large batch of samples to control the variance of the empirical mean.
2. Online/Iterative Phase (Variance-Reduced Updates): Subsequent iterations estimate only the difference between the current Bellman backup and the anchor. By reusing the same samples for the difference term, the variance is significantly reduced, allowing for smaller sample sizes per iteration.

The algorithm uses a "halving" strategy, repeatedly calling a subroutine (HALFERR) that halves the error bound $u$ until the target accuracy $\epsilon$ is reached.

2.3 Handling Unknown Dynamics

In the fully unknown regime, the algorithm cannot compute the Bellman backup exactly even with a sampled context. Instead, for every sampled context $Z$, the algorithm draws a transition sample $X' \sim P(\cdot | x, \pi(x, Z), Z)$ to form an unbiased estimate of the backup. This introduces an additional variance term related to the transition kernel, which is rigorously bounded using a martingale decomposition that accounts for both context and transition randomness.

3. Key Results

3.1 Known Dynamics, Unknown Context Distribution ($P, r$ known, $\mu$ unknown)

When the transition kernel and rewards are known, the paper establishes that the sample complexity is independent of the context space size $|\mathcal{Z}|$.

• Complexity: The algorithm achieves the objectives with $(n_{\text{learn}}, m_{\text{query}}) = (\tilde{O}(\beta^3/\epsilon^2), 0)$, where $\beta = (1-\gamma)^{-1}$.
• Optimality: A minimax lower bound of $\Omega(\beta^3/\epsilon^2)$ is proven, showing the upper bound is tight up to logarithmic factors.
• Implication: The learner does not need to estimate values for every $(x, z)$ pair; estimating the averaged value $\bar{V}$ suffices. The query phase requires no additional samples ($m_{\text{query}}=0$) because the lift from $\bar{V}$ to $V(x, z)$ is a deterministic Bellman backup using the known $P$ and $r$.

3.2 Application to Perfect One-Step Lookahead

The framework is applied to the perfect one-step transition lookahead setting (where the agent observes all possible next states for all actions before acting).

• This setting is shown to be a special case of the exogenous contextual MDP where the context space is the set of all possible transition tensors.
• The result yields a sample complexity of $\tilde{O}(\beta^3/\epsilon^2)$ lookahead samples.
• In terms of individual tabular transitions, this corresponds to $\tilde{O}(\beta^3 |\mathcal{X}||\mathcal{A}|/\epsilon^2)$, improving the previous bound by Lu et al. (2025) by a factor of $\beta$.

3.3 Fully Unknown Regime ($P$ and $\mu$ unknown)

When both the dynamics and the context distribution are unknown, the paper focuses on Policy Evaluation (PE).

• Complexity: The algorithm achieves $(n_{\text{learn}}, m_{\text{query}}) = (\tilde{O}(|\mathcal{X}|\beta^3/\epsilon^2), \tilde{O}(\beta^2/\epsilon^2))$.
• Trade-off: The offline phase learns the averaged value $\bar{V}^\pi$ (costing $|\mathcal{X}|$ transitions per context sample). The query phase estimates the final one-step transition expectation at the realized $(x, z)$ using fresh samples.
• Lower Bounds: The paper proves matching lower bounds, demonstrating that:
  • $m_{\text{query}} = \Omega(\beta^2/\epsilon^2)$ is necessary for decision-time accuracy.
  • $n_{\text{learn}} + |\mathcal{X}| m_{\text{query}} = \Omega(|\mathcal{X}|\beta^3/\epsilon^2)$ is necessary for offline learning.
• Significance: Even in the fully unknown regime, the complexity remains independent of $|\mathcal{Z}|$. The unknown dynamics do not force estimation on the augmented space; they only introduce stochasticity in the decision-time "lift" from $\bar{V}$ to $V(x, z)$.

4. Significance and Claims

The paper claims to provide the first minimax-optimal PAC bounds for learning in exogenous contextual MDPs. Its primary contributions are:

1. Decoupling from Context Size: It demonstrates that the statistical price of large context spaces is avoidable. The sample complexity depends on the controlled state space $|\mathcal{X}|$ but not on the context space $|\mathcal{Z}|$.
2. Optimal Rates: The derived rates ($\tilde{O}(\beta^3/\epsilon^2)$ for known dynamics and $\tilde{O}(|\mathcal{X}|\beta^3/\epsilon^2)$ for fully unknown) are shown to be minimax optimal up to logarithmic factors.
3. Unified Framework: The framework subsumes and improves upon existing results for perfect one-step lookahead and provides a rigorous analysis of the offline/online sample trade-off.

The authors note that while the lower bound is tight at the purely online point ($n_{\text{learn}}=0$) and the reusable-oracle point (their proposed algorithm), the tightness of the bound for intermediate trade-offs remains an open question. Additionally, extending these results to Best-Policy Extraction (BPE) in the fully unknown regime is left as future work due to the difficulty of handling selection bias when estimating the maximum over actions from transition samples.