The Transition System

pylcm models involve two independent layers of transitions that compose during backward induction:

1. State transitions — how a state variable arrived at its current value (attached to grids)

2. Regime transitions — which regime the agent enters next period (attached to the regime itself)

The two layers have opposite orientations:

• Regime transitions are forward-looking. The transition field on a Regime answers “where does the agent go next?” It lives on the source regime and points toward the future.

• State transitions are backward-looking. The transition parameter on a grid answers “how did this state variable reach its current value?” It lives on the grid that receives the value. In multi-regime models, per-boundary mappings are placed on the target regime’s grid.

This notebook explains how each layer works, how per-boundary transitions resolve cross-regime state mismatches, and how everything composes in the Bellman equation.

State Transition Mechanics¶

State transitions are attached directly to grid objects via the transition parameter. There are four cases:

Deterministic state transitions¶

A state transition function defines how the current state value was determined from last period’s states, actions, and parameters. The function’s argument names are resolved from the regime’s namespace.

Fixed states and identity transitions¶

When transition=None, pylcm auto-generates an _IdentityTransition internally. This shows up in regime.get_all_functions() under the key "next_<state_name>".

Both fixed states produce _IdentityTransition objects. These are marked with _is_auto_identity = True so that validation can distinguish them from user-provided transitions.

Stochastic state transitions (DiscreteMarkovGrid)¶

For DiscreteMarkovGrid, the transition function returns a probability array over the categories. During the solve step, pylcm computes a probability-weighted expectation over next-period states:

Shock grids¶

Shock grids (from lcm.shocks.iid and lcm.shocks.ar1) have intrinsic transitions computed from the distribution. For IID shocks, the transition probabilities are the same regardless of the current value. For AR(1) shocks, probabilities depend on the current state.

Shock grids do not accept a transition parameter — their transitions are built-in.

Per-Boundary State Transitions¶

When a discrete state has different categories across regimes, a simple callable transition is not enough — you need to map from one category set to another at the regime boundary.

The solution: a mapping transition keyed by (source_regime, target_regime) pairs, placed on the target regime’s grid.

Example: different health categories¶

Suppose working life has three health states (disabled, bad, good) but retirement only has two (bad, good). The transition from working to retired needs an explicit mapping.

The mapping is placed on the target regime’s grid:

Resolution priority¶

When resolving which transition function to use at a regime boundary (source, target), pylcm checks (in order):

1. Target grid mapping for (source, target)

2. Source grid mapping for (source, target)

3. Source grid’s callable transition

4. Target grid’s callable transition

5. Auto-generated identity (if categories match)

If the categories differ across regimes and no explicit mapping is found, ModelInitializationError is raised.

Parameterized per-boundary transitions¶

Per-boundary mapping functions can take parameters beyond the state variable itself. A common use case is a continuous state whose transition law differs across regime boundaries — for example, wealth that grows at a rate specific to the target regime.

When pylcm resolves a per-boundary transition from the target grid’s mapping (priority 1 above), any parameters in that function are looked up in the target regime’s parameter template. This means the user specifies the parameter value under the target regime in the params dict, and pylcm automatically routes it to the transition function at the boundary.

The rule is simple: whoever owns the mapping owns the parameters. Since per-boundary mappings live on the target regime’s grid, their parameters come from the target regime.

Example: regime-specific growth rate¶

Consider a two-regime model (phase 1 → phase 2) where wealth grows at a rate that is specific to phase 2. The transition function on phase 2’s wealth grid takes a growth_rate parameter:

Because the mapping {("phase1", "phase2"): next_wealth_at_boundary} lives on phase 2’s grid, the growth_rate parameter appears in phase 2’s parameter template. The user supplies it under "phase2" in the params dict:

params = {
    "phase1": {...},
    "phase2": {
        "next_wealth": {"growth_rate": 0.05},
        ...
    },
}

Internally, pylcm detects that the transition was resolved from the target grid’s mapping and renames the parameter to a cross-boundary qualified name (e.g., phase2__next_wealth__growth_rate). At solve and simulation time, the value is looked up from internal_params["phase2"] — not from "phase1" — even though the transition is evaluated as part of phase 1’s backward induction step.

Internals: How Cross-Boundary Parameters Are Processed¶

The sections above describe the user-facing rule: whoever owns the mapping owns the parameters. This section explains the internal machinery that makes it work.

The problem: flat function signatures¶

All compiled Q functions have flat parameter signatures. The dags library composes utility, constraints, and transitions into a single function whose parameters are flat qualified names (qnames) like utility__risk_aversion or next_wealth__interest_rate. These are passed as **kwargs at the call site:

# In solve_brute.py:
V_arr = max_Q_over_a(
    **state_action_space.states,
    **state_action_space.actions,
    next_V_arr=next_V_arr,
    **internal_params[regime_name],  # flat kwargs
)

When a per-boundary transition is target-originated (e.g., phase2’s grid owns the (phase1, phase2) mapping), its transition function is compiled into phase1’s Q function — but the parameter values live in internal_params["phase2"]. The flat signature means we can’t just pass a nested dict; the value must appear as a named kwarg in phase1’s params.

Step 1: Transition resolution and origin tracking¶

_resolve_state_transition (in regime_processing.py) resolves which transition function to use at each regime boundary. When the winning function comes from the target grid’s mapping dict (priority 1), it is recorded in a target_originated_transitions set for that source regime.

For the phase1 → phase2 example, phase1’s resolved transitions include the function next_wealth_at_boundary under the flat name "phase2__next_wealth". Because it came from phase2’s grid, it is marked as target-originated.

Step 2: Parameter renaming (_rename_target_originated_transition)¶

