The Phase Grammar

Every pylcm model lives in two phases: solve (backward induction over the state-action grid) and simulate (forward sampling of subjects). Most regime slots mean the same thing in both phases — but some quantities genuinely differ between them, and the phase grammar lets you say so in one place.

The grammar is one idea: phase is a broadcast dimension of the regime specification.

A bare value broadcasts to both phases — write a function or a grid once and both phases use it.
Phased(solve=..., simulate=...) specifies each phase explicitly.

Phased is accepted where a per-phase variant makes sense:

functions — per-phase implementations.
state_transitions — per-phase laws of motion.
transition — per-phase regime transitions (matching forms; for per-target dicts, identical key sets).
koopmans_aggregator — one callable per phase.
states — only the combination Phased(solve=callable, simulate=Grid), the carried state described below.
joint_transitions[target][kernel] — around the whole JointTransition; both sides keep the same outputs, support size, and static support schema.

constraints, actions, active, and derived_categoricals are phase-invariant — solve and simulate must agree on what is feasible and what can be chosen, otherwise simulated agents would face a different problem than the one their policy was computed for. Phased is rejected there with an explanation.

Ordinary declarations keep Phased outermost: it never nests inside a per-target transition dict. Structured objects own three explicit inner seams instead: CollectiveUtility.utilities[stakeholder], StakeholderRoute.fallback, and the whole joint kernel above. Those named fields are exceptions, not a general license to place Phased inside arbitrary mappings. Phased itself never nests.

Carried States¶

A carried state is a state that the policy does not condition on, but whose true value the simulation must track — e.g. actual wealth in private pensions when the policy was solved on a value imputed from public pensions etc. in order to keep the state space manageable. It is spelled Phased(solve=callable, simulate=Grid) in states, giving one quantity two roles:

solve: a derived function — the quantity is computed from other states by the callable and never becomes a grid dimension, so the value-function grid does not grow;
simulate: a genuine state — seeded from the initial conditions and evolved each period by its ordinary state_transitions law, with the Grid as its domain.

Decisions are evaluated at the solve-phase imputation — the value the solved policy was computed for; every other simulate consumer reads the carried true value.

This pays off when a state matters for subjects’ histories but is well approximated by a function of other states for decision-making: you keep the per-subject dynamics in simulation without paying for another axis of the value function.

The example below tracks pension wealth. During solve it is imputed from average earnings (aime); during simulation it is a real state that compounds at a fixed rate.

import pprint

import jax.numpy as jnp

from lcm import AgeGrid, LinSpacedGrid, Model, Phased, Regime, categorical
from lcm.typing import FloatND, ScalarInt

RETIREMENT_AGE = 62
MAX_AGE = 63


@categorical(ordered=False)
class RegimeId:
    working: ScalarInt
    dead: ScalarInt


def next_regime(age: float) -> ScalarInt:
    return jnp.where(age >= RETIREMENT_AGE, RegimeId.dead, RegimeId.working)


def impute_pension_wealth(aime: float) -> float:
    """Solve-phase pension wealth: imputed from average earnings."""
    return 0.1 * aime


def evolve_pension_wealth(pension_wealth: float) -> float:
    """Simulate-phase law of motion: compounds at a fixed rate."""
    return 1.03 * pension_wealth


def utility(consumption: float) -> FloatND:
    return jnp.log(consumption)


def next_wealth(*, wealth: float, consumption: float, pension_wealth: float) -> float:
    return wealth - consumption + pension_wealth


def next_aime(aime: float) -> float:
    return aime


def consumption_feasible(*, consumption: float, wealth: float) -> bool:
    return consumption <= wealth

working = Regime(
    transition=next_regime,
    active=lambda age: age < MAX_AGE,
    states={
        "wealth": LinSpacedGrid(start=1.0, stop=100.0, n_points=10),
        "aime": LinSpacedGrid(start=1.0, stop=50.0, n_points=5),
        # The carried state: derived during solve, a real state in simulation.
        "pension_wealth": Phased(
            solve=impute_pension_wealth,
            simulate=LinSpacedGrid(start=0.0, stop=20.0, n_points=4),
        ),
    },
    state_transitions={
        "wealth": next_wealth,
        "aime": next_aime,
        # The carried state's law of motion is an ordinary entry.
        "pension_wealth": evolve_pension_wealth,
    },
    actions={"consumption": LinSpacedGrid(start=1.0, stop=10.0, n_points=5)},
    constraints={"consumption_feasible": consumption_feasible},
    functions={"utility": utility},
)

dead = Regime(transition=None, functions={"utility": lambda: 0.0})

model = Model(
    regimes={"working": working, "dead": dead},
    ages=AgeGrid(start=60, stop=63, step="Y"),
    regime_id_class=RegimeId,
)

The Params Template Unions Both Phases¶

