Technical Summary: Adaptive Estimation and Optimal Control in Offline Contextual MDPs without Stationarity

1. Problem Formulation

The paper addresses the challenge of estimating transition mechanisms and optimizing costs in Contextual Markov Decision Processes (MDPs) using offline datasets. Contextual MDPs generalize standard MDPs by incorporating exogenous information ("context") that modulates both state transitions and costs. While these models are critical for applications in healthcare, economics, and operations, existing statistical theory for offline learning in this domain relies heavily on restrictive assumptions:

• Stationarity/Ergodicity: Most methods assume the data generating process is stationary or ergodic, which is frequently violated in real-world logs (e.g., evolving patient health conditions).
• Parametric Modeling: Many approaches assume a specific parametric form for the transition density, leading to misspecification errors.
• Smoothness: Standard non-parametric methods often require strong smoothness assumptions (e.g., Hölder continuity) and fixed bandwidth selection, which is impractical when the state, control, and context spaces are high-dimensional or uncountably infinite.

The authors aim to develop a non-parametric framework that operates without stationarity or ergodicity assumptions, provides oracle risk bounds in an empirical Hellinger metric, and transfers these guarantees to optimal control selection via a plug-in scheme.

2. Methodology

The core of the proposed methodology is the application of T-estimation, a robust statistical technique rooted in model selection theory (Baraud & Birgé, 2009; Baraud, 2011).

2.1 The T-Estimator Construction

Let $s(x, a, g, y)$ denote the true transition density. The authors construct an estimator $\hat{s}$ by minimizing a penalized pairwise comparison functional over a rich, countable model class $\mathcal{S}$:
$$\hat{s} \in \arg\min_{f \in \mathcal{S}} \left( \sup_{f' \in \mathcal{S}} \left[ \alpha H^2(f, f') + T(f, f') \right] + \frac{L \Delta_{\mathcal{S}}(f)}{n} \right)$$
Where:

• $H^2$ is the empirical Hellinger loss.
• $T(f, f')$ is a comparison functional that evaluates candidates against observed trajectories. It includes a density-correcting term and compares how well $f$ and $f'$ fit the process $(X_i, a_i, G_i, X_{i+1})$.
• $\Delta_{\mathcal{S}}(f)$ is a data-independent complexity penalty (e.g., related to the dimension of the model or number of bins).
• $\alpha$ and $L$ are universal constants.

2.2 Technical Innovations

To derive guarantees without stationarity, the analysis relies on two key technical tools:

1. Martingale Bernstein Inequality: Used to control empirical deviations without assuming mixing or stationarity.
2. Polarization Identity: Used to relate risk in the squared-root density space to the Hellinger loss, enabling high-probability comparisons that convert to expected oracle inequalities.

2.3 From Estimation to Control

Once the transition density $\hat{s}$ is estimated, the paper addresses optimal control. Given a cost function $L(x, a, g, y)$, the empirical cost is defined as:
$$\hat{C}_n(a, g) = \int L(x, a, g, y) \hat{s}(x, a, g, y) \lambda_n^{(1)}(dx, dy)$$
The optimal control $\hat{a}^*(g)$ is selected by minimizing this estimated cost. The framework is policy-agnostic, meaning any procedure that optimizes this plug-in cost inherits the statistical rates of the estimator.

3. Key Contributions and Results

3.1 Non-Parametric Estimation with Oracle Inequalities

The paper establishes an oracle risk bound for the estimator $\hat{s}$ in the empirical Hellinger metric. Under mild assumptions (boundedness of a dominating density and no stationarity), the expected risk satisfies:
$$\mathbb{E}[H^2(s, \hat{s})] \lesssim \inf_{f \in \mathcal{S}} \left\{ H^2(s, f) + \frac{\Delta_{\mathcal{S}}(f) + \text{dim}(f)\log n}{n} \right\}$$
This result implies that the estimator performs as well as the best candidate in the model class $\mathcal{S}$, up to a penalty term, without requiring the data to be stationary.

3.2 Minimax-Rate Adaptivity

For standard smoothness classes (e.g., Hölder or Besov spaces), the procedure achieves the optimal non-parametric rate (up to poly-logarithmic factors):
$$\mathbb{E}[H^2(s, \hat{s})] \lesssim \left( \frac{\log n}{n} \right)^{\frac{\bar{\sigma}}{4+\bar{\sigma}}}$$
This demonstrates adaptivity: the estimator automatically adjusts to the unknown smoothness of the true density without prior knowledge of the bandwidth or smoothness parameter.

3.3 Finite-Sample Guarantees for Optimal Control

The paper provides non-asymptotic bounds for the cost gap between the control selected by the estimated surrogate and the true optimal control. The bound scales as:
$$O\left( \sqrt{\frac{\max(\Delta_{\mathcal{S}}(f), \log n)}{n}} \right)$$
Furthermore, the authors prove a minimax lower bound of $\Omega(1/\sqrt{n})$ for the optimal cost, showing that their plug-in scheme is rate-optimal. This links the difficulty of estimating the transition kernel directly to the difficulty of selecting the optimal control.

3.4 Robustness to Distribution Shift

The framework is extended to settings with distribution shift (covariate shift between training and deployment). The authors derive bounds showing that the estimation error under shift depends on the Hellinger distance between the true and shifted distributions, as well as the total variation distance between their invariant measures.

4. Significance and Claims

The paper claims to close a significant gap between the realities of offline contextual MDPs (irregular, non-stationary dynamics) and existing theory. Its primary significance lies in:

• Removing Restrictive Assumptions: It provides the first theoretical framework for offline contextual MDPs that does not rely on stationarity, ergodicity, or parametric modeling.
• Robustness: The T-estimator is inherently robust to model irregularity and non-stationarity, making it suitable for real-world logs where dynamics shift over time.
• Unified Framework: It unifies density estimation and optimal control, showing that a robust estimator directly translates to optimal control selection with rigorous finite-sample guarantees.
• Practical Applicability: The authors note that while T-estimators have computational costs (exponential in dimension for some classes), they offer a "practical, assumption-light route" to decision-making in complex domains like healthcare and economics where data is often messy and non-stationary.

The authors conclude that this method serves as an effective tool for inferential purposes in MDPs and optimal policy determination in offline settings, specifically addressing the limitations of non-stationarity and irregularity that plague current approaches.