Causal Models and Interventions › Code Implementation

Michael BrenndoerferJune 21, 202610 min read
World Models HandbookAll chapters
Part 4 of 12Space, Physics, and Causality
Chapter 21Causal Models and Interventions › Code Implementation
Reading section 10 of 14 · Code Implementation · 10 minSection 10 of 14 · 10 min

Book contents

Section 10 of 14 · 10 min read

Causal Models and Interventions › Code ImplementationOne section · Sections in this chapter10 / 14 ↓

We will simulate this SCM, fit models on observational and interventional data, compute a counterfactual, and evaluate both planners in the true system. The examples use NumPy for computation and Matplotlib with the book's plotting helpers for figures, and they run on a CPU in seconds. The progression gives numerical counterparts for the chapter's central worked claims: ordinary least squares produces a biased observational slope and approximately recovers the causal slope from randomized data, while the same-episode calculation and planner comparison expose the downstream consequences.

Now we define the SCM and sample a historical log. The heat load UU is drawn but kept "hidden" from the modeler: we will only pass it to the data-generating process, never to the naive model. This discipline, keeping the hidden variable present in the generator and absent from every estimator, is what makes the experiment an honest simulation of partial observability rather than a rigged comparison.

In[11]:
Code
import numpy as np

rng = np.random.default_rng(11)

CAUSAL_SLOPE = -1.5  # true effect of cooling on temperature
CONFOUNDING = 0.8  # how strongly operators react to heat load
SIGMA_A = 0.5  # noise in operator behavior
SIGMA_Y = 1.0  # noise in temperature
BASELINE = 2.0  # baseline temperature deviation
U_VAR = 1.0  # prior variance of the hidden heat load

n_obs = 4000


def sample_observation(n, rng):
    """Historical log: operators react to the hidden heat load."""
    u = rng.normal(size=n)  # hidden heat load
    a = CONFOUNDING * u + SIGMA_A * rng.normal(size=n)  # operators' cooling
    y = CAUSAL_SLOPE * a + u + BASELINE + SIGMA_Y * rng.normal(size=n)
    return u, a, y


u_obs, a_obs, y_obs = sample_observation(n_obs, rng)

Next we fit ordinary least squares on the log. This is a linear analogue of a dynamics model trained to minimize prediction error on an offline dataset. Nothing in the estimator encodes how the actions were assigned, which is the point of the comparison.

In[12]:
Code
def fit_ols(x, y):
    X = np.column_stack([x, np.ones_like(x)])
    coef, *_ = np.linalg.lstsq(X, y, rcond=None)
    return coef[0], coef[1]  # slope, intercept


slope_obs, intercept_obs = fit_ols(a_obs, y_obs)
Out[13]:
Console
Naive observational fit:   Y ≈ -0.618·A + 1.989
True causal mechanism:     Y ≈ -1.500·A + 2.000

The fitted slope is around −0.60-0.60: the naive model underestimates the power of cooling by more than half, exactly as the population calculation predicted. Ordinary predictive diagnostics cannot by themselves establish whether that slope has a causal interpretation, although they can still reveal weak fit or misspecification. Notice also what the intercept is doing: it sits near the baseline of 2.02.0, so the model gets the "no cooling" prediction roughly right. The error is concentrated in the slope, which is the quantity that matters under intervention, a reminder that predictive quality and interventional quality can decouple one parameter at a time.

Now we generate interventional data: instead of letting operators choose the cooling, we randomize it uniformly, as a controlled experiment would. On this data, the backdoor is severed and regression recovers the mechanism.

In[14]:
Code
n_int = 4000
u_int = rng.normal(size=n_int)
a_int = rng.uniform(0.0, 3.0, size=n_int)  # do(A): randomized cooling
y_int = (
    CAUSAL_SLOPE * a_int + u_int + BASELINE + SIGMA_Y * rng.normal(size=n_int)
)

slope_int, intercept_int = fit_ols(a_int, y_int)
Out[15]:
Console
Interventional fit:        Y ≈ -1.503·A + 2.018

The interventional slope lands on the true −1.5-1.5. This contrast is the chapter in one line: same world, same sample size, different data-collection mechanism, entirely different model. The regression did not get smarter, the data did not get bigger, and the noise did not get smaller. The only change is that the action was drawn independently of the hidden load, which deleted the backdoor path and turned the observational distribution into an interventional one. The visual comparison makes the bias concrete.

Out[16]:
Visualization
Scatter plot of temperature versus cooling with a flat dashed naive fit and a steeper solid causal line.
Observational log data with the naive regression (dashed) and the true causal mechanism (solid). Reactive operator behavior confounds and flattens the apparent effect of cooling.
Scatter plot of randomized cooling versus temperature where the dashed naive fit disagrees with the fitted interventional line.
Interventional data from randomized cooling. The regression now recovers the true mechanism, and the naive line is visibly wrong over the operating range.

In the left panel, the naive regression line is much shallower than the true mechanism: the hot-day confounder props up the temperature whenever cooling is high. In the right panel, randomization breaks that pairing, and the same regression procedure lands close to the truth in this finite sample. The observational actions span roughly [−3.6,3.6][-3.6, 3.6], while the randomized actions fill [0,3][0, 3], so the two distributions overlap only partially. The naive model's unconstrained prescription of 3.333.33 lies in a weakly supported upper tail (about 0.03 percent of logged actions exceed it) and is clipped to 3.03.0. Predictions there combine high-leverage, limited-support risk with the global confounding bias in the fitted slope.