The params template reads the regime in user vocabulary, before the phase split. Where a slot differs by phase, the parameters of both variants appear in the template — a parameter needed by only one phase is still a parameter of the model. Below, pension_wealth (the solve-phase imputation) and next_pension_wealth (the simulate-phase law) both surface:

pprint.pprint(model.get_params_template())

{'dead': {'utility': {}},
 'working': {'certainty_equivalent': {},
             'consumption_feasible': {},
             'koopmans_aggregator': {'discount_factor': 'FloatND'},
             'next_aime': {},
             'next_pension_wealth': {},
             'next_regime': {},
             'next_wealth': {},
             'pension_wealth': {},
             'utility': {}}}

Both Phases in Action¶

Solving uses the imputation (no pension_wealth axis in the value function); simulation seeds pension wealth from the initial conditions and compounds it at 3% per period:

result = model.simulate(
    params={"discount_factor": 0.95},
    initial_conditions={
        "age": jnp.array([60.0, 60.0]),
        "wealth": jnp.array([20.0, 70.0]),
        "aime": jnp.array([10.0, 40.0]),
        "pension_wealth": jnp.array([2.0, 8.0]),
        "regime_id": jnp.array([RegimeId.working] * 2),
    },
    log_level="warning",
)
result.to_dataframe()[
    ["period", "subject_id", "regime_name", "wealth", "pension_wealth"]
]

The pension_wealth column starts at the seeded values (2.0 and 8.0) and grows by the factor 1.03 each period — the simulate-phase law — while the solve phase never saw a pension-wealth grid axis at all.

Wrong Beliefs: Different Transitions in Solve and Simulate¶

The motivating use case for a Phased law of motion: agents believe a state evolves one way, while the data-generating process differs — the policy is solved under the belief, the simulation evolves the truth. Beliefs and truth are distinct parameters, so name them apart with dags.rename_arguments; everything that keeps one name stays one shared parameter. The params template unions both sides:

rho_belief binds only in the solve variant,
rho_true binds only in the simulate variant,
sigma keeps one name, so both phases share its value.

The rational-expectations counterfactual is one line: set rho_belief equal to rho_true in the params dict.

from dags import rename_arguments


def next_income(*, income: float, rho: float, sigma: float) -> float:
    return rho * income + sigma


believer = Regime(
    transition=next_regime,
    active=lambda age: age < MAX_AGE,
    states={"income": LinSpacedGrid(start=0.0, stop=10.0, n_points=11)},
    state_transitions={
        "income": Phased(
            solve=rename_arguments(next_income, mapper={"rho": "rho_belief"}),
            simulate=rename_arguments(next_income, mapper={"rho": "rho_true"}),
        ),
    },
    actions={"consumption": LinSpacedGrid(start=0.1, stop=1.0, n_points=3)},
    functions={
        "utility": lambda consumption, income: jnp.log(consumption + 0.1 * income)
    },
)

beliefs_model = Model(
    regimes={"working": believer, "dead": dead},
    ages=AgeGrid(start=60, stop=63, step="Y"),
    regime_id_class=RegimeId,
)
pprint.pprint(beliefs_model.get_params_template()["working"]["next_income"])

{'rho_belief': 'float', 'rho_true': 'float', 'sigma': 'float'}

beliefs_result = beliefs_model.simulate(
    params={
        "discount_factor": 0.95,
        "working": {
            "next_income": {"rho_belief": 0.95, "rho_true": 0.8, "sigma": 0.5},
        },
    },
    initial_conditions={
        "age": jnp.array([60.0, 60.0]),
        "income": jnp.array([2.0, 4.0]),
        "regime_id": jnp.array([RegimeId.working] * 2),
    },
    log_level="warning",
)
beliefs_result.to_dataframe()[["period", "subject_id", "regime_name", "income"]]

The realized income paths follow the true law (0.8 * income + 0.5: 2.0 → 2.1 and 4.0 → 3.7), while the policy was solved under the believed persistence of 0.95 — the simulated agents act on beliefs and live in the truth.

Which Phase Supplies Which Part of the Decision¶

A simulated agent acts on its beliefs about the future and lives in the truth now. That splits the state-action value $Q$ it maximizes into two halves, each taken from a different phase:

part of $Q$	taken from	why
period utility, feasibility, the Koopmans aggregator	simulate	today’s payoff and today’s feasible set are known when the action is chosen
the continuation — next-period state kernels, regime-transition probabilities, and every helper they read	solve	the future is only perceived, and the value function was solved under those beliefs

The realized next state is drawn from the simulate laws, so the trajectory follows the truth whatever the agent believed. In the vocabulary of the model class, solve supplies the solution phase, while simulate supplies both the decision phase — how today’s actions are ranked — and the realization phase, which transitions actually occur.

