pyepo

PyEPO provides several families of end-to-end gradient surrogates for training predictors:

Solution-returning modules like DPO (which use randomized/perturbed losses) require an explicit key= argument when used inside a jax.jit decorated function. If you call a randomized loss inside jax.jit without providing a key, it will raise an error rather than silently freezing a single noise draw.

When using jax.jit, the key becomes a traced argument that must be split in each iteration of your training loop.

from pyepo.func.jax import DPO

dpo = DPO(optmodel, n_samples=10, sigma=0.5)

def loss_fn(p, k):
    # 'key=k' is required for jitted randomized losses
    we = dpo(predmodel.apply(p, xj), key=k) 
    return jnp.mean((we - wj) ** 2)

step = jax.jit(jax.grad(loss_fn))

key = jax.random.PRNGKey(0)
for epoch in range(10):
    key, subkey = jax.random.split(key)
    grads = step(params, subkey)
    updates, opt_state = optimizer.update(grads, opt_state)
    params = optax.apply_updates(params, updates)

The Solution Pool is a mechanism used in end-to-end predict-then-optimize training to speed up the process by approximating the feasible region. Instead of solving the full optimization problem (e.g., a linear or integer program) for every training instance, PyEPO stores previously computed optimal solutions in a pool $S$.

When the pool is active, PyEPO selects the best solution from the cached pool $S$ based on the predicted cost (minimizing the objective for minimization problems or maximizing for maximization) instead of calling a solver. This acts as an inner approximation of the feasible region.

Key components of the algorithm:

• Predictor weights ($\omega$): The parameters being trained.
• Predicted cost ($\hat{c}$): The output of the model $m(\omega, x)$.
• Solution Pool ($S$): A set of previously solved optimal solutions.
• solve_ratio ($p_{\text{solve}}$): The probability of choosing to solve the actual optimization problem versus selecting from the pool.
• Cost Transform ($t(\cdot)$): An optional transformation applied to predicted costs.

The pyepo.func.jax module provides JAX-compatible versions of the training methods (SPOPlus, DPO, PFY). These losses use jax.custom_vjp to handle the gradient through the optimization model.

Importing:

• PyTorch: from pyepo.func import SPOPlus, DPO, PFY
• JAX: from pyepo.func.jax import SPOPlus, DPO, PFY

Solver Backends:

• MPAX: Solved natively. The PDHG solve is JAX-traceable, allowing the entire training step to be accelerated with jax.jit.
• Non-MPAX (GurobiPy, COPT, Pyomo, OR-Tools): These are reached via jax.pure_callback, wrapping the CPU solver. While the training step can still be wrapped in jax.jit, the solver itself will run on the CPU.

import jax
import jax.numpy as jnp
import optax
from flax import linen as nn
import pyepo
from pyepo.data.dataset import optDataset
from pyepo.func.jax import SPOPlus

# Example setup for a 5x5 grid shortest path
grid = (5, 5)
optmodel = pyepo.model.shortestPathModel(grid)

# Synthetic data generation
x, c = pyepo.data.shortestpath.genData(
    num_data=1000, num_features=5, grid=grid, deg=4, noise_width=0.5, seed=135,
)

ds = optDataset(optmodel, x, c)
xj = jnp.asarray(x, jnp.float32)
cj, wj, zj = (jnp.asarray(a, jnp.float32) for a in (ds.costs, ds.sols, ds.objs))

# Predictor and Loss
predmodel = nn.Dense(optmodel.num_cost)
params = predmodel.init(jax.random.PRNGKey(0), xj[:1])
spo = SPOPlus(optmodel, reduction="mean")
optimizer = optax.adam(1e-2)
opt_state = optimizer.init(params)

# End-to-end training loop
for epoch in range(10):
    grads = jax.grad(lambda p: spo(predmodel.apply(p, xj), cj, wj, zj))(params)
    updates, opt_state = optimizer.update(grads, opt_state)
    params = optax.apply_updates(params, updates)