Normal (source-originated) transitions rename growth_rate → next_wealth__growth_rate — a standard qname within the source regime’s namespace.

Target-originated transitions use a cross-boundary qname that embeds the target regime name. _rename_target_originated_transition does the following:

1. Extracts the target regime name from the flat function name ("phase2__next_wealth" → "phase2").

2. Looks up the target regime’s parameter template to find which parameters the function uses (e.g., growth_rate).

3. Renames each parameter to a cross-boundary qname: growth_rate → phase2__next_wealth__growth_rate.

4. Records the mapping in a cross_boundary_params dict:

   cross_boundary_params["phase2__next_wealth__growth_rate"] = (
       "phase2",                      # target regime
       "next_wealth__growth_rate",    # qname in target's namespace
   )

This renaming avoids collisions with phase1’s own parameters (phase1 might have its own next_wealth__growth_rate for within-regime transitions).

Step 3: flat_param_names and function compilation¶

After all transitions are renamed, process_regimes computes flat_param_names — the model-level set of all regime-prefixed flat parameter names:

flat_param_names = frozenset(flatten_to_qnames(params_template))

This flattens the nested params_template (keyed by regime → function → arg) into qualified names like "phase1__utility__risk_aversion" or "phase2__next_wealth__growth_rate". Cross-boundary qnames are included because _rename_target_originated_transition already inserted them into the compiled function signatures. The frozenset is stored on InternalRegime.flat_param_names and passed to the function compilation pipeline (build_Q_and_F_functions, build_next_state_simulation_functions).

The compilation pipeline uses flat_param_names for one critical purpose: determining which arguments to vmap over. In _get_vmap_params:

non_vmap = {"period", "age"} | flat_param_names
return tuple(arg for arg in all_args if arg not in non_vmap)

Parameters (including cross-boundary ones) are scalar — they must not be vmapped. Everything else (states, actions) gets vmapped.

Step 4: Runtime resolution (merge_cross_boundary_params)¶

At the top of solve() and simulate(), before any backward induction or forward simulation begins, merge_cross_boundary_params is called once:

internal_params = merge_cross_boundary_params(internal_params, internal_regimes)

This copies values from the target regime’s params into the source regime’s params dict under the cross-boundary qname:

# For phase1's cross_boundary_params:
#   "phase2__next_wealth__growth_rate" → ("phase2", "next_wealth__growth_rate")
#
# merge_cross_boundary_params does:
internal_params["phase1"]["phase2__next_wealth__growth_rate"] = (
    internal_params["phase2"]["next_wealth__growth_rate"]
)

After merging, every call site just uses **internal_params[regime_name] — no per-call-site cross-boundary logic needed.

Summary of the internal pipeline¶

User defines:       phase2 grid transition={("phase1","phase2"): f(wealth, growth_rate)}
                                                        │
                                                        ▼
Resolution:         _resolve_state_transition marks it as target-originated
                                                        │
                                                        ▼
Renaming:           _rename_target_originated_transition renames
                    growth_rate → phase2__next_wealth__growth_rate
                    and records cross_boundary_params mapping
                                                        │
                                                        ▼
Compilation:        flat_param_names = own_params ∪ cross_boundary_params.keys()
                    Q function compiled with phase2__next_wealth__growth_rate
                    in its flat signature
                                                        │
                                                        ▼
Runtime:            merge_cross_boundary_params copies
                    internal_params["phase2"]["next_wealth__growth_rate"]
                    → internal_params["phase1"]["phase2__next_wealth__growth_rate"]
                                                        │
                                                        ▼
Call site:          max_Q_over_a(**internal_params["phase1"])
                    — the cross-boundary value is now a regular flat kwarg

Key data structures:

Why not just index the target regime directly? The compiled Q function has a flat signature determined at model init time. It expects named kwargs — it cannot accept a nested dict or index into internal_params at runtime. The cross-boundary machinery bridges the gap between “params are owned by the target regime” and “compiled functions have flat signatures.”

Regime Transition Mechanics¶

Regime transitions determine which regime the agent enters next period. Internally, both deterministic and stochastic transitions are converted to a uniform probability array format.

Deterministic transitions → one-hot encoding¶

A RegimeTransition wraps a function that returns an integer regime ID. Internally, _wrap_deterministic_regime_transition converts this to a one-hot probability array using jax.nn.one_hot:

Stochastic transitions → probability array¶

A MarkovRegimeTransition wraps a function that directly returns a probability array. No conversion is needed — the array is used as-is.

After wrapping, the probability array is further converted to a dictionary keyed by regime name (via _wrap_regime_transition_probs), giving a uniform internal representation regardless of whether the original transition was deterministic or stochastic.

How Transitions Compose in the Bellman Equation¶

The value function computation depends on the regime type:

Terminal regimes¶

No continuation value. The value function equals the utility directly:

Non-terminal with deterministic regime transition¶

The continuation value comes from a single next-period regime:

where $r'$ is the deterministically chosen next regime and $s' = g(s, a)$ is the next-period state.

Non-terminal with stochastic regime transition¶

The continuation value is an expectation over possible next regimes:

where $p_r$ is the probability of transitioning to regime $r$.

Adding stochastic state transitions¶

When a state has a Markov transition (DiscreteMarkovGrid) or shock grid, an additional layer of expectation is added inside the max:

The inner sum handles the stochastic state transition; the outer sum handles the stochastic regime transition. When either is deterministic, its corresponding sum collapses to a single term.

Summary¶

The uniform probability format means the backward induction algorithm treats all transitions the same way — deterministic transitions are just the special case where one probability is 1 and the rest are 0.