optora.core.dro_base¶

Bases: Module, ABC

Set of distributions within a bounded divergence of a nominal distribution.

An ambiguity set pairs a Divergence with a radius: every distribution q inside the set satisfies divergence(q, nominal) <= radius. Modules in optora.dro subclass AmbiguitySet to implement the inner maximization of the DRO minimax problem for a specific divergence geometry (for example KL, a general phi-divergence, or Wasserstein).

Inherits from torch.nn.Module (rather than a plain ABC) so nominal is registered as a buffer and divergence as a submodule: a single .to(device)/.cuda() call then moves the nominal distribution and any tensor state the divergence holds (for example SinkhornDivergence's ground-cost matrix) together, and both surface through state_dict().

Attributes:

Bases: AmbiguitySet

Ambiguity set whose worst-case expectation is a low-dimensional dual solve.

Every divergence-based ambiguity set that reformulates its inner supremum as a convex dual minimization over a handful of dual variables (optora.dro.KLAmbiguitySet over \(\log(\eta)\), optora.dro.PhiAmbiguitySet over \((\log(\eta), \lambda)\)) runs the same machinery around a formulation-specific dual objective, so that machinery lives here: hold an injected dual_solver, start it from a dual point, and re-evaluate the dual objective at the returned optimum so the result stays differentiable with respect to loss by the envelope theorem.

Repeated calls are warm-started. The dual optimum moves only slightly between consecutive worst_case_expectation calls on a slowly changing loss — exactly what an optora.dro.MinimaxSolver outer iteration produces — so _solve_dual caches the detached optimum of each solve and starts the next solve from it instead of from initial_dual_point. That makes the second and later solves of such a sequence converge in far fewer inner iterations at the same optimum, since each dual is convex and its solution does not depend on where the iteration started. Call reset_warm_start() to discard the cache and go back to initial_dual_point, for example before evaluating an unrelated loss or after a diverged solve.

A batched loss of shape (..., n) is solved with one set of dual variables per batch element, since the duals decouple across batch elements. _solve_dual therefore hands dual_solver the sum of the per-element dual objectives: the sum's gradient with respect to any one element's dual variables is that element's own gradient, so a single joint solve reproduces independent per-element solves exactly. Convergence is judged on the joint gradient norm over the whole batch, so every element keeps iterating until the batch as a whole is stationary (see progress/decisions.md).

Attributes: