omnibias.struct.torch.plan
Differentiable dynamic programs for sequences and parses, structured distributions, and finite-temperature bounds.
import omnibias.struct.torch.planModule contract & mathematical conventions
Differentiable soft value iteration on an acyclic MDP (torch).
Bit-identical twin of omnibias.struct.jax.plan (float64). 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 (the
temperature axis), 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.torch.soft_value_iterationSource ↗soft_value_iteration(rewards: Tensor, mdp: AcyclicMDP, beta: float=1.0) -> TensorSoft-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 | TensorRequired | 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. |
TensorExplicit 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" -qRegression example · soft value iteration gradient is action visitation
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_gradient_is_action_visitation() -> None:
mdp = _diamond_mdp()
rewards = torch.tensor([0.5, -0.2, 0.3, 0.1, 0.7], requires_grad=True)
tplan.soft_value_iteration(rewards, mdp, 3.0).backward()
grad = rewards.grad
assert grad is not None
# Every trajectory takes the single terminal action 3->4, so its visitation is 1.
assert abs(float(grad[4]) - 1.0) < 1e-9
# Sibling actions out of the start share the whole visitation mass.
assert abs(float(grad[0] + grad[1]) - 1.0) < 1e-9python -m pytest packages/omnibias-struct/tests/test_plan_attn.py -k test_soft_value_iteration_gradient_is_action_visitation -qRead the implementation
def soft_value_iteration(rewards: Tensor, mdp: AcyclicMDP, beta: float = 1.0) -> Tensor:
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[Tensor | 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] = torch.zeros((), dtype=rewards.dtype, device=rewards.device)
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(torch.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.torch.soft_value_iterationomnibias.struct.torch.plan.soft_value_iteration
soft_value_iteration_batched
omnibias.struct.torch.soft_value_iteration_batchedSource ↗soft_value_iteration_batched(rewards: Tensor, mdp: AcyclicMDP, beta: float=1.0) -> TensorBatched soft_value_iteration -> (B,) for rewards (B, num_actions).
Maps the soft-Bellman backup over the leading batch axis with torch.func.vmap (shared
mdp structure) -- bit-identical to looping the per-example layer.
Parameters
| Name | Type & default | Contract |
|---|---|---|
rewardspositional-or-keyword | TensorRequired | 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. |
TensorExecuted regression · soft mas and value iteration batched torch
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_mas_and_value_iteration_batched_torch() -> None:
torch = pytest.importorskip("torch")
torch.set_default_dtype(torch.float64)
from omnibias.struct.torch import (
soft_mas,
soft_mas_batched,
soft_value_iteration,
soft_value_iteration_batched,
)
rng = np.random.default_rng(1)
score = torch.tensor(rng.standard_normal((4, 3, 4)))
b_mas = soft_mas_batched(score, 2.0)
loop_mas = torch.stack([soft_mas(score[i], 2.0) for i in range(4)])
assert b_mas.shape == (4,)
assert torch.max(torch.abs(b_mas - loop_mas)).item() < ATOL
mdp = _mdp()
rewards = torch.tensor(rng.standard_normal((4, len(mdp.actions))))
b_vi = soft_value_iteration_batched(rewards, mdp, 1.5)
loop_vi = torch.stack([soft_value_iteration(rewards[i], mdp, 1.5) for i in range(4)])
assert torch.max(torch.abs(b_vi - loop_vi)).item() < ATOLpython -m pytest "packages/omnibias-struct/tests/test_batch_families.py::test_soft_mas_and_value_iteration_batched_torch" -qRegression example · families batched torch jax parity
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 -k test_families_batched_torch_jax_parity -qRead the implementation
def soft_value_iteration_batched(rewards: Tensor, mdp: AcyclicMDP, beta: float = 1.0) -> Tensor:
r"""Batched :func:`soft_value_iteration` -> ``(B,)`` for ``rewards`` ``(B, num_actions)``.
Maps the soft-Bellman backup over the leading batch axis with ``torch.func.vmap`` (shared
``mdp`` structure) -- bit-identical to looping the per-example layer.
"""
from torch.func import vmap
def fwd(r: Tensor) -> Tensor:
return soft_value_iteration(r, mdp, beta)
out: Tensor = vmap(fwd)(rewards)
return outOther public import paths
omnibias.struct.torch.soft_value_iteration_batchedomnibias.struct.torch.plan.soft_value_iteration_batched