KL-DRO radius sweep¶

This is the smallest honest DRO experiment I could write: fix a nominal distribution, fix a per-scenario loss, and watch what happens to

\[ \sup_{q \,:\, D_{\mathrm{KL}}(q \,\|\, \mathrm{nominal}) \le \rho} \mathbb{E}_q[\mathrm{loss}] \]

as the radius \(\rho\) grows. Nothing is being optimized over a decision variable here — the decision is out of the picture entirely, so whatever the curve does is the ambiguity set's doing and nothing else's.

Why it is worth plotting¶

The formulation has two endpoints you can write down without solving anything, which makes it a genuinely useful sanity check rather than a pretty picture:

Radius	What the ball is	Worst case
\(\rho = 0\)	just \(\{\mathrm{nominal}\}\)	\(\sum_i \mathrm{nominal}_i \,\mathrm{loss}_i\)
\(\rho \to \infty\)	eventually contains the point mass on the worst scenario	\(\max_i \mathrm{loss}_i\)

The script asserts the first one to machine precision (Optora special-cases \(\rho = 0\) and returns the nominal expectation exactly, rather than sending the dual variable \(\eta \to \infty\) and hoping). The second is only approached: the dual optimum runs off to the boundary, so a finite-iteration solve gets close and stops. That asymmetry is expected, and seeing the curve flatten out just below the red line is the point.

In between, the curve should be nondecreasing — a bigger ball cannot contain a less adversarial distribution — and the script checks that too, on the sampled radii, rather than taking it on faith.

Reading the output¶

Three things to look at:

The printed comparison between \(\mathbb{E}_{\mathrm{nominal}}[\mathrm{loss}]\) and worst_case_expectation(radius=0.0); they should agree to about \(10^{-12}\).
The monotonically nondecreasing in radius: True line.
The figure, on a log-scaled radius axis: a flat start near the nominal expectation, a steep middle where robustness actually costs something, and a plateau at \(\max_i \mathrm{loss}_i\).

Everything is float64 here

The tolerances above are only defensible in double precision. In float32 the gradient norms in the dual solve plateau around \(10^{-6}\), which is enough to make a tight convergence check fail even though the iterate is sitting on the answer.

Source¶

examples/01_kl_dro_radius_sweep.py

"""KL-DRO worst-case expectation across a growing ambiguity radius.

Isolates the core DRO primitive: for a fixed nominal distribution and a
fixed per-scenario loss, how does the worst-case expected loss

    sup_{q: D_KL(q || nominal) <= radius} E_q[loss]

grow as the KL-ball radius grows? Two closed-form limits bound the curve,
and both are checked numerically below rather than only asserted:

- radius = 0: the ball collapses to `{nominal}`, so the worst case is
  exactly the nominal expectation `sum(nominal * loss)`.
- radius -> inf: the ball eventually contains a distribution putting all
  its mass on the single worst-case scenario, so the worst case saturates
  at `loss.max()` (approached in the limit, never exactly attained by a
  finite-iteration dual solve -- see `progress/decisions.md`).

Run:
    python examples/01_kl_dro_radius_sweep.py
"""

import matplotlib.pyplot as plt
import torch
from _plotting import save_figure

from optora.dro import KLAmbiguitySet
from optora.solvers import GradientDescent

NOMINAL = torch.tensor([0.4, 0.3, 0.2, 0.1], dtype=torch.float64)
LOSS = torch.tensor([0.5, 1.5, 3.0, 8.0], dtype=torch.float64)
DUAL_SOLVER = GradientDescent(step_size=0.1, max_iter=2000, tol=1e-9)


def main() -> None:
    nominal_expectation = torch.sum(NOMINAL * LOSS).item()
    exact_zero_radius = KLAmbiguitySet(
        nominal=NOMINAL, radius=0.0, dual_solver=DUAL_SOLVER
    )
    zero_radius_value = exact_zero_radius.worst_case_expectation(LOSS).item()
    print(f"E_nominal[loss]                    = {nominal_expectation:.6f}")
    print(f"worst_case_expectation(radius=0.0) = {zero_radius_value:.6f}")
    assert abs(nominal_expectation - zero_radius_value) < 1e-12

    radii = torch.logspace(-3, 1, steps=25, dtype=torch.float64)
    worst_case = [
        KLAmbiguitySet(nominal=NOMINAL, radius=radius.item(), dual_solver=DUAL_SOLVER)
        .worst_case_expectation(LOSS)
        .item()
        for radius in radii
    ]
    worst_possible = LOSS.max().item()
    print(f"worst case at radius={radii[0].item():.4f}  -> {worst_case[0]:.6f}")
    print(f"worst case at radius={radii[-1].item():.4f} -> {worst_case[-1]:.6f}")
    print(f"loss.max() (saturation limit)      = {worst_possible:.6f}")

    increments = torch.diff(torch.tensor(worst_case))
    is_nondecreasing = bool(torch.all(increments >= -1e-9))
    print(f"monotonically nondecreasing in radius: {is_nondecreasing}")

    fig, ax = plt.subplots(figsize=(6, 4))
    ax.plot(radii, worst_case, marker="o", label="worst-case E_q[loss]")
    ax.axhline(
        nominal_expectation, color="gray", linestyle="--", label="E_nominal[loss]"
    )
    ax.axhline(worst_possible, color="firebrick", linestyle=":", label="max_i loss_i")
    ax.set_xscale("log")
    ax.set_xlabel("KL ambiguity radius")
    ax.set_ylabel("worst-case expected loss")
    ax.set_title("KL-DRO worst-case expectation vs. ambiguity radius")
    ax.legend()
    fig.tight_layout()
    save_figure(fig, "01_kl_dro_radius_sweep")


if __name__ == "__main__":
    main()