Robust decisions across geometries¶

Once you accept that you should optimize against a worst case, the next question is which worst case. This example takes one small decision problem and solves it four times, changing only the shape of the ambiguity set:

\[ \min_x \; \sup_{q \in \mathcal{Q}(\mathrm{nominal}, \rho)} \mathbb{E}_q\!\left[(\mathrm{outcomes} - x)^2\right] \]

with \(\mathcal{Q}\) ranging over the KL, \(\chi^2\), total-variation, and Wasserstein balls. The outcomes are [1, 2, 3, 10] under a uniform nominal — three ordinary values and one outlier — because the whole point is to see how hard each geometry leans on that 10.

What the comparison shows¶

At \(\rho = 0\) every ball collapses to \(\{\mathrm{nominal}\}\), so all four formulations must return the same answer: the plain empirical-risk decision, which for squared error is just the nominal mean, \(4.0\). That shared starting point is what makes the rest of the plot comparable at all.

As \(\rho\) grows, each robust decision \(x^\star\) is dragged toward the outlier, but at different rates, and the ordering is a property of the divergence rather than of the solver:

Total variation moves in discrete-feeling jumps — it can relocate a block of mass wholesale, so it reacts fast and early.
KL is the most reluctant to fully abandon a scenario, since driving \(q_i \to 0\) costs it unboundedly.
\(\chi^2\) sits in between, and in the interior regime it matches the familiar \(\mathbb{E}[\cdot] + \sqrt{\rho \operatorname{Var}[\cdot]}\) variance-penalty form.
Wasserstein is the only one that cares about the geometry of the outcome space rather than just the labels: it is given the squared-distance cost matrix, so moving mass from \(1\) to \(10\) costs far more than moving it from \(3\) to \(10\).

That last distinction is the one I find easiest to forget. The \(\phi\)-divergence family is blind to how far apart the scenarios are; it only sees probabilities. Wasserstein is not.

A note on cost¶

Every point on every curve is a nested solve — an outer MinimaxSolver whose objective internally runs an ambiguity set's own dual solve on each iteration. Iteration counts multiply, so the budgets at the top of the script are cut down hard from the defaults, and each set gets its own step size.

The \(\chi^2\) step size is not arbitrary

Its conjugate \(\phi^\ast(s) = s + s^2/4\) grows faster than KL's logsumexp, and an aggressive step size makes the joint dual solve diverge to NaN silently rather than loudly. KL tolerates \(0.1\) here; \(\chi^2\) gets \(0.05\) for a reason.

Also worth knowing, since it looks like it should not work: the outer gradient flows correctly through worst_case_expectation even though the inner dual variable is solved for under detach. That is the envelope theorem — at the inner optimum the gradient contribution through the dual variable vanishes, so only the final re-evaluation needs to stay attached to the graph.

Source¶

examples/03_robust_decision_across_ambiguity_sets.py

"""Comparing the robust decision across all four ambiguity-set geometries.

Problem: choose a single scalar decision `x` to minimize the worst-case
expected squared error against four observed outcomes,

    min_x sup_{q in ambiguity_set(nominal, radius)} E_q[(outcomes - x) ** 2]

At `radius = 0` every ambiguity set collapses to `{nominal}`, so every
formulation recovers the same ordinary empirical-risk decision (the
nominal-weighted mean of `outcomes`). As `radius` grows, each formulation's
robust decision `x*` is pulled toward the outcome that would hurt most if
an adversary reallocated probability mass there; how quickly, and how far,
depends on the divergence geometry underneath the ambiguity set.

Run:
    python examples/03_robust_decision_across_ambiguity_sets.py
"""

from collections.abc import Callable

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

from optora.dro import (
    ChiSquareAmbiguitySet,
    KLAmbiguitySet,
    MinimaxProblem,
    MinimaxSolver,
    TotalVariationAmbiguitySet,
    WassersteinAmbiguitySet,
)
from optora.solvers import GradientDescent

OUTCOMES = torch.tensor([1.0, 2.0, 3.0, 10.0], dtype=torch.float64)
NOMINAL = torch.full_like(OUTCOMES, 1.0 / OUTCOMES.numel())
COST = (OUTCOMES.unsqueeze(0) - OUTCOMES.unsqueeze(1)) ** 2

# Nested outer/inner solver composition is expensive in eager PyTorch (see
# `progress/decisions.md`), so both budgets below are tuned down
# aggressively from each solver's own defaults; this is a demo of
# qualitative behavior, not a tight-tolerance convergence test.
# The outer solver uses `check_interval=1` because one of its steps costs a
# full inner dual solve, which dwarfs the synchronization a check costs.
OUTER_SOLVER = GradientDescent(step_size=0.01, max_iter=120, tol=1e-7, check_interval=1)
KL_DUAL_SOLVER = GradientDescent(step_size=0.1, max_iter=200, tol=1e-8)
CHI_SQUARE_DUAL_SOLVER = GradientDescent(step_size=0.05, max_iter=300, tol=1e-8)


def loss_fn(x: torch.Tensor) -> torch.Tensor:
    return (OUTCOMES - x) ** 2


def robust_decision(ambiguity_set: object, radius: float) -> float:
    problem = MinimaxProblem(
        ambiguity_set=ambiguity_set,  # type: ignore[arg-type]
        loss_fn=loss_fn,
        initial_point=torch.tensor(OUTCOMES.mean().item(), dtype=torch.float64),
    )
    result = MinimaxSolver(solver=OUTER_SOLVER).solve(problem)
    return result.point.item()


def main() -> None:
    erm_decision = OUTCOMES.mean().item()
    print(
        f"empirical-risk decision (every formulation at radius=0): {erm_decision:.4f}"
    )

    families: dict[str, Callable[[float], object]] = {
        "KL": lambda radius: KLAmbiguitySet(
            nominal=NOMINAL, radius=radius, dual_solver=KL_DUAL_SOLVER
        ),
        "chi-square": lambda radius: ChiSquareAmbiguitySet(
            nominal=NOMINAL, radius=radius, dual_solver=CHI_SQUARE_DUAL_SOLVER
        ),
        "total variation": lambda radius: TotalVariationAmbiguitySet(
            nominal=NOMINAL, radius=radius
        ),
        "Wasserstein": lambda radius: WassersteinAmbiguitySet(
            nominal=NOMINAL,
            cost=COST,
            radius=radius,
        ),
    }
    radii = [0.0, 0.02, 0.05, 0.1, 0.2, 0.4]

    fig, ax = plt.subplots(figsize=(7, 4))
    for name, build_ambiguity_set in families.items():
        decisions = [
            robust_decision(build_ambiguity_set(radius), radius) for radius in radii
        ]
        print(f"{name:>16}: " + " -> ".join(f"{value:.3f}" for value in decisions))
        ax.plot(radii, decisions, marker="o", label=name)

    ax.axhline(
        erm_decision, color="gray", linestyle="--", label="empirical risk (radius=0)"
    )
    ax.set_xlabel("ambiguity radius")
    ax.set_ylabel("robust decision x*")
    ax.set_title("robust decision vs. radius, across ambiguity-set geometries")
    ax.legend()
    fig.tight_layout()
    save_figure(fig, "03_robust_decision_across_ambiguity_sets")


if __name__ == "__main__":
    main()