Heterogenous agents...in SPACE!

1

I don’t want to clutter Aiyagari example, Transition - #24 by MichaelTiemann with topics too far off-piste from the core transition path topic, so here’s a new thread. I spent yesterday with “Geography vs Income: The Heterogenous effects of carbon taxation” by Labrousse and Perereau (2026, DOI: Redirecting).

The first puzzler, easily solved, was how to present a spatial dimension to the toolkit (we know how to build grids, we know how to present age-dependent vectors, and we can even work with an array of agent types). The answer is to register location-specific data using CreateParamVectorIndexes as follows:

k_grid = (1:K)'; % The different locations/regions we want to model
n_k = K;
Params.w_k = [w_1, w_2, ..., w_K]; % 1-by-K vector of location-dependent wages
% Maps Params.w_k to the index/values of endogenous state variable 'k'
ParamVectorIndexes = CreateParamVectorIndexes(Params, 'w_k', 'k');

Awesome. But what happens when we take it to the next (dimensional) level? In the paper, moving from one location to another is an AR process with a migration probability matrix (easy) and a migration cost matrix (a percentage of income, w_k*z_w). A suggestion for how to get both the probability and the cost in a consistent way is to create a z_grid of indexes and then pull the correct value from each of the matrices with that grid value:

% Suppose Rouwenhorst or Tauchen gives you these 5 states:
actual_z_values = [0.8, 0.9, 1.0, 1.1, 1.2]; 

% 1. Set the VFI Toolkit grid to be indices, not the raw values
n_z = length(actual_z_values);
z_grid = 1:n_z; 
pi_z = ... % Your transition matrix from the AR process remains unchanged

% 2. Store the actual values and your kappa_m data as parameters
Param.z_val = actual_z_values;
Param.kappa_m = [0.01, 0.02, 0.03, 0.04, 0.05]; % Corresponding costs

But we cannot walk around the Kronecker product within the return function. We must make everything ultimately map to scalars, so:

% After building your 75x75 pi_joint matrix...

n_Regions = 5;
n_z = 5; % Number of income shocks
n_Joint = n_Regions^2 * n_z;

% Initialize the 1 x 75 vectors to hold the scalar values for each state
Param.z_grid_wage   = zeros(1, n_Joint);
Param.z_grid_kappa  = zeros(1, n_Joint);
Param.z_grid_shock  = zeros(1, n_Joint);
Param.z_grid_T      = zeros(1, n_Joint);
Param.z_grid_P_E    = zeros(1, n_Joint); % Regional energy price index
Param.z_grid_e_bar  = zeros(1, n_Joint); % Regional subsistence energy

for i = 1:n_Joint
    % Unpack the state exactly once on the CPU
    Pair_Idx = ceil(i / n_z);
    k        = mod(Pair_Idx - 1, n_Regions) + 1;
    k_prime  = ceil(Pair_Idx / n_Regions);
    z_idx    = mod(i - 1, n_z) + 1;
    
    % Populate the scalar value for this specific point on the Markov chain
    Param.z_grid_wage(i)  = w_baseline(k); 
    Param.z_grid_kappa(i) = kappa_matrix(k, k_prime);
    Param.z_grid_shock(i) = z_val_matrix(k, z_idx);
    Param.z_grid_T(i)     = T_matrix(k, z_idx);
    Param.z_grid_P_E(i)   = P_E_regional(k);
    Param.z_grid_e_bar(i) = e_bar_regional(k);
end

To ensure the arrayfun engine automatically broadcasts these 1x75 vectors across the state space, we must ensure they are dimensionally aligned with the exogenous state z.

In standard VFI Toolkit models, if you define a parameter in ReturnFnParamNames and its dimensions match the corresponding grid (e.g., shaping it as a [1, 1, 75] array to match the z dimension), the toolkit expands it perfectly:

