| """Sparse dynamic programs for endpoint-normalized path laws. |
| |
| No root-to-terminal path enumeration is needed. Runtime is linear in the |
| materialized graph size. Full catalog expansion remains a separate cost. |
| """ |
|
|
| from dataclasses import dataclass |
| import numpy as np |
| from scipy.special import logsumexp |
| from .graph import Graph |
|
|
|
|
| @dataclass |
| class Solution: |
| """Exact node log partitions, edge probabilities, and endpoint weights.""" |
|
|
| log_prefix: np.ndarray |
| backward: np.ndarray |
| forward: np.ndarray |
| target: dict[str, float] |
| log_z: float |
| temperature: float |
|
|
|
|
| def normalize_rewards(graph, log_rewards): |
| """Normalize finite log weights over canonical outcomes and return log Z.""" |
| keys = set(graph.terminals) |
| if set(log_rewards) != keys: |
| raise ValueError("Reward keys must equal the terminal outcome set") |
| vals = np.array([log_rewards[y] for y in keys], dtype=float) |
| if not np.all(np.isfinite(vals)): |
| raise ValueError("Log rewards must be finite") |
| z = float(logsumexp(vals)) |
| return {y: float(np.exp(log_rewards[y] - z)) for y in graph.terminals}, z |
|
|
|
|
| def prefix_values(graph, temperature=1.0): |
| """Compute one log partition per node in topological order.""" |
| if not np.isfinite(temperature) or temperature <= 0: |
| raise ValueError("Temperature must be positive") |
| a = np.full(len(graph.nodes), -np.inf) |
| a[graph.root_index] = 0.0 |
| for v in graph.order: |
| if v == graph.root_index: |
| continue |
| a[v] = logsumexp( |
| [ |
| a[graph.node_index[graph.edges[i].source]] |
| - graph.edges[i].cost / temperature |
| for i in graph.incoming[v] |
| ] |
| ) |
| return a |
|
|
|
|
| def backward_policy(graph, values, temperature): |
| """Local normalization preserves a complete conditional route law.""" |
| q = np.zeros(len(graph.edges)) |
| for v in graph.order: |
| ids = graph.incoming[v] |
| if not ids: |
| continue |
| logits = np.array( |
| [ |
| values[graph.node_index[graph.edges[i].source]] |
| - graph.edges[i].cost / temperature |
| for i in ids |
| ] |
| ) |
| q[ids] = np.exp(logits - logsumexp(logits)) |
| return q |
|
|
|
|
| def forward_from_backward(graph, q, target): |
| """Propagate endpoint probability backward, then reverse each edge flow.""" |
| logmass = np.full(len(graph.nodes), -np.inf) |
| for y, v in graph.terminals.items(): |
| logmass[v] = np.log(target[y]) |
| for v in reversed(graph.order): |
| for ei in graph.incoming[v]: |
| u = graph.node_index[graph.edges[ei].source] |
| if q[ei] > 0: |
| logmass[u] = np.logaddexp(logmass[u], logmass[v] + np.log(q[ei])) |
| p = np.zeros(len(graph.edges)) |
| for u in graph.order: |
| ids = graph.outgoing[u] |
| if not ids: |
| continue |
| lp = np.array( |
| [ |
| np.log(q[i]) + logmass[graph.node_index[graph.edges[i].target]] |
| for i in ids |
| ] |
| ) |
| p[ids] = np.exp(lp - logsumexp(lp)) |
| return p |
|
|
|
|
| def solve(graph: Graph, log_rewards, temperature=1.0): |
| """Construct the exact joint law on a finite stored executable graph.""" |
| target, z = normalize_rewards(graph, log_rewards) |
| a = prefix_values(graph, temperature) |
| q = backward_policy(graph, a, temperature) |
| return Solution( |
| a, q, forward_from_backward(graph, q, target), target, z, temperature |
| ) |
|
|
|
|
| def endpoint_distribution(graph, forward): |
| """Propagate a normalized edge policy and return exact terminal masses.""" |
| mass = np.zeros(len(graph.nodes)) |
| mass[graph.root_index] = 1.0 |
| for u in graph.order: |
| ids = graph.outgoing[u] |
| if ids and not np.isclose(np.sum(forward[ids]), 1, atol=1e-6): |
| raise ValueError("Invalid forward normalization") |
| for ei in ids: |
| mass[graph.node_index[graph.edges[ei].target]] += mass[u] * forward[ei] |
| return {y: float(mass[v]) for y, v in graph.terminals.items()} |
|
|
|
|
| def sample(graph, forward, n=100, seed=0): |
| """Draw n complete paths, including edge IDs, action records, and costs.""" |
| rng = np.random.default_rng(seed) |
| result = [] |
| for _ in range(n): |
| u = graph.root_index |
| path = [] |
| cost = 0.0 |
| lp = 0.0 |
| while graph.outgoing[u]: |
| ids = graph.outgoing[u] |
| probs = np.asarray(forward[ids]) |
| probs = probs / probs.sum() |
| j = int(rng.choice(len(ids), p=probs)) |
| ei = ids[j] |
| e = graph.edges[ei] |
| path.append(ei) |
| cost += e.cost |
| lp += np.log(probs[j]) |
| u = graph.node_index[e.target] |
| result.append( |
| { |
| "outcome": graph.nodes[u].outcome, |
| "edge_ids": [graph.edges[i].id for i in path], |
| "edge_indices": path, |
| "cost": cost, |
| "log_probability": float(lp), |
| "actions": [graph.edges[i].action for i in path], |
| } |
| ) |
| return result |
|
|
|
|
| def expected_cost(graph, forward): |
| """Compute mean additive execution cost by propagating node masses.""" |
| mass = np.zeros(len(graph.nodes)) |
| mass[graph.root_index] = 1.0 |
| cost = 0.0 |
| for u in graph.order: |
| for ei in graph.outgoing[u]: |
| e = graph.edges[ei] |
| f = mass[u] * forward[ei] |
| mass[graph.node_index[e.target]] += f |
| cost += f * e.cost |
| return float(cost) |
|
|
|
|
| def uniform_policy(graph): |
| """Return equal probabilities across each node's outgoing actions.""" |
| p = np.zeros(len(graph.edges)) |
| for ids in graph.outgoing: |
| if ids: |
| p[ids] = 1 / len(ids) |
| return p |
|
|
|
|
| def tilted_reference(graph, log_rewards, temperature=None): |
| """Reference endpoint tilting; optional path-cost penalty is explicit.""" |
| p = uniform_policy(graph) |
| h = np.full(len(graph.nodes), -np.inf) |
| for y, v in graph.terminals.items(): |
| h[v] = log_rewards[y] |
| for u in reversed(graph.order): |
| ids = graph.outgoing[u] |
| if ids: |
| logits = np.array( |
| [ |
| np.log(p[i]) |
| + h[graph.node_index[graph.edges[i].target]] |
| - (graph.edges[i].cost / temperature if temperature else 0) |
| for i in ids |
| ] |
| ) |
| h[u] = logsumexp(logits) |
| p[ids] = np.exp(logits - h[u]) |
| return p |
|
|
