Uncertainty-Mass Penalty (UNM) for Robust Planning

• UNM is a principled metric that quantifies the expected cost increase when true system parameters deviate from estimates, integrating both distribution and mass of uncertainty.
• It employs statistical modeling techniques, including kernel density estimation and Sobol sensitivity analysis, to manage uncertainty in high-dimensional optimization tasks.
• The framework has been validated in query optimization and safe reinforcement learning, significantly reducing penalty costs and constraint violations with notable performance gains.

The Uncertainty-Mass Penalty (UNM) is a principled metric and algorithmic framework for managing risk under model uncertainty in optimization and control tasks. It quantifies the expected cost increase (“penalty”) incurred by a decision or policy when true system parameters deviate from nominal (estimated) values, accounting for both the distribution and mass of uncertainty. The UNM has been formalized and deployed in both query optimization (Xiu et al., 2024) and model-based safe reinforcement learning (Ma et al., 2021), each employing application-specific instantiations but sharing the core principle of risk-aware, expected-penalty minimization.

1. Formal Definition and Mathematical Formulation

Query Optimization Framework

Let $Q$ denote a query template, $\pi$ a candidate execution plan, $\hat s \in [0,1]^d$ the vector of estimated selectivities, and $\sigma \in [0,1]^d$ the true (but unknown) selectivities. Define $C(\pi, \sigma)$ as the realized cost of plan $\pi$ at selectivity $\sigma$, and $C^*(\sigma)$ as the cost of the optimal plan at $\sigma$.

A penalty function $\mathrm{pen}(\pi, \sigma)$ measures plan suboptimality relative to the optimum. For example, with tolerance $\tau$: $$\mathrm{pen}_\tau(\pi,\sigma) = \begin{cases} 0, & C(\pi,\sigma) \le (1+\tau) C^*(\sigma)\ C(\pi,\sigma) - C^*(\sigma), & \text{otherwise} \end{cases}$$ Since $\sigma$ is unknown, it is modeled as a random variable with density $f(\sigma\mid \hat s)$. The Uncertainty-Mass Penalty is the expected penalty: $$\mathbb{E}_{\sigma \sim f(\cdot \mid \hat s)}[\mathrm{pen}_\tau(\pi, \sigma)] = \int_{[0,1]^d}\mathrm{pen}_\tau(\pi,\sigma)f(\sigma \mid \hat s)d\sigma$$ The robust planning objective is to select $\pi$ minimizing this expected penalty (Xiu et al., 2024).

Safe Reinforcement Learning Formulation

Given a per-step cost $c(s,a)$ and a model epistemic uncertainty estimator $u_{\hat T}(s,a)$, the UNM-augmented cost is: $\tilde{c}(s,a) = c(s,a) + \kappa\, u_{\hat T}(s,a)$ where $\kappa \ge 0$ is a penalty coefficient. The robust policy optimization objective becomes: $$\max_\pi \; E_{\pi,\hat T}\left[\sum_{t=0}^\infty \gamma^t r(s_t,a_t)\right] \quad \text{subject to} \quad E_{\pi,\hat T}\left[\sum_{t=0}^\infty \gamma^t \left(c(s_t,a_t) + \kappa u_{\hat T}(s_t,a_t)\right)\right] \leq C$$ Equivalently, the dual occupancy-measure perspective expresses the cost and uncertainty-mass via stationary visitation measures (Ma et al., 2021).

2. Statistical Modeling of Uncertainty

Query Optimization: Selectivity Model

Workload-informed error models are constructed by profiling querylets (small subqueries) to gather (estimated, actual) selectivity pairs. Error profiles are bucketized (“low”, “high”) per selectivity dimension and modeled using kernel density estimators on $\log (\hat s_i / \sigma_i)$. Assuming inter-bucket independence, the joint conditional density factors as: $$f(\sigma | \hat s) \approx \prod_{i=1}^d g_i (\log (\hat s_i/\sigma_i) | \hat s_i )$$ This approach rigorously captures the empirical error mass and propagates it into the penalty computation (Xiu et al., 2024).

Reinforcement Learning: Epistemic Uncertainty Estimation

In model-based RL, $u_{\hat T}(s,a)$ serves as an upper bound on epistemic uncertainty. In tabular cases, IPM-based (e.g., total-variation) bounds derived from concentration inequalities yield: $$u_{\hat T}(s,a) = \sqrt{\frac{|S|}{8 n(s,a)} \ln \frac{4 |S| |A|}{\delta}}$$ For high-dimensional (ensemble) models, $u_{\hat T}(s,a)$ is the maximum predictive covariance's Frobenius norm across an ensemble or the variance of bootstrap predictions in latent (e.g., pixel-based) state representations (Ma et al., 2021).