% Reshape to ensure VFIToolkit broadcasts them along the exogenous 'z' dimension
% (Assuming 'z' is the 3rd dimension in your specific FHorz setup: a, aprime, z)
Param.z_grid_wage  = reshape(Param.z_grid_wage, 1, 1, n_Joint);
Param.z_grid_kappa = reshape(Param.z_grid_kappa, 1, 1, n_Joint);
% ... reshape the rest similarly

% Add them to the return function signature list
ReturnFnParamNames = {'a', 'aprime', 'z_grid_shock', 'z_grid_wage', ...
                      'z_grid_kappa', 'z_grid_T', 'z_grid_P_E', 'z_grid_e_bar', ...
                      'r', 'j', 'Param'};

Now here’s where the AI came up with something I’ve never seen: the integration of a root-finding algorithm into a return function:

function F = HouseholdReturnFn(a, aprime, shock_val, wage_val, kappa_val, T_val, P_E_val, e_bar_val, r, e_j, Param)
    
    % All inputs are now SCALARS provided perfectly by the arrayfun expansion!
    
    % 1. Discretionary Budget 
    % (Assuming e_j is handled natively by your age loop)
    Income = wage_val * shock_val * e_j;
    Total_Resources = Income + (1+r)*a + T_val;
    
    % 2. Deduct Savings, Mobility Frictions, and Subsistence Energy
    Subsistence_Cost = P_E_val * e_bar_val;
    Mobility_Cost    = kappa_val * Income;
    
    X_tilde = Total_Resources - aprime - Mobility_Cost - Subsistence_Cost;
    
    % 3. Mask Infeasibility
    if X_tilde <= 0
        F = -Inf;
        return; % Fast exit for this scalar thread
    end
    
    % 4. Vectorized Newton-Raphson (Now strictly scalar math)
    U = max(X_tilde, 1e-6); 
    tol = 1e-6;      
    diff = inf;
    
    % Pre-compute scalar constants for this thread
    Term_C = Param.Lambda_C * (1.0)^(1 - Param.sigma);
    Term_H = Param.Lambda_H * (Param.P_H_val)^(1 - Param.sigma); % Assuming P_H is passed similarly
    Term_E = Param.Lambda_E * (P_E_val)^(1 - Param.sigma);
    X_term = X_tilde^(1 - Param.sigma);
    
    pow_C = Param.epsilon_C * (1 - Param.sigma);
    pow_H = Param.epsilon_H * (1 - Param.sigma);
    pow_E = Param.epsilon_E * (1 - Param.sigma);
    
    iter = 0;
    while diff > tol && iter < 50
        F_U = Term_C * U^(pow_C) + Term_H * U^(pow_H) + Term_E * U^(pow_E) - X_term;
        dF_dU = Term_C * pow_C * U^(pow_C - 1) + Term_H * pow_H * U^(pow_H - 1) + Term_E * pow_E * U^(pow_E - 1);
        
        U_new = U - (F_U / dF_dU);
        diff = abs(U_new - U);
        U = U_new;
        iter = iter + 1;
    end
    
    F = U;
end

What I think it is doing is making the non-homothetic choice of consumption vs. housing vs. energy following Comin, Lashkari and Mestieri (2021) as cited when elaborating Equation 1 of the Households.

I’ve got a lot of code to stich together, but wanted to share this particular bit for comment, criticism, etc.

2

And before people get too excited about this, I’ve been working from an open-access pre-print of the paper that is missing some key functionality in the published version. So the above models apply to the pre-print, not the final, in case you were wondering.

The AI guidance (after reading info related to the published paper, not what’s above) says:

The Gumbel Preference Shock

The mention of the Gumbel distribution (Ferriere and Navarro, 2025) is the secret weapon that makes this discrete choice framework computationally viable.

Without the Gumbel shock, the value function would have sharp, discontinuous “kinks” wherever a household is exactly indifferent between renting or owning, or staying or moving. These kinks cause standard root-finding GE solvers to fail. The Gumbel shock smooths these kinks out by introducing idiosyncratic preference variance (\rho).

