ENH: State-dependent shock distributions

Successor to item 3 of #6.

Current formulation

The Bellman operator is computed as $v(s) = \max_x, f(s,x) + \beta \sum_k w_k, v(g(s,x,e_k))$ with a single fixed node set shocks and weight vector weights used at every $(s, x)$: the distribution of $e$ is state- and action-independent.

What "state-dependent" can mean

Three flavors, with sharply different implementation costs:

1. Dependence through a smooth transformation of a fixed base shock — already expressible today. A shock $e \sim N(\mu(s), \sigma(s)^2)$ (wealth-dependent income risk, stochastic volatility, heteroskedastic AR processes) needs no new interface: write $e = \mu(s) + \sigma(s)u$ with a fixed standardized base shock $u$, put the fixed quadrature nodes/weights for $u$ in shocks/weights, and absorb the dependence into the transition, g(s, x, u) = h(s, x, mu(s) + sigma(s)*u). For location-scale families the quadrature accuracy equals the state-independent case; in principle any continuous conditional distribution is $e = F_s^{-1}(u)$ for a fixed uniform $u$. Proposed action: this is a documentation task — add a short docs section with the pattern.

2. State-dependent weights over a fixed support. Example: disaster risk — with probability $p(s)$ a crash shock, otherwise a normal draw, $p$ varying continuously with the state. The inverse-CDF trick technically applies but $F_s^{-1}$ is discontinuous in $u$: with $N$ nodes the mixing probability is effectively rounded to multiples of $1/N$ and is flat in $s$ almost everywhere, so this flavor genuinely needs an interface extension — weights as a function of the state (fixed shocks, $w_k(s)$). Mechanically mild: the sweeps and evaluate_policy! read $w_k(s_i)$ per node; the FOC inner solver is unaffected since $w$ is constant in $x$. Care points: the allocation-free sweep contract (the user function must not allocate per state — an in-place or tuple-returning convention), and LQA needs a defined behavior (e.g. evaluate the weights at the approximation point).

3. Action-dependent weights. Examples: search effort raising the offer arrival probability $\lambda(x)$; borrowing choices shifting default probabilities. Now $w_k(s, x)$, and the FOC gradient assembly is wrong as written — $\partial_x \sum_k w_k, \hat v(g)$ acquires $\sum_k \partial_x w_k \cdot \hat v(g)$ terms the current code does not compute. A first version would restrict the FOC path to $x$-independent weights and fall back to Brent.

General mechanism

The clean endpoint subsuming flavors 2–3 is a Bellman-kernel refactor: the kernel consumes a per-$(s,x)$ finite weighted support $((s'_k, w_k))_k$, with the current g + shocks + weights API as the constant-kernel special case. This is the same refactor identified in the #89 plan (the optional "generic explicit-finite models" step), so the two would share one implementation. Standing caveat inherited from #97: with state-dependent supports, verifying that all branch next-states stay inside the interpolation domain becomes harder, since the reachable set varies with $s$.

Relation to the POMDPs.jl interface (#89)

• POMDPs.jl covers all three flavors representationally, by construction: transition(m, s, a) builds a distribution object per call (SparseCat, ImplicitDistribution), so state/action dependence is the default rather than a feature; explicit solvers consume it uniformly via support/pdf/weighted_iterator.
• The planned as_mdp bridge inherits exactly this package's coverage: it constructs SparseCat(g.(s, x, shocks), weights), whose support varies with $(s, x)$ but whose weights are the constant quadrature vector — flavor 1 only. A native flavor-2 extension would flow through the bridge with a one-line change (SparseCat(g.(s, x, shocks), weights(s))).
• The reverse entrance (solving generic explicit-finite POMDPs models) is the kernel refactor above seen from the input side. The POMDPs interface cannot express smoothness in $x$ or weight-constancy, so the derivative-free inner solver is the default there — the price of that generality.

Suggested order: (i) document flavor 1 (cheap, closes the most common use cases); (ii) flavor 2 as an incremental enhancement behind the existing API; (iii) flavor 3 together with the kernel refactor, coordinated with the #89 sequencing.

🤖 Generated with Claude Code (Claude Fable 5)

Start by reading the current Bellman sweeps, evaluate_policy!, FOC path, and LQA behavior described in the issue, then compare the planned #89 kernel refactor and caveat from #97. The proposed sequence is to document flavor 1 first, add state-dependent weights next, and coordinate action-dependent weights with the generic finite-support kernel; done means the selected scope is implemented with allocation and interpolation behavior defined.