Skip to content

optora.dro.kl_dro

KL-divergence-constrained ambiguity set (KL-DRO).

KLAmbiguitySet

Bases: DualAmbiguitySet

KL-divergence-constrained ambiguity set for KL-DRO.

Bounds every candidate distribution q by \(D_{\mathrm{KL}}(q \,\|\, \mathrm{nominal}) \le \mathrm{radius}\). The worst-case expected loss over this set admits a convex dual (Hu and Hong 2013; Ben-Tal et al. 2013):

\[ \sup_{q:\, D_{\mathrm{KL}}(q \,\|\, \mathrm{nominal}) \,\le\, \mathrm{radius}} \mathbb{E}_q[\mathrm{loss}] = \inf_{\eta > 0} \; \eta \cdot \mathrm{radius} + \eta \log \mathbb{E}_{\mathrm{nominal}}\!\left[\exp\!\left(\frac{\mathrm{loss}}{\eta}\right)\right] \]

reducing the worst-case expectation to a one-dimensional convex minimization over the dual variable \(\eta\). dual_solver solves this minimization over \(\log(\eta)\) rather than \(\eta\) directly, so the unconstrained GradientDescent solver keeps \(\eta\) strictly positive throughout the iteration. This formulation needs no optimal-transport machinery, making it the simplest DRO formulation to build (see progress/architecture.md).

Attributes:

Name Type Description
nominal Tensor

Reference distribution the ambiguity set is centered on.
divergence

KLDivergence instance measuring distance from nominal.
radius float | Tensor

Nonnegative bound on the KL divergence of any distribution inside the ambiguity set from nominal, either a float or a tensor of radii evaluated as one batch.
eps

Small positive constant used to clamp nominal away from zero before taking the logarithm inside the dual objective.
dual_solver

Solver minimizing the dual objective over log(eta).
initial_dual_point Tensor

Value of log(eta) the first dual solve starts from; later solves warm-start from the previous optimum (see optora.core.dro_base.DualAmbiguitySet).
log_nominal Tensor

Elementwise logarithm of the clamped nominal, a constant of the dual objective cached once rather than recomputed on every call.

__init__(nominal, radius, eps=1e-12, dual_solver=None, initial_log_eta=0.0, validate=False)

Initialize the KL-DRO ambiguity set.

Parameters:

Name Type Description Default
nominal Tensor

Reference distribution the ambiguity set is centered on, a nonnegative tensor that sums to one along its last dimension.

 required
radius float | Tensor

Nonnegative bound on the KL divergence of any distribution inside the ambiguity set from nominal. A tensor radius is broadcast against the batch shape of worst_case_expectation's loss, evaluating a sweep of radii in one solve.

 required
eps float

Small positive constant used to clamp nominal away from zero before taking the logarithm inside the dual objective, and passed through to the underlying KLDivergence.

 1e-12
dual_solver Solver[MinimizationProblem, MinimizationResult] | None

Solver minimizing the dual objective over log(eta). Required when evaluating a positive-radius set.

 None
initial_log_eta float

Value of log(eta) the first dual solve starts from. Later calls warm-start from the previous solve's optimum unless reset_warm_start() is called.

 0.0
validate bool

Whether to check that nominal is nonnegative and sums to one. The check synchronizes with the device, so it is opt-in and off by default.

 False

Raises:

Type Description
ValueError

If radius is a negative float, if eps is not positive, or if validate is set and nominal is not a valid probability distribution.

worst_case_expectation(loss)

Compute the worst-case expected loss over the KL ambiguity set.

Parameters:

Name Type Description Default
loss Tensor

Per-scenario loss values of shape (..., n), one trailing entry per element of nominal's support. Leading dimensions are a batch of independent loss vectors, each solved with its own dual variable log(eta) in a single joint solve.

 required

Returns:

Type Description
Tensor

A tensor of shape (...) holding the worst-case expected loss:
Tensor

the exact sum(nominal * loss) when radius is zero (the
Tensor

ambiguity set then contains only nominal), otherwise the
Tensor

convex dual objective evaluated at the log(eta) found by
Tensor

dual_solver.

Raises:

Type Description
ValueError

If loss's trailing dimension does not match nominal's support size, or its batch shape does not broadcast against nominal and radius.
RuntimeError

If radius is positive and dual_solver is None.

contains(candidate)

Check whether a candidate distribution lies inside the ambiguity set.

The answer is returned as a boolean tensor on candidate's device rather than as a Python bool, so membership can be used as a mask or composed with further tensor work without forcing a device-to-host synchronization. Call bool(...) on the result only where a host-side branch is genuinely needed.

Parameters:

Name Type Description Default
candidate Tensor

Candidate distribution with the same shape as nominal.

 required

Returns:

Type Description
Tensor

A boolean tensor that is True where the divergence of
Tensor

candidate from nominal does not exceed radius.

reset_warm_start()

Discard the cached dual optimum so the next solve starts cold.

After this call the next worst_case_expectation starts from initial_dual_point again, as the first one did.