Fitted planners, continuous states, and set-valued interventions

[raphaelrrcoelho]

Four commits developed from 29ec740, rebased onto current main, plus the integration fixes that rebase required. This is the RL surface the library was missing — real planners, continuous states, and set-valued actions — and it closes an interface dishonesty that a separate library-wide audit found independently.

What it adds

Real fitted planners. FittedQIteration (agents/fitted.py) is DOVI's backward induction with the (H,S,A) table and (S,A,S) transition tensor replaced by one fitted regressor per (step, action). Its act() reads the step from the observation, computes Q-values and argmaxes with random tie-breaking — it conditions on the observation, and it refuses to act before fitting rather than returning the optimism cap for every action. BoundedFittedQIteration (agents/bounded_fitted.py) recovers the causal bound the fitted backup had to give up, propagating a Manski envelope through the recursion — valid because max and conditional expectation are both monotone.

Continuous states. causalrl.state supplies StateEncoder, OneHotEncoder, IdentityEncoder, RBFEncoder, FeatureTransition and a shared unpack_transitions/TransitionBatch. The tabular case is contained, not discarded: an indicator basis spans every function on a finite state set, so OneHotEncoder plus a least-squares learner reproduces the tabular backup to ~1e-8, asserted directly against a hand-rolled tabular recursion. A generalisation that does not contain its own special case is not trustworthy, so that property is a test rather than a claim.

Function-valued bounds. FunctionalManskiBounds (bounds/functional.py) is causal_q_bounds where feature space has no cells: lower(x,a) = mu·e + r_min·(1−e), upper(x,a) = mu·e + r_max·(1−e), reproducing the tabular function to ~1e-4 with indicator features. Nuisances are cross-fitted, which is load-bearing rather than decorative — an in-sample plug-in makes the interval optimistically tight, the anti-conservative direction and the dangerous one for a bound. OverlapDiagnostic separates the two regimes that look alike from outside: propensities near zero (wide, honest, uninformative) from propensities a flexible model drove near one (tight, and only as good as that model).

Set-valued interventions. do() has always taken a Mapping and pomis() has always returned sets, but Agent.act returned an arm index — so a multi-variable intervention could be identified and never executed. InterventionSpace states which variables are manipulable in a context and to which values; assignments() turns an intervention set into arms; InterventionalAgent returns an intervention; ScalarAgentAdapter lifts the existing arm-indexed agents in unchanged. AdmissibleInterventions recomputes POMIS per context, and its test pins a witness graph where restricting the manipulable set is a latent projection, not a filter — unconstrained POMIS is [{C,D}] but the correct answer under manipulable={D} is [{}, {D}], so a filtering implementation returns nothing and loses optimality.

Pinned mechanisms. PinnedMechanism deploys a known structural equation while fit_scm learns the rest. Mixed models carry a new provenance="mixed", gated for L3 exactly like "fitted" — a model is only as identified as its weakest node. Pinned nodes are still scored, and that score tests the asserted equation rather than measuring a fit.

An honest agent interface. update() was an empty body in ten subclasses — every *BackdoorAgent, CertifiedPolicyAgent, CounterfactualOptimalPolicy, NaiveOffline, and both fitted agents. These are batch learners whose policy comes from fit/ingest_offline; a single reward carries nothing they can act on, so the ABC was asserting something untrue of a third of its implementers. BatchAgent supplies the no-op once and says why.

Documentation completeness, enforced. 127 of 253 exported names had no API-reference entry — among them Agent, CausalMBRLAgent and CausalThompsonSampling, the README's own headline example. mkdocs --strict never caught this because mkdocstrings validates the references that exist, never the ones that are absent. All 127 are written, and test_every_export_appears_in_the_api_reference keeps it closed, resolving through the lazy export map so a renamed export is checked against the attribute it points at.

What the rebase required

The branch predates phase 2 (PoissonGLMFit, BayesianLinearFit, a lean pass), so two conflicts needed semantic resolution rather than a mechanical pick:

• scm/scm.py — the branch widened the L3 abduction guard to provenance in ("fitted", "mixed") and reports which provenance in the message; main's lean pass had bound sorted(non_invertible_nodes()) once instead of twice. Merged version keeps both. Verified against the built library: a "mixed" model with a non-invertible learned node refuses with witness=['X']; an all-pinned model is "specified" and abducts normally; a fitted model with three bad nodes yields witness=['A','Y','Z'], which actually discriminates the sorted binding.
• scm/fit.py — _FAMILY_NAMES gained poisson_glm/bayesian_linear on main and pinned on the branch. All three kept; _family_name verified to return the snake_case name for all seven shipped fitters.

Then the branch's own new test caught the integration gap it was written to catch: PoissonGLMFit and BayesianLinearFit reached __all__ after that test existed, so they had no API tier and no ::: entry. Fixed by placing them, not by relaxing the test.

Pre-existing defects fixed

Surfaced by the gate as pre-existing rather than introduced:

• from causalrl.scm import PinnedMechanism raised ImportError — it was the only fitter missing from that lazy export map.
• _provenance([]) returned "fitted" — it counted pinned nodes and fell through, so a model with no nodes asserted its equations were learned from data. An empty model has no non-invertible node so the L3 guard is unaffected; this is honest labelling, and this library already deleted certify_mean for emitting query="do" on a computation with no intervention.
• FunctionalManskiBounds._fit_fold still annotated its actions array FloatArray after fit was corrected to NDArray[np.int_].

Verification

ruff check clean · ruff format --check 353 files · pyright src 0 errors, 0 warnings, 0 informations · pytest --cov-fail-under=90 1052 passed, 4 skipped, coverage 97.04% (main: 929 passed, 96.79%).

The 4 skips are module-scope importorskips for extras the dev+data matrix omits (numpyro ×2, jax, stable_baselines3) — identical to main. CI's py3.11 numpyro lane is the gate for one assertion here: the bayesian_linear family name, one of the three the fit.py conflict kept, is only exercised there.

Known boundary

causalrl.interference (Aronow & Samii direct/spillover/total effects) is included but is not reachable from agents/ or envs/ — it is causal inference rather than RL, and a concurrent audit is removing exactly that shape of island. It is flagged for a separate decision rather than quietly kept or quietly dropped.