3. Algorithmic Workflow and Robust Plan Selection

PARQO with UNM Objective

1. Reduce Problem Dimensionality: Sobol sensitivity analysis identifies the $k$ most “sensitive” selectivity dimensions (those whose variation most strongly affects penalty).
2. Sample Space Construction: Draw $S$ samples from $f(\sigma|\hat s)$ across these $k$ dimensions.
3. Plan Enumeration: For each sample $\sigma^{(j)}$, compute the true-optimal plan $\pi_j$ and cost $C^*(\sigma^{(j)})$, adding unique $\pi_j$ to the candidate set $P$.
4. Penalty Evaluation: For each $\pi \in P$, estimate sample-average expected penalty:

$$\widehat{\mathbb{E}}[\mathrm{pen}(\pi)] = \frac{1}{S} \sum_{j=1}^S \mathrm{pen}_\tau(\pi, \sigma^{(j)})$$

1. Plan Selection: Return the plan $\pi$ minimizing this expected penalty.

Extensions to parametric QO (PQO) reuse samples and candidate sets under bounded KL divergence between distributions $f(\cdot|\hat s)$ and $f(\cdot|\hat s')$ (Xiu et al., 2024).

CAP in Safe RL

1. Model Update: Fit transition model to current data.
2. Planning/Policy Optimization: Optimize policy with cost constraint that includes the UNM penalty.
3. Data Collection: Deploy policy in the real environment, augment buffer.
4. Penalty Adaptation: Update $\kappa$ with proportional-integral feedback based on realized cost, enforcing the constraint $C$.

This adaptive loop yields conservative policy updates by directly penalizing actions with high model uncertainty (Ma et al., 2021).

4. Sensitivity and Dimension Reduction Techniques

Sobol's variance decomposition is leveraged in PARQO to attribute total penalty variance to individual and joint selectivity dimensions. The first-order Sobol index for dimension $i$,

$$S_i = \frac{\mathrm{Var}\left[\mathbb{E}[Y | \sigma_i]\right]}{\mathrm{Var}[Y]}$$

quantifies the expected penalty contribution from variations in $\sigma_i$. The top $k$ dimensions by $S_i$ are selected for subsequent robust plan search, facilitating scalable optimization in high-dimensional settings (Xiu et al., 2024).

5. Theoretical Guarantees and Safety Properties

In model-based RL, the uncertainty-mass penalty ensures high-probability safe exploration:

• For finite state-action spaces and bounded costs, if a policy $\pi$ satisfies the conservative constraint under the model (with IPM-derived bounds), it also satisfies the cost constraint in the true environment with probability at least $1-\delta$.
• Under a union bound, zero-violation property extends to $K$ intermediate policies over training.

A simulation-lemma argument ties the difference in true vs. nominal cost to the model uncertainty term, regulating the safety via the penalization coefficient $\kappa$ and statistical confidence $\delta$ (Ma et al., 2021).

6. Empirical Evidence and Practical Impact

Empirical studies in both domains highlight the efficacy of UNM-based robustness:

Benchmark/Domain	Metric/Effect	UNM/Robust Baseline
JOB, DSB, STATS-CEB (QO)	Templates won/speedup	PARQO: 3.23×, 2.01×, 1.36×
PostgreSQL vs. PARQO (QO)	Max per-query gain	Up to 425×
IMDB time-sliced (QO)	Cross-slice plan speedup	∼3.8×
JOB PQO (QO)	End-to-end speedup, queries	2.4×, 33,000 queries
Gridworld (RL)	Training violations	CAP: Zero violations
HalfCheetah (RL)	Violations/steps vs. FOCOPS	1.7 vs. hundreds
Car Racing (RL)	Constraint violations	Dramatically reduced

In query optimization, robust plans delivered lower large penalties and outperformed baseline optimizers whenever selectivity errors induced large deviations. In RL, incorporating the UNM dramatically reduced constraint violations during training and enhanced sample efficiency (Xiu et al., 2024, Ma et al., 2021).

7. Related Methodologies and Extensions

• Distributionally Robust Optimization: The UNM metric operationalizes risk-averse planning analogous to expected regret minimization under plausible uncertainty distributions.
• Adaptive Safe RL: The dynamic penalty updating in CAP extends static robust approaches and ensures constraint adherence without excessive conservatism.
• Parametric QO (PQO): Efficient amortization of profiling cost via shared samples and plans furthers the scalability of robust query optimization (Xiu et al., 2024).

A plausible implication is that the UNM principle generalizes to any decision-making under epistemic model uncertainty, providing both a risk-sensitive objective and actionable algorithmic workflows. The distinction between model-driven (profiled) and data-driven (ensemble/statistical) uncertainty mass estimation enables flexible deployment across domains.