test_benchmark_model_simulates_end_to_end fails with NaN in dead regime V at age 51

[hmgaudecker]

Summary

tests/test_benchmark.py::test_benchmark_model_simulates_end_to_end (marked @pytest.mark.long_running, so deselected from default pixi run -e tests-cpu tests) fails on main with:

lcm.exceptions.InvalidValueFunctionError: Value function at age 51 in regime 'dead': all values are NaN.

This is independent of the runtime-consumption-points work in PR #8 — the dead regime has no consumption action, and the same failure reproduces with both the static-points and runtime-points versions of the consumption grid.

Provenance

• The test was originally masked by a different failure: AttributeError: 'PiecewiseLinSpacedGrid' object has no attribute 'start' in get_benchmark_initial_conditions (introduced when AIME became piecewise; the .start/.stop reads no longer work). PR Anchor consumption grid lower bound to consumption_floor parameter #8 incidentally fixes that by switching to to_jax().min()/.max().
• With that fixed, the next failure is the NaN-in-V described above.

Suggested next step

Re-run with debug logging to identify whether the NaN enters via utility.dead → preferences.bequest, utility_scale_factor, or somewhere else:

model.solve(params=params, log_level="debug", log_path="./debug/")

The benchmark coefficient_rra for type_1 is 0.999077, which sits just outside jnp.isclose(gamma, 1.0)'s default tolerance (~1e-5). Both branches of the where in bequest/utility_scale_factor are evaluated under JAX, and the log-branch hits log(0) if assets_shifted ever lands on a zero — worth inspecting first.

Reproducer

pixi run -e tests-cpu tests aca-model/tests/test_benchmark.py -m long_running

[hmgaudecker]

CI on feature/runtime-action-grids (pylcm) 03ba800 reproduces this — same NaN-in-dead-V

Run https://github.com/OpenSourceEconomics/pylcm/actions/runs/25109160612/job/73578235664 ran on the latest pylcm tip (03ba800, includes cf00e99 "Fail loudly when reading runtime IrregSpacedGrid before substitution") and still hits:

lcm.exceptions.InvalidValueFunctionError: Value function at age 51 in regime 'dead': all values are NaN.

So this is not a missing-points bug — cf00e99 would surface that as a clearer error. The dead-regime NaN is a real aca-model bug, masked until now by the PiecewiseLinSpacedGrid.start AttributeError in get_benchmark_initial_conditions.

Local reproducer

"""Run from /home/hmg/econ/aca-dev/pylcm:
    pixi run -e benchmarks-cuda12 python repro_aca_bench_dead_nan.py
"""
import sys, jax, jax.numpy as jnp

def _log(msg): print(msg, flush=True, file=sys.stderr)

def main():
    from aca_model.benchmark import (
        create_benchmark_model, get_benchmark_initial_conditions, get_benchmark_params,
    )
    from aca_model.agent import preferences

    model = create_benchmark_model()
    fixed_params, params = get_benchmark_params(model=model)
    initial_conditions = get_benchmark_initial_conditions(model=model, n_subjects=20, seed=0)

    # Manual bequest at every (assets, pref_type) gridpoint
    usf = preferences.utility_scale_factor(
        average_consumption=fixed_params["average_consumption"],
        consumption_weight=jnp.asarray(params["consumption_weight"].to_numpy()),
        coefficient_rra=jnp.asarray(params["coefficient_rra"].to_numpy()),
        time_endowment=params["time_endowment"],
        fixed_cost_of_work_intercept=params["fixed_cost_of_work_intercept"],
        fixed_cost_of_work_age_trend=params["fixed_cost_of_work_age_trend"],
        scale_reference_hours=fixed_params["scale_reference_hours"],
        reference_age=fixed_params["reference_age"],
        scale_reference_age=fixed_params["scale_reference_age"],
    )
    _log(f"utility_scale_factor per type = {usf}")
    for t in (0, 1):
        for a in (0.0, 250_000.0, 500_000.0):
            v = preferences.bequest(
                assets=a, pref_type=t,
                bequest_shifter=params["bequest_shifter"],
                scaled_bequest_weight=params["scaled_bequest_weight"],
                consumption_weight=jnp.asarray(params["consumption_weight"].to_numpy()),
                coefficient_rra=jnp.asarray(params["coefficient_rra"].to_numpy()),
                utility_scale_factor=usf,
            )
            _log(f"bequest(a={a:>9.0f}, t={t}) = {float(v):.6e}")

    model.simulate(
        params=params, initial_conditions=initial_conditions,
        period_to_regime_to_V_arr=None, log_level="off",
        check_initial_conditions=False,
    )

if __name__ == "__main__":
    main()

Findings

• utility_scale_factor per type = [6.376e+11, 9.147e-04] — large dynamic range but finite.
• Manual preferences.bequest(assets=..., pref_type=..., ..., utility_scale_factor=usf) at every (assets ∈ {0, 250k, 500k}, pref_type ∈ {0, 1}) gridpoint returns finite, well-behaved values:

  bequest(a=        0, t=0) = -1.483e-01
  bequest(a=   250000, t=0) = -5.062e-02
  bequest(a=   500000, t=0) = -2.550e-02
  bequest(a=        0, t=1) =  2.865e-02
  bequest(a=   250000, t=1) =  2.867e-02
  bequest(a=   500000, t=1) =  2.867e-02
  

• Yet model.simulate(...) immediately raises with all V values NaN.

The bequest math is fine on these inputs, so the NaN must be entering through how the dead regime composes utility.dead with the pref_type-indexed utility_scale_factor registered as a regime function. Most plausible mechanism: pylcm broadcasts consumption_weight / coefficient_rra / utility_scale_factor down to per-cell scalars before they reach bequest, where the code then does consumption_weight[pref_type] (scalar-indexed-by-int → bad shape / NaN).

If that's right, the fix is in aca-model: either pass full-Series args through to bequest without DAG broadcasting, or move the pref_type indexing inside utility.dead rather than in bequest. Either way, the dead regime appears to mishandle the pref-type-indexed scaling because it's both a state-indexed input and a regime function output simultaneously.