Worst-case distribution shift¶

The AmbiguitySet contract is deliberately narrow: worst_case_expectation hands back the scalar \(\mathbb{E}_{q^\star}[\mathrm{loss}]\) and nothing else. That is the right API — the worst-case distribution \(q^\star\) is an implementation detail of each formulation, and for the dual-solved sets it never exists explicitly in memory at all. But it does mean you never get to look at the adversary.

For total variation you can, because the solution is combinatorial rather than variational. This example reconstructs \(q^\star\) by hand and then plots where the probability mass goes.

The closed form being reconstructed¶

Inside a TV ball the adversary is solving a linear program with a mass budget, and the optimal strategy is embarrassingly simple: sort the scenarios by loss, then move mass from the cheapest scenarios into the single most expensive one until the budget

\[ \tfrac{1}{2}\sum_i |q_i - \mathrm{nominal}_i| \le \rho \]

runs out. Cheapest scenario first, drained completely before touching the next. Optora implements exactly this, vectorized, in TotalVariationAmbiguitySet — no iterative solver, no step size to tune, exact.

The worst_case_distribution helper in the script mirrors that logic in about ten lines. The subtle bit is cumsum(rest) - rest, which gives the mass already consumed before each scenario; the more obvious cat([zeros(1), cumsum(rest)[:-1]]) is off by one when there is only a single scenario to drain.

Trust, but cross-check¶

A reconstruction that merely looks right is worthless, so the script recomputes \(\mathbb{E}_{q^\star}[\mathrm{loss}]\) from its own \(q^\star\) and compares it against the library's worst_case_expectation across forty radii. It asserts agreement below \(10^{-9}\). If the two ever drift apart, either the reconstruction is wrong or the library is — and the assertion tells you to go find out which.

The radius sweep stops at the total movable mass \(1 - \mathrm{nominal}_{\arg\max \mathrm{loss}}\), because past that point everything is already piled onto the worst scenario and the plot has nothing left to say.

Reading the figure¶

The stackplot is the interesting artefact here. Each band is one scenario's probability mass as a function of radius, and you watch the low-loss bands collapse one at a time — never gradually and never simultaneously — while the worst-loss band swells to absorb them. That staircase is the LP's vertex-hopping made visible, and it is a much better intuition for "what does robustness actually assume about the world" than any scalar curve.

Source¶

examples/02_worst_case_distribution_shift.py

"""Tracking convergence across probability space: the TV-DRO worst-case
candidate distribution as the ambiguity radius grows.

`TotalVariationAmbiguitySet.worst_case_expectation` only returns the scalar
`E_q*[loss]` at the worst-case candidate distribution `q*`, not `q*`
itself -- the `AmbiguitySet` API is deliberately narrow (see
`optora/core/dro_base.py`). This example reconstructs `q*` directly from
the closed-form combinatorial solution the library documents internally
(sort scenarios by loss, then reallocate probability mass from the
cheapest scenarios to the single most expensive one, up to the radius
budget -- see `optora/dro/phi_dro.py` and `progress/decisions.md`), then
cross-checks `E_q*[loss]` against the library's own `worst_case_expectation`
to confirm the reconstruction is exact, not just plausible.

Run:
    python examples/02_worst_case_distribution_shift.py
"""

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

from optora.dro import TotalVariationAmbiguitySet

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)


def worst_case_distribution(
    nominal: torch.Tensor, loss: torch.Tensor, radius: float
) -> torch.Tensor:
    """Reconstruct the TV-DRO worst-case candidate distribution `q*`."""
    order = torch.argsort(loss)
    rest_idx, worst_idx = order[:-1], order[-1]
    rest_nominal = nominal[rest_idx]

    mass_available_before = torch.cumsum(rest_nominal, dim=0) - rest_nominal
    remaining_budget = torch.clamp(radius - mass_available_before, min=0.0)
    moved = torch.minimum(remaining_budget, rest_nominal)

    q = nominal.clone()
    q[rest_idx] = nominal[rest_idx] - moved
    q[worst_idx] = nominal[worst_idx] + moved.sum()
    return q


def main() -> None:
    max_movable_mass = (torch.sum(NOMINAL) - NOMINAL[torch.argmax(LOSS)]).item()
    radii = torch.linspace(0.0, max_movable_mass, steps=40, dtype=torch.float64)

    reconstructed = torch.stack(
        [worst_case_distribution(NOMINAL, LOSS, radius.item()) for radius in radii]
    )
    library_values = torch.stack(
        [
            TotalVariationAmbiguitySet(nominal=NOMINAL, radius=radius.item())
            .worst_case_expectation(LOSS)
            .detach()
            for radius in radii
        ]
    )
    reconstructed_values = reconstructed @ LOSS
    max_abs_diff = torch.max(torch.abs(reconstructed_values - library_values)).item()
    print(
        "max |reconstructed E_q*[loss] - library worst_case_expectation| = "
        f"{max_abs_diff:.2e}"
    )
    assert max_abs_diff < 1e-9

    fig, ax = plt.subplots(figsize=(7, 4))
    ax.stackplot(
        radii,
        *reconstructed.T,
        labels=[f"scenario {i} (loss={value:.1f})" for i, value in enumerate(LOSS)],
    )
    ax.set_xlabel("total-variation ambiguity radius")
    ax.set_ylabel("candidate probability mass q*_i(radius)")
    ax.set_title("worst-case distribution shift as the TV ball grows")
    ax.legend(loc="center left", bbox_to_anchor=(1.0, 0.5))
    fig.tight_layout()
    save_figure(fig, "02_worst_case_distribution_shift")


if __name__ == "__main__":
    main()