The agent is naive: it never anticipates that the world will differ from the model it solved. A sophisticated agent, who does anticipate it, is a different object — its continuation is no longer the solved $V_{t+1}$ — and is not expressible this way.

One deterministic consequence stays on the truth side: a next_<state> that this period’s utility reads is the result of the agent’s own action, known when the action is taken, so it resolves to its simulate variant. Misperception enters only at the continuation boundary.

Perceived Risk: A `Phased` Stochastic Law¶

The belief above was about a deterministic law. The same works for risk: give each phase its own MarkovTransition and the agent solves under a perceived kernel while the simulation draws from the true one. This is the natural home for perceived mortality, perceived health risk, or a policy rule the agent misreads.

Below, a worker believes good health persists, while in truth it does not. Bad health costs MEDICAL_COST out of wealth, so what the worker believes changes how much it saves.

import pandas as pd

from lcm import DiscreteGrid, MarkovTransition
from lcm.typing import BoolND, DiscreteState

MEDICAL_COST = 2.0


@categorical(ordered=False)
class Health:
    good: ScalarInt
    bad: ScalarInt


def health_believed(health: DiscreteState) -> FloatND:
    """Perceived: good health persists."""
    return jnp.where(
        health == Health.good,
        jnp.array([1.0, 0.0]),
        jnp.array([0.0, 1.0]),
    )


def health_true(health: DiscreteState) -> FloatND:
    """Realized: good health does not persist."""
    return jnp.where(
        health == Health.good,
        jnp.array([0.0, 1.0]),
        jnp.array([0.0, 1.0]),
    )


def resources(*, wealth: float, health: DiscreteState) -> FloatND:
    """Wealth net of the medical expense bad health imposes."""
    return wealth - MEDICAL_COST * (health == Health.bad)


def health_next_wealth(*, resources: FloatND, consumption: float) -> FloatND:
    return 1.02 * (resources - consumption)


def health_feasible(*, consumption: float, resources: FloatND) -> BoolND:
    return consumption <= resources


def build_health_model(
    health_law: MarkovTransition | Phased[MarkovTransition, MarkovTransition],
) -> Model:
    worker = Regime(
        transition=next_regime,
        active=lambda age: age < MAX_AGE,
        states={
            "wealth": LinSpacedGrid(start=4.0, stop=60.0, n_points=57),
            "health": DiscreteGrid(category_class=Health),
        },
        state_transitions={"wealth": health_next_wealth, "health": health_law},
        actions={"consumption": LinSpacedGrid(start=0.5, stop=30.0, n_points=60)},
        constraints={"consumption_feasible": health_feasible},
        functions={"utility": utility, "resources": resources},
    )
    return Model(
        regimes={"working": worker, "dead": dead},
        ages=AgeGrid(start=60, stop=63, step="Y"),
        regime_id_class=RegimeId,
    )

naive = build_health_model(
    Phased(
        solve=MarkovTransition(health_believed),
        simulate=MarkovTransition(health_true),
    )
)
rational = build_health_model(MarkovTransition(health_true))

health_initial_conditions = {
    "age": jnp.array([60.0]),
    "wealth": jnp.array([20.0]),
    "health": jnp.array([Health.good]),
    "regime_id": jnp.array([RegimeId.working]),
}

paths = {
    label: health_model.simulate(
        params={"discount_factor": 0.95},
        initial_conditions=health_initial_conditions,
        log_level="warning",
    )
    .to_dataframe()
    .set_index("period")[["health", "wealth", "consumption"]]
    for label, health_model in (("naive", naive), ("rational", rational))
}

pd.concat(paths, names=["beliefs"])

The naive worker consumes more in the first period than the rational one, who knows the medical expense is coming and carries the difference forward as precautionary saving. Both fall ill in the next period regardless: the draw follows the true kernel either way. The belief changed the policy, not the world. The last row is the terminal regime, where nothing is chosen.

Set the two variants equal — or drop the Phased wrapper — and the naive worker becomes the rational one, which is the counterfactual worth reporting alongside any result built on misperception.

The two phases need not agree on whether a law is stochastic¶

A deterministic law is a degenerate kernel, not a different kind of state, so Phased(solve=MarkovTransition(...), simulate=<deterministic>) and its reverse are both legal. An agent may perceive risk where there is none, or treat as certain a transition that is not.

What stays phase-invariant¶

Feasibility is a primitive of the model the agent solved, so constraints may not vary by phase — not directly, and not through a helper or a law of motion they read. A phase-specific feasible set would let the simulated agent choose actions its value function was never computed for. actions, active, and derived_categoricals are phase-invariant for the same reason.

Solver support can be narrower than the public phase grammar. NNBEGM replays its solve-time candidate bank, so a bare declaration or Phased(solve=f, simulate=f) is accepted only when both fields contain the same object. Genuine phase variation is rejected during Model(...) construction. See NNBEGM replay capability.