Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Distributionally Robust Control with End-to-End Statistically Guaranteed Metric Learning

Jingyi Wu, Chao Ning, Yang Shi

TL;DR

This work tackles distributionally robust control (DRC) under distributional uncertainty by identifying a mismatch between data-driven ambiguity-set design and downstream control performance. It introduces an end-to-end framework that learns an anisotropic Wasserstein ambiguity set, parameterized by a positive definite matrix $\Lambda$, within a bilevel optimization that couples inner DRC optimization with outer metric learning. A convex reformulation and a stochastic augmented Lagrangian algorithm enable tractable solution and provable guarantees, including finite-sample concentration for the anisotropic ball, continuity with respect to $\Lambda$, and non-asymptotic convergence to Clarke stationary points. Case studies on numerical and inventory-control tasks show that the proposed regionwise, learning-based approach yields superior closed-loop performance and better generalization across initial conditions compared to conventional Wasserstein DRC and pointwise end-to-end methods, highlighting the practical impact of task-aware ambiguity-set design.

Abstract

Wasserstein distributionally robust control (DRC) recently emerges as a principled paradigm for handling uncertainty in stochastic dynamical systems. However, it constructs data-driven ambiguity sets via uniform distribution shifts before sequentially incorporating them into downstream control synthesis. This segregation between ambiguity set construction and control objectives inherently introduces a structural misalignment, which undesirably leads to conservative control policies with sub-optimal performance. To address this limitation, we propose a novel end-to-end finite-horizon Wasserstein DRC framework that integrates the learning of anisotropic Wasserstein metrics with downstream control tasks in a closed-loop manner, thus enabling ambiguity sets to be systematically adjusted along performance-critical directions and yielding more effective control policies. This framework is formulated as a bilevel program: the inner level characterizes dynamical system evolution under DRC, while the outer level refines the anisotropic metric leveraging control-performance feedback across a range of initial conditions. To solve this program efficiently, we develop a stochastic augmented Lagrangian algorithm tailored to the bilevel structure. Theoretically, we prove that the learned ambiguity sets preserve statistical finite-sample guarantees under a novel radius adjustment mechanism, and we establish the well-posedness of the bilevel formulation by demonstrating its continuity with respect to the learnable metric. Furthermore, we show that the algorithm converges to stationary points of the outer level problem, which are statistically consistent with the optimal metric at a non-asymptotic convergence rate. Experiments on both numerical and inventory control tasks verify that the proposed framework achieves superior closed-loop performance and robustness compared against state-of-the-art methods.

Distributionally Robust Control with End-to-End Statistically Guaranteed Metric Learning

TL;DR

This work tackles distributionally robust control (DRC) under distributional uncertainty by identifying a mismatch between data-driven ambiguity-set design and downstream control performance. It introduces an end-to-end framework that learns an anisotropic Wasserstein ambiguity set, parameterized by a positive definite matrix , within a bilevel optimization that couples inner DRC optimization with outer metric learning. A convex reformulation and a stochastic augmented Lagrangian algorithm enable tractable solution and provable guarantees, including finite-sample concentration for the anisotropic ball, continuity with respect to , and non-asymptotic convergence to Clarke stationary points. Case studies on numerical and inventory-control tasks show that the proposed regionwise, learning-based approach yields superior closed-loop performance and better generalization across initial conditions compared to conventional Wasserstein DRC and pointwise end-to-end methods, highlighting the practical impact of task-aware ambiguity-set design.

Abstract

Wasserstein distributionally robust control (DRC) recently emerges as a principled paradigm for handling uncertainty in stochastic dynamical systems. However, it constructs data-driven ambiguity sets via uniform distribution shifts before sequentially incorporating them into downstream control synthesis. This segregation between ambiguity set construction and control objectives inherently introduces a structural misalignment, which undesirably leads to conservative control policies with sub-optimal performance. To address this limitation, we propose a novel end-to-end finite-horizon Wasserstein DRC framework that integrates the learning of anisotropic Wasserstein metrics with downstream control tasks in a closed-loop manner, thus enabling ambiguity sets to be systematically adjusted along performance-critical directions and yielding more effective control policies. This framework is formulated as a bilevel program: the inner level characterizes dynamical system evolution under DRC, while the outer level refines the anisotropic metric leveraging control-performance feedback across a range of initial conditions. To solve this program efficiently, we develop a stochastic augmented Lagrangian algorithm tailored to the bilevel structure. Theoretically, we prove that the learned ambiguity sets preserve statistical finite-sample guarantees under a novel radius adjustment mechanism, and we establish the well-posedness of the bilevel formulation by demonstrating its continuity with respect to the learnable metric. Furthermore, we show that the algorithm converges to stationary points of the outer level problem, which are statistically consistent with the optimal metric at a non-asymptotic convergence rate. Experiments on both numerical and inventory control tasks verify that the proposed framework achieves superior closed-loop performance and robustness compared against state-of-the-art methods.

Paper Structure

This paper contains 29 sections, 14 theorems, 113 equations, 2 figures, 1 table, 2 algorithms.

Table of Contents

  1. Introduction
  2. Problem Formulation
  3. System dynamics and constraints
  4. Conventional Wasserstein distributionally robust controller
  5. The Proposed Metric Learning Method
  6. Anisotropic Wasserstein ambiguity set with statistical finite-sample guarantee
  7. Anisotropic Wasserstein DRC and convex reformulation
  8. Bilevel learning framework
  9. Solution Algorithm
  10. Sample approximation problem
  11. Differentiate the controller
  12. Differentiate the augmented Lagrangian function
  13. Two-level algorithm for stochastic augmented Lagrangian framework
  14. Theoretical Analysis
  15. Continuity properties
  16. ...and 14 more sections

Key Result

Proposition 1

Let $\Lambda \in \mathbb{S}_{++}^d$ be a positive definite matrix. Then the anisotropic Wasserstein distance between $\mathbb{P}$ and $\mathbb{Q}$ satisfies where $\Lambda_{\#} \mathbb{P}$ denotes the pushforward measure of $\mathbb{P}$ under the linear map $z \mapsto \Lambda z$.

Figures (2)

  • Figure 1: Box plot of 10-step closed-loop costs for five random initial states across 100 disturbance realizations in numerical experiments
  • Figure 2: Box plot of 10-step closed-loop costs for five random initial states across 100 disturbance realizations in inventory control simulations

Theorems & Definitions (20)

  • Definition 1: Conditional Value-at-Risk
  • Definition 2: Anisotropic Wasserstein Distance
  • Proposition 1
  • Theorem 1: Statistical finite-sample guarantee
  • Theorem 2
  • Remark 1
  • Remark 2
  • Definition 3: Conservative Jacobian, differentiable-optimization-4
  • Proposition 2
  • Theorem 3
  • ...and 10 more