In the VFI Toolkit, you do not need to manually code the Gumbel draws into your return function. The toolkit has built-in extreme value/logit smoothing architectures for discrete choices that calculate the expected maximum (Emax) continuation value automatically. You will just need to pass the variance parameter \rho into the toolkit’s specific solver options when setting up the model.

3

This is a hallucination. The use of Gumbel shocks to convert the discrete choice into an effectively continuous one is not something VFI Toolkit will do. I’m not saying you cannot get the toolkit to do it, I’ve never tried to think about how you might do this, but it is certainly not a feature of the toolkit.

The approach of the toolkit, because it uses GPUs, it typically just to solve for the value function with sharp discountinous kinks. As it says though, this does make GE a little trickier as you cannot use derivative-based methods like Newton-methods nearly as easily. That said, there are plenty of other approaches like simplex, shooting, and CMA-ES that do still work well.

1 Like
4

My Claude-read of the paper. It is a largely standard infinite horizon value fn problem, plus stationary general eqm. V(a,h,z,k), where a is savings, h is own/rent, z is markov labor-productivity, k is (5-values) location. Two complications: z is semi-endogenous shock, and some GE conditions are ‘per location’.

The location/spacial aspect is a 5-valued markov, from the value fn perspective; although as you highlight it requires being rather careful with parameter values.

The tricky part is how the spacial works for GE: " Note also that the asset market is national (one interest rate r, capital is mobile) while labor and housing markets are segmented by region — that segmentation, plus the migration frictions, is exactly what prevents the carbon tax burden from being arbitraged away across space." [Claude quote]
This could be done in the toolkit by careful use of setting up FnsToEvaluate that use ‘location indicators’, which will work but is rather in-elegant.
[It is like what the toolkit does for ‘GE by ptype’, but done on a specific markov dimension; I will think is there is some general concept here that I could code as a toolkit feature. Maybe it is just about ‘GE by conditional restriction’, and then you can just set up five conditional restrictions that are the five locations, this would be very general in terms of how the user can set what level GE is solved at. I feel like this makes sense as a highly flexible approach?]

The location, k, indexes some prices and preference parameters, but this is easy enough to handle, just need five values for each parameter and an if/else inside the ReturnFn to figure out which one to use. Bit verbose to code, but not hard.

There is one main complication, which is that pi\_{}z (the markov transition matrix for labor productivity units) differs by region k [paper is unclear on if z\_{}grid also differs by location or not]. In principle there is the concept of a ‘semi-endogneous shock’ in VFI Toolkit for handling precisely this issue: an exogenous shock that depends on the current value of an endogenous state, but it is only available for plain-vanilla problems so would need dusting off and an overhaul to the modern toolkit standard.
[Many years ago I got an email from someone who wanted to do a model where the shocks depend on size of firm, which I thought was a cool idea and so I implemented semi-endogenous shock for this purpose so he could use it. I haven’t touched it since.]

2 Likes
5

Claude-read the Appendix B about how they compute. Is exactly what I said above, but they also do some transition paths.

Their state-space is “tiny” (as Claude put it)
S = A × Z × K × H = 80 × 7 × 5 × 2 = 5,600 grid points. Assets on an exponential grid over [0, 70]; a̲ = 0, so no borrowing at all. Which is why they can solve the stationary eqm in 30s (see below).

T=140 for the transition path, and the hard part is that there are 14 GE condns.

“Each steady state takes 30 seconds to compute on a personal computer, and a few minutes for a non-linear transition between two distinct steady states. Computing the Jacobian used to update transition’s guesses take up to 30 minutes. The entire code has been written from scratch on Matlab.” [copy-paste from paper appendix]
Sounds like they should have avoided the Broyden-method and done shooting instead, would likely have been faster as avoid the 30 minutes for Jacobian.

6

I read that bit about doing it all from scratch using MATLAB and shed a single tear that they didn’t know how much the VFI Toolkit could have helped. :laughing:

1 Like