Now the counterfactual. We pick one logged episode, apply abduction to infer its hidden heat load, and ask what the temperature would have been under a stronger intervention in that same episode. Because we simulated the system, we can also reveal the true hidden load and check the answer. The abduction step is implemented as a weighted least-squares estimate, which for Gaussian noise and a Gaussian prior is exactly the posterior mean of UU given the observations; the weights are the inverse variances of the three information sources, the action equation, the temperature equation, and the prior.

In[17]:
Code
## Take one specific logged episode (A = 1.02, Y = 2.94 after rounding)
i = 1694
a0, y0, u_true = a_obs[i], y_obs[i], u_obs[i]

## Abduction: weighted least-squares estimate of u from the observed pair.
## Residuals: eta_A = a0 - 0.8*u  (variance 0.25)
##            eta_Y = y0 + 1.5*a0 - u - 2  (variance 1.0)
## Prior:     u ~ N(0, 1)
prec_A, prec_Y, prec_U = 1 / SIGMA_A**2, 1 / SIGMA_Y**2, 1 / U_VAR
H = CONFOUNDING**2 * prec_A + prec_Y + prec_U
b = CONFOUNDING * a0 * prec_A + (y0 - CAUSAL_SLOPE * a0 - BASELINE) * prec_Y
u_hat = b / H

## Action: replace the mechanism with do(A = 1.5). For a same-episode
## counterfactual, retain every realized disturbance. In this additive SCM,
## differencing the two structural equations retains eta_Y automatically.
a_cf = 1.5
y_cf = y0 + CAUSAL_SLOPE * (a_cf - a0)
Out[18]:
Console
Observed episode: A = 1.02, Y = 2.94 (true hidden U = 1.22)
Abducted heat load:  û = 1.25  (truth 1.22)
Counterfactual: Y had we done A = 1.5: 2.22
Naive fitted prediction at A = 1.5: 1.06

Three numbers deserve attention. The posterior mean u^\hat{u} happens to land close to the simulated hidden load for this episode, but its posterior standard deviation is about 0.470.47; it is an uncertain estimate of one background factor, not recovery of an individual's identity. The same-episode counterfactual retains the realized outcome disturbance through ya′=y0−1.5(a′−a0)y_{a'}=y_0-1.5(a'-a_0) and therefore answers the stated room-and-day question. The finite-sample naive fitted prediction uses the confounded observational regression and does not condition on this episode. Counterfactual reasoning is appropriate when the estimand explicitly concerns the same episode under an alternative action, and its validity remains conditional on the structural model.

Finally, we close the loop and let each model plan. The planner itself is identical in both cases: pick the constant cooling effort that drives the predicted temperature to the comfort target. Only the model handed to the planner differs. We evaluate both plans in the true system. Holding the planner fixed is what isolates the causal variable in this experiment: any difference in outcome is attributable to the model, not to the planning algorithm.

In[19]:
Code
target = 0.0


def plan_constant(slope, intercept):
    """Solve slope * a + intercept = target, clipped to the actuator range."""
    return np.clip((target - intercept) / slope, 0.0, 3.0)


a_naive = plan_constant(slope_obs, intercept_obs)
a_causal = plan_constant(slope_int, intercept_int)


def true_expected_cost(a):
    """E[(Y - target)^2] under do(A = a) in the true system."""
    mean_y = CAUSAL_SLOPE * a + BASELINE  # E[U] = 0
    var_y = U_VAR + SIGMA_Y**2  # Var(U) + Var(noise)
    return (mean_y - target) ** 2 + var_y


## The naive model's believed residual variance, estimated from its own fit
naive_resid_var = np.var(y_obs - (slope_obs * a_obs + intercept_obs))
Out[20]:
Console
Naive plan:  A = 3.00  ->  true cost 8.25
Causal plan: A = 1.34  ->  true cost 2.00
Causal plan cost predicted by naive model: 2.62

Read the last line carefully: by the naive model's own accounting, its plan is better, yet in the true system that plan costs several times more than the causal plan. The planner is optimal for the fitted model, but the fitted slope does not represent the intervention. Robust constraints or pessimism could limit the damage without making that slope causally correct. The code estimates in-distribution residual variance; it does not construct calibrated predictive uncertainty, so this example shows that residual fit fails to test interventional validity rather than proving a particular uncertainty estimate is miscalibrated.

In[21]:
Code
costs = [true_expected_cost(a_naive), true_expected_cost(a_causal)]
Out[22]:
Visualization
Bar chart comparing true planning cost for naive versus causal model plans, with the naive bar roughly four times taller.
True expected comfort cost of the plan chosen by each model. The naive (confounded) model selects far too much cooling, and its plan performs several times worse than the causal plan, while predicting success.

The bar chart quantifies the penalty: the naive plan's true cost is roughly four times the causal plan's, even though the naive model assigned its own plan a near-minimal predicted cost. The asymmetry between the two bars is the chapter's final lesson in numerical form. The difference in cost is not driven by noise, which both plans share equally through the Var(U)+SIGMAY2\text{Var}(U) + \text{SIGMA}_Y^2 term in the expected cost; it is driven entirely by the squared bias term, the distance between each plan's mean temperature and the target. Bias is what confounding buys you, and squared bias is what it costs.