choice.estimate_transition_matrix

Guide & worked example · Source

choice.estimate_transition_matrix(
    states,
    actions,
    next_states,
    n_states,
    n_choices,
    method='frequency',
    bandwidth=None,
)

Estimate state transition probabilities P(x’|x,a) from panel data.

Three approaches supported: 1. “frequency”: Nonparametric frequency estimator (implemented) 2. “kernel”: Kernel smoothing for continuous states (future) 3. “parametric”: Ordered probit/logit for ordinal transitions (future)

Args: states: (n_obs,) observed states x_t actions: (n_obs,) observed actions a_t next_states: (n_obs,) observed next states x_{t+1} n_states: Number of discrete states n_choices: Number of actions method: Estimation approach bandwidth: For kernel methods (not yet implemented)

Returns: P: (n_states, n_choices, n_states) estimated transition matrix where P[x, a, x’] ≈ Pr(x_{t+1}=x’ | x_t=x, a_t=a)

Example: >>> # Estimate transitions from bus maintenance data >>> P_hat = estimate_transition_matrix( … states=data.states, … actions=data.actions, … next_states=data.next_states, … n_states=90, … n_choices=2, … method=“frequency” … ) >>> # Verify probabilities sum to 1 >>> assert torch.allclose(P_hat.sum(dim=2), torch.ones(90, 2))

References: Rawat & Rust (2025), Section 3: RL methods work with estimated transitions