omnibias.struct.jax.plan
Differentiable dynamic programs for sequences and parses, structured distributions, and finite-temperature bounds.
import omnibias.struct.jax.planModule contract & mathematical conventions
Differentiable soft value iteration on an acyclic MDP (jax).
Bit-identical twin of omnibias.struct.torch.plan (float64; needs jax_enable_x64).
The soft-Bellman backup V_beta(s) = lse_beta_a [ r(s, a) + V_beta(next) ] is the
entropy-regularised value; it reuses the shared lse_beta and anneals to hard value
iteration as beta -> inf, with a certified log(N)/beta suboptimality. Differentiable
in the reward vector. Do not conflate the beta -> inf relaxation with the delta -> 0 tower.
Shared formulas do not imply universal floating-point identity. Verify dtype, shape, derivative-order support, and the stated numerical assumptions for the path you use.
soft_value_iteration
omnibias.struct.jax.soft_value_iterationSource ↗soft_value_iteration(rewards: Array, mdp: AcyclicMDP, beta: float=1.0) -> ArraySoft-Bellman value of mdp.start for a reward vector rewards (num_actions,).
V_beta(s) = lse_beta_a [ r(s, a) + V_beta(next(s, a)) ] (0 at terminals);
-> hard value iteration as beta -> inf. Differentiable in rewards (its
gradient is the soft action-visitation).
Parameters
| Name | Type & default | Contract |
|---|---|---|
rewardspositional-or-keyword | ArrayRequired | See the operation contract and implementation below. |
mdppositional-or-keyword | AcyclicMDPRequired | See the operation contract and implementation below. |
betapositional-or-keyword | floatDefault: 1.0 | See the operation contract and implementation below. |
ArrayExplicit validation & failure cases
- When
start_value is None, raisesValueError('start state value was not computed').
These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.
Executed regression · soft value iteration parity
Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.
def test_soft_value_iteration_parity() -> None:
mdp = _diamond_mdp()
rng = np.random.default_rng(1)
rewards = rng.standard_normal(len(mdp.actions))
for beta in (1.0, 8.0):
v_t = float(tplan.soft_value_iteration(torch.tensor(rewards), mdp, beta))
v_j = float(jplan.soft_value_iteration(jnp.asarray(rewards), mdp, beta))
assert abs(v_t - v_j) < 1e-9python -m pytest "packages/omnibias-struct/tests/test_plan_attn.py::test_soft_value_iteration_parity" -qRead the implementation
def soft_value_iteration(rewards: Array, mdp: AcyclicMDP, beta: float = 1.0) -> Array:
r"""Soft-Bellman value of ``mdp.start`` for a reward vector ``rewards`` ``(num_actions,)``.
``V_beta(s) = lse_beta_a [ r(s, a) + V_beta(next(s, a)) ]`` (``0`` at terminals);
``-> hard value iteration`` as ``beta -> inf``. Differentiable in ``rewards`` (its
gradient is the soft action-visitation).
"""
values: list[Array | None] = [None] * mdp.num_states
for s in range(mdp.num_states - 1, -1, -1):
acts = mdp.actions_of(s)
if not acts:
values[s] = jnp.zeros((), dtype=rewards.dtype)
continue
contribs = []
for i in acts:
nxt = values[mdp.actions[i][1]]
assert nxt is not None # next state has a higher index -> already computed
contribs.append(rewards[i] + nxt)
values[s] = logsumexp_beta(jnp.stack(contribs), beta, axis=-1)
start_value = values[mdp.start]
if start_value is None: # pragma: no cover - start is always in range
raise ValueError("start state value was not computed")
return start_valueOther public import paths
omnibias.struct.jax.soft_value_iterationomnibias.struct.jax.plan.soft_value_iteration
soft_value_iteration_batched
omnibias.struct.jax.soft_value_iteration_batchedSource ↗soft_value_iteration_batched(rewards: Array, mdp: AcyclicMDP, beta: float=1.0) -> ArrayBatched soft_value_iteration -> (B,) for rewards (B, num_actions) (jax.vmap).
Parameters
| Name | Type & default | Contract |
|---|---|---|
rewardspositional-or-keyword | ArrayRequired | See the operation contract and implementation below. |
mdppositional-or-keyword | AcyclicMDPRequired | See the operation contract and implementation below. |
betapositional-or-keyword | floatDefault: 1.0 | See the operation contract and implementation below. |
ArrayExecuted regression · families batched torch jax parity
Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.
def test_families_batched_torch_jax_parity() -> None:
pytest.importorskip("torch")
jax = pytest.importorskip("jax")
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
import torch
from omnibias.struct.jax import soft_eisner_batched as jax_ei
from omnibias.struct.jax import soft_value_iteration_batched as jax_vi
from omnibias.struct.torch import soft_eisner_batched as torch_ei
from omnibias.struct.torch import soft_value_iteration_batched as torch_vi
torch.set_default_dtype(torch.float64)
rng = np.random.default_rng(5)
arc = rng.standard_normal((4, 4, 4))
et = torch_ei(torch.tensor(arc), 2.0).numpy()
ej = np.asarray(jax_ei(jnp.asarray(arc), 2.0))
assert np.max(np.abs(et - ej)) < ATOL
mdp = _mdp()
rewards = rng.standard_normal((4, len(mdp.actions)))
vt = torch_vi(torch.tensor(rewards), mdp, 1.5).numpy()
vj = np.asarray(jax_vi(jnp.asarray(rewards), mdp, 1.5))
assert np.max(np.abs(vt - vj)) < ATOLpython -m pytest "packages/omnibias-struct/tests/test_batch_families.py::test_families_batched_torch_jax_parity" -qRead the implementation
def soft_value_iteration_batched(rewards: Array, mdp: AcyclicMDP, beta: float = 1.0) -> Array:
r"""Batched :func:`soft_value_iteration` -> ``(B,)`` for ``rewards`` ``(B, num_actions)`` (``jax.vmap``)."""
import jax
out: Array = jax.vmap(lambda r: soft_value_iteration(r, mdp, beta))(rewards)
return outOther public import paths
omnibias.struct.jax.soft_value_iteration_batchedomnibias.struct.jax.plan.soft_value_iteration_batched