ENH: Add action_values to DiscreteDP; add util.index_dict

[oyamad]

• Add action_values to DiscreteDP, building on ENH: Add state_values to DiscreteDP #832 which added state_values
• Also add index_dict which builds a Numba typed dict (so usable within a jitted function) mapping each value in an array of unique values to its index
• Julia version: ENH: Add state_values/action_values to DiscreteDP QuantEcon.jl#402

Definition of the set of actions

I proposed to defer the decision in #832 (comment), but of the two approaches as described in QuantEcon/QuantEcon.jl#94 (comment) I am proposing Approach 2 here: Universal action set A with A(s) \subset A for each s; the union of all A(s)s may be a proper subset of A.

• For the product formulation: in this case, the n x m reward matrix R shares one action axis for all states, so each action a is the same action at every state: thus simply, self._num_actions = self.R.shape[1].
• For the SA pair formulation: in this case, the length of action_values if supplied defines the number of actions, so self._num_actions = values.shape[0] (where values = np.asarray(action_values)); otherwise, self._num_actions = self.a_indices.max() + 1.

Example

Old Aiyagari code:

import numpy as np
from numba import njit
import quantecon as qe

# Indexing: (a_i, z_i) pairs are mapped to s_i = a_i * z_size + z_i;
# to invert, a_i = s_i // z_size, z_i = s_i % z_size

@njit
def populate_R(R, a_size, z_size, a_vals, z_vals, r, w):
    n = a_size * z_size
    for s_i in range(n):
        a_i = s_i // z_size
        z_i = s_i % z_size
        a = a_vals[a_i]
        z = z_vals[z_i]
        for new_a_i in range(a_size):
            a_new = a_vals[new_a_i]
            c = w * z + (1 + r) * a - a_new
            if c > 0:
                R[s_i, new_a_i] = np.log(c)  # Utility

@njit
def populate_Q(Q, a_size, z_size, Pi):
    n = a_size * z_size
    for s_i in range(n):
        z_i = s_i % z_size
        for a_i in range(a_size):
            for next_z_i in range(z_size):
                Q[s_i, a_i, a_i * z_size + next_z_i] = Pi[z_i, next_z_i]

@njit
def asset_marginal(s_probs, a_size, z_size):
    a_probs = np.zeros(a_size)
    for a_i in range(a_size):
        for z_i in range(z_size):
            a_probs[a_i] += s_probs[a_i * z_size + z_i]
    return a_probs

r, w, beta = 0.03, 0.956, 0.96
Pi = np.array([[0.9, 0.1], [0.1, 0.9]])
z_vals = np.array([0.1, 1.0])
a_size, z_size = 200, 2
a_vals = np.linspace(1e-10, 18, a_size)
n = a_size * z_size

R = np.full((n, a_size), -np.inf)
populate_R(R, a_size, z_size, a_vals, z_vals, r, w)
Q = np.zeros((n, a_size, n))
populate_Q(Q, a_size, z_size, Pi)

ddp = qe.markov.DiscreteDP(R, Q, beta)
results = ddp.solve(method='policy_iteration')

# Decode the optimal policy by inverting the flattening formula
a_star = np.empty((z_size, a_size))
for s_i in range(n):
    a_i = s_i // z_size
    z_i = s_i % z_size
    a_star[z_i, a_i] = a_vals[results.sigma[s_i]]

# Aggregate capital, via the marginal distribution over assets
s_probs = results.mc.stationary_distributions[0]
asset_probs = asset_marginal(s_probs, a_size, z_size)
K = np.sum(asset_probs * a_vals)

This PR:

import numpy as np
from numba import njit
import quantecon as qe

@njit
def build_R_Q(s_vals, a_vals, z_vals, Pi, smap, zmap, r, w):
    n, m = s_vals.shape[0], a_vals.shape[0]
    R = np.full((n, m), -np.inf)
    Q = np.zeros((n, m, n))
    for s_i in range(n):
        a, z = s_vals[s_i, 0], s_vals[s_i, 1]
        z_i = zmap[z]
        for a_i in range(m):
            a_new = a_vals[a_i]
            c = w * z + (1 + r) * a - a_new
            if c > 0:
                R[s_i, a_i] = np.log(c)  # Utility
            for z_new in z_vals:
                Q[s_i, a_i, smap[(a_new, z_new)]] = Pi[z_i, zmap[z_new]]
    return R, Q

r, w, beta = 0.03, 0.956, 0.96
Pi = np.array([[0.9, 0.1], [0.1, 0.9]])
z_vals = np.array([0.1, 1.0])
a_vals = np.linspace(1e-10, 18, 200)

s_vals = np.array([(a, z) for a in a_vals for z in z_vals])  # states: (a, z)
smap = qe.index_dict(s_vals)   # numba.typed.Dict: (a', z') -> next-state index
zmap = qe.index_dict(z_vals)   # numba.typed.Dict: z -> row index of Pi

R, Q = build_R_Q(s_vals, a_vals, z_vals, Pi, smap, zmap, r, w)

ddp = qe.markov.DiscreteDP(R, Q, beta,
                           state_values=s_vals, action_values=a_vals)
results = ddp.solve(method='policy_iteration')

# Optimal policy decoded to asset values, paired with s_vals rows
a_star = results.sigma_values

# Aggregate capital, directly from the state values
s_probs = results.mc.stationary_distributions[0]
K = s_probs @ results.state_values[:, 0]

AI assisted with Claude Code Fable 5

[coveralls]

coverage: 90.78% (+0.2%) from 90.606% — ddp-state-action-values into main

[Copilot]

Pull request overview

Adds human-readable action values to DiscreteDP and introduces a Numba-compatible value-to-index utility.

Changes:

• Propagates state/action values through DP conversions and solve results.
• Adds sigma_values for decoded optimal policies.
• Adds and documents index_dict, with comprehensive tests.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
quantecon/markov/ddp.py	Implements action values and result propagation.
quantecon/markov/core.py	Clarifies state-value documentation.
quantecon/markov/tests/test_ddp.py	Tests state/action value behavior.
quantecon/util/indexing.py	Implements index_dict.
quantecon/util/tests/test_indexing.py	Tests indexing behavior and validation.
quantecon/util/__init__.py	Exports index_dict.
quantecon/__init__.py	Exposes index_dict publicly.
docs/source/util/indexing.rst	Adds utility API documentation.
docs/source/util.rst	Includes the indexing documentation page.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[mmcky]

thanks